Separate active YAMN code from Katzenpost PoC

6 files changed, 316 insertions, 0 deletions

diff --git a/katzenpost/dispatcher/codec.go b/katzenpost/dispatcher/codec.go
new file mode 100644
index 0000000..3053dde
--- /dev/null
+++ b/katzenpost/dispatcher/codec.go
@@ -0,0 +1,29 @@
+package dispatcher
+
+import (
+	"fmt"
+
+	"github.com/fxamacker/cbor/v2"
+)
+
+func EncodeEnvelope(envelope Envelope) ([]byte, error) {
+	encoded, err := cbor.Marshal(envelope)
+	if err != nil {
+		return nil, fmt.Errorf("encode envelope: %w", err)
+	}
+	if len(encoded) > MaxEnvelopeBytes {
+		return nil, fmt.Errorf("encoded envelope is too large")
+	}
+	return encoded, nil
+}
+
+func DecodeEnvelope(encoded []byte) (Envelope, error) {
+	if len(encoded) == 0 || len(encoded) > MaxEnvelopeBytes {
+		return Envelope{}, fmt.Errorf("invalid encoded envelope size")
+	}
+	var envelope Envelope
+	if err := cbor.Unmarshal(encoded, &envelope); err != nil {
+		return Envelope{}, fmt.Errorf("decode envelope: %w", err)
+	}
+	return envelope, nil
+}
diff --git a/katzenpost/dispatcher/envelope.go b/katzenpost/dispatcher/envelope.go
new file mode 100644
index 0000000..b9ff7b8
--- /dev/null
+++ b/katzenpost/dispatcher/envelope.go
@@ -0,0 +1,45 @@
+package dispatcher
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"net/mail"
+	"strings"
+)
+
+const MaxEnvelopeBytes = 64 * 1024
+
+// Envelope is the only application payload accepted by the dispatcher. Message
+// must already be a complete YAMN transport envelope, not user supplied mail.
+type Envelope struct {
+	EntryAddress string `json:"entry_address"`
+	Message      []byte `json:"message"`
+}
+
+func (e Envelope) Validate(allowedDomains map[string]struct{}) error {
+	if len(e.Message) == 0 || len(e.Message) > MaxEnvelopeBytes {
+		return fmt.Errorf("invalid YAMN envelope size")
+	}
+	if bytes.Contains(e.Message, []byte("\r")) {
+		return errors.New("YAMN envelope must use LF line endings")
+	}
+	address, err := mail.ParseAddress(e.EntryAddress)
+	if err != nil || address.Address != e.EntryAddress {
+		return errors.New("invalid entry remailer address")
+	}
+	parts := strings.Split(address.Address, "@")
+	if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
+		return errors.New("invalid entry remailer address")
+	}
+	if _, ok := allowedDomains[strings.ToLower(parts[1])]; !ok {
+		return errors.New("entry remailer domain is not allowed")
+	}
+	if !bytes.HasPrefix(e.Message, []byte("To: "+e.EntryAddress+"\n")) {
+		return errors.New("YAMN envelope recipient does not match entry remailer")
+	}
+	if !bytes.Contains(e.Message, []byte("-----BEGIN REMAILER MESSAGE-----\n")) {
+		return errors.New("invalid YAMN envelope")
+	}
+	return nil
+}
diff --git a/katzenpost/dispatcher/envelope_test.go b/katzenpost/dispatcher/envelope_test.go
new file mode 100644
index 0000000..712071c
--- /dev/null
+++ b/katzenpost/dispatcher/envelope_test.go
@@ -0,0 +1,21 @@
+package dispatcher
+
+import "testing"
+
+func TestEnvelopeValidate(t *testing.T) {
+	allowed := map[string]struct{}{"remailer.example": {}}
+	valid := Envelope{
+		EntryAddress: "entry@remailer.example",
+		Message:      []byte("To: entry@remailer.example\nFrom: mix@nowhere.invalid\n\n-----BEGIN REMAILER MESSAGE-----\nabc\n"),
+	}
+	if err := valid.Validate(allowed); err != nil {
+		t.Fatalf("expected valid envelope: %v", err)
+	}
+
+	invalidDomain := valid
+	invalidDomain.EntryAddress = "entry@other.example"
+	invalidDomain.Message = []byte("To: entry@other.example\n-----BEGIN REMAILER MESSAGE-----\n")
+	if err := invalidDomain.Validate(allowed); err == nil {
+		t.Fatal("expected unknown domain rejection")
+	}
+}
diff --git a/katzenpost/dispatcher/frame.go b/katzenpost/dispatcher/frame.go
new file mode 100644
index 0000000..43f547c
--- /dev/null
+++ b/katzenpost/dispatcher/frame.go
@@ -0,0 +1,84 @@
+package dispatcher
+
+import (
+	"crypto/rand"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+
+	"github.com/fxamacker/cbor/v2"
+)
+
+const (
+	ProtocolVersion  = 1
+	MaxSphinxPayload = 2000
+	FrameDataBytes   = 1500
+	MaxFrames        = (MaxEnvelopeBytes + FrameDataBytes - 1) / FrameDataBytes
+)
+
+type Frame struct {
+	Version uint8    `cbor:"1,keyasint"`
+	ID      [16]byte `cbor:"2,keyasint"`
+	Digest  [32]byte `cbor:"3,keyasint"`
+	Index   uint16   `cbor:"4,keyasint"`
+	Total   uint16   `cbor:"5,keyasint"`
+	Data    []byte   `cbor:"6,keyasint"`
+}
+
+func SplitEnvelope(encoded []byte) ([][]byte, error) {
+	if len(encoded) == 0 || len(encoded) > MaxEnvelopeBytes {
+		return nil, errors.New("invalid encoded envelope size")
+	}
+
+	var id [16]byte
+	if _, err := rand.Read(id[:]); err != nil {
+		return nil, fmt.Errorf("generate message id: %w", err)
+	}
+	digest := sha256.Sum256(encoded)
+	total := (len(encoded) + FrameDataBytes - 1) / FrameDataBytes
+	if total > MaxFrames {
+		return nil, errors.New("encoded envelope requires too many frames")
+	}
+
+	frames := make([][]byte, 0, total)
+	for index, start := 0, 0; start < len(encoded); index, start = index+1, start+FrameDataBytes {
+		end := min(start+FrameDataBytes, len(encoded))
+		frame := Frame{
+			Version: ProtocolVersion,
+			ID:      id,
+			Digest:  digest,
+			Index:   uint16(index),
+			Total:   uint16(total),
+			Data:    encoded[start:end],
+		}
+		payload, err := cbor.Marshal(frame)
+		if err != nil {
+			return nil, fmt.Errorf("encode frame %d: %w", index, err)
+		}
+		if len(payload) > MaxSphinxPayload {
+			return nil, fmt.Errorf("frame %d exceeds Sphinx payload limit", index)
+		}
+		frames = append(frames, payload)
+	}
+	return frames, nil
+}
+
+func DecodeFrame(payload []byte) (Frame, error) {
+	if len(payload) == 0 || len(payload) > MaxSphinxPayload {
+		return Frame{}, errors.New("invalid frame size")
+	}
+	var frame Frame
+	if err := cbor.Unmarshal(payload, &frame); err != nil {
+		return Frame{}, fmt.Errorf("decode frame: %w", err)
+	}
+	if frame.Version != ProtocolVersion {
+		return Frame{}, errors.New("unsupported frame version")
+	}
+	if frame.Total == 0 || int(frame.Total) > MaxFrames || frame.Index >= frame.Total {
+		return Frame{}, errors.New("invalid frame sequence")
+	}
+	if len(frame.Data) == 0 || len(frame.Data) > FrameDataBytes {
+		return Frame{}, errors.New("invalid frame data size")
+	}
+	return frame, nil
+}
diff --git a/katzenpost/dispatcher/frame_test.go b/katzenpost/dispatcher/frame_test.go
new file mode 100644
index 0000000..21c84ba
--- /dev/null
+++ b/katzenpost/dispatcher/frame_test.go
@@ -0,0 +1,46 @@
+package dispatcher
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+func TestSplitAndReassembleEnvelope(t *testing.T) {
+	encoded := bytes.Repeat([]byte("yamn-envelope\n"), 900)
+	payloads, err := SplitEnvelope(encoded)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(payloads) < 2 {
+		t.Fatal("expected multiple frames")
+	}
+
+	reassembler := NewReassembler(time.Minute)
+	var result []byte
+	for _, payload := range payloads {
+		if len(payload) > MaxSphinxPayload {
+			t.Fatalf("payload exceeds limit: %d", len(payload))
+		}
+		frame, err := DecodeFrame(payload)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assembled, complete, err := reassembler.Add(frame)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if complete {
+			result = assembled
+		}
+	}
+	if !bytes.Equal(result, encoded) {
+		t.Fatal("reassembled envelope differs")
+	}
+}
+
+func TestDecodeFrameRejectsOversizedPayload(t *testing.T) {
+	if _, err := DecodeFrame(make([]byte, MaxSphinxPayload+1)); err == nil {
+		t.Fatal("expected oversized payload rejection")
+	}
+}
diff --git a/katzenpost/dispatcher/reassembler.go b/katzenpost/dispatcher/reassembler.go
new file mode 100644
index 0000000..87af98c
--- /dev/null
+++ b/katzenpost/dispatcher/reassembler.go
@@ -0,0 +1,91 @@
+package dispatcher
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"errors"
+	"sync"
+	"time"
+)
+
+const MaxPendingMessages = 64
+
+type pendingMessage struct {
+	digest  [32]byte
+	total   uint16
+	parts   map[uint16][]byte
+	created time.Time
+}
+
+type Reassembler struct {
+	mu      sync.Mutex
+	pending map[[16]byte]*pendingMessage
+	ttl     time.Duration
+	now     func() time.Time
+}
+
+func NewReassembler(ttl time.Duration) *Reassembler {
+	return &Reassembler{
+		pending: make(map[[16]byte]*pendingMessage),
+		ttl:     ttl,
+		now:     time.Now,
+	}
+}
+
+func (r *Reassembler) Add(frame Frame) ([]byte, bool, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.expireLocked()
+	pending, ok := r.pending[frame.ID]
+	if !ok {
+		if len(r.pending) >= MaxPendingMessages {
+			return nil, false, errors.New("too many pending messages")
+		}
+		pending = &pendingMessage{
+			digest:  frame.Digest,
+			total:   frame.Total,
+			parts:   make(map[uint16][]byte, frame.Total),
+			created: r.now(),
+		}
+		r.pending[frame.ID] = pending
+	}
+	if pending.total != frame.Total || pending.digest != frame.Digest {
+		delete(r.pending, frame.ID)
+		return nil, false, errors.New("inconsistent frame metadata")
+	}
+	if _, duplicate := pending.parts[frame.Index]; !duplicate {
+		pending.parts[frame.Index] = bytes.Clone(frame.Data)
+	}
+	if len(pending.parts) != int(pending.total) {
+		return nil, false, nil
+	}
+
+	var assembled bytes.Buffer
+	for index := uint16(0); index < pending.total; index++ {
+		part, exists := pending.parts[index]
+		if !exists {
+			return nil, false, nil
+		}
+		if assembled.Len()+len(part) > MaxEnvelopeBytes {
+			delete(r.pending, frame.ID)
+			return nil, false, errors.New("reassembled envelope is too large")
+		}
+		assembled.Write(part)
+	}
+	delete(r.pending, frame.ID)
+	result := assembled.Bytes()
+	if sha256.Sum256(result) != pending.digest {
+		return nil, false, errors.New("reassembled envelope digest mismatch")
+	}
+	return result, true, nil
+}
+
+func (r *Reassembler) expireLocked() {
+	cutoff := r.now().Add(-r.ttl)
+	for id, pending := range r.pending {
+		if pending.created.Before(cutoff) {
+			delete(r.pending, id)
+		}
+	}
+}