refactor(wrap): переписать как мимикрию под DTLS 1.2 ApplicationData с AEAD

Заменяет наивный wire-формат [12B nonce | ChaCha20-XOR] на фейковый заголовок DTLS 1.2 record, проходящий DPI-инспекцию: [content_type=0x17 | version=0xFEFD | epoch | seq | length | AEAD ciphertext | tag] Ключевые изменения: - ChaCha20-Poly1305 AEAD вместо неаутентифицированного ChaCha20 XOR - Padding длины по фиксированным корзинам {128..1300} ломает fingerprint по размеру - Direction-safe epochs (MSB=0 клиент, MSB=1 сервер) предотвращают reuse nonce - Nonce выводится из XOR(SHA-256(key), поля заголовка), без константного префикса - Горячий путь сервера использует sync.Pool, убирая аллокацию на пакет Пользовательские флаги (-wrap, -wrap-key, -gen-wrap-key) не изменены.
2026-10-12 09:41:53 +02:00 · 2026-05-18 19:32:59 +03:00
parent b3f81add0e
commit 39f95a8e0c
4 changed files with 456 additions and 148 deletions
@@ -1162,7 +1162,16 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD
 		// Do not set conn2 deadline (conn2 can sometimes be listenConn if direct mode is used)
 	})
 	var internalPipeAddr atomic.Value
-	useWrap := len(turnParams.wrapKey) == wrapKeyLen
+	var wc *wrapConn
+	if len(turnParams.wrapKey) == wrapKeyLen {
+		var wcErr error
+		wc, wcErr = newWrapConn(turnParams.wrapKey, false)
+		if wcErr != nil {
+			log.Printf("[STREAM %d] WRAP init failed: %v", streamID, wcErr)
+			turncancel()
+			return
+		}
+	}

 	go func() {
 		defer turncancel()
@@ -1170,8 +1179,8 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD
 		// Reusable scratch buffer for wrapped wire bytes; sized once per
 		// stream so the hot-path TX loop performs zero allocations.
 		var wireBuf []byte
-		if useWrap {
-			wireBuf = make([]byte, wrapNonceLen+len(buf))
+		if wc != nil {
+			wireBuf = make([]byte, wrapMaxWire(len(buf)))
 		}
 		for {
 			if turnctx.Err() != nil {
@@ -1188,8 +1197,8 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD
 			internalPipeAddr.Store(addr1)

 			out := buf[:n]
-			if useWrap {
-				written, wrapErr := wrapInto(wireBuf, turnParams.wrapKey, out)
+			if wc != nil {
+				written, wrapErr := wc.wrapInto(wireBuf, out)
 				if wrapErr != nil {
 					log.Printf("[STREAM %d] WRAP failed: %v", streamID, wrapErr)
 					return
@@ -1208,8 +1217,8 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD
 	wg.Go(func() {
 		defer turncancel()
 		readBufLen := 1600
-		if useWrap {
-			readBufLen += wrapNonceLen
+		if wc != nil {
+			readBufLen = wrapMaxWire(1600)
 		}
 		buf := make([]byte, readBufLen)
 		plain := make([]byte, 1600)
@@ -1225,8 +1234,8 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD

 			if addr, ok := addr1.(net.Addr); ok {
 				payload := buf[:n]
-				if useWrap {
-					m, wrapErr := unwrapPacket(turnParams.wrapKey, payload, plain)
+				if wc != nil {
+					m, wrapErr := wc.unwrapPacket(payload, plain)
 					if wrapErr != nil {
 						log.Printf("[STREAM %d] UNWRAP failed: %v (n=%d)", streamID, wrapErr, n)
 						continue
@@ -2156,7 +2165,15 @@ func createSmuxSession(ctx context.Context, tp *turnParams, peer *net.UDPAddr, i
 		cleanup()
 		return nil, nil, fmt.Errorf("generate cert: %w", err)
 	}
-	dtlsPC := &relayPacketConn{relay: relayConn, peer: peer, wrapKey: tp.wrapKey}
+	var relayWC *wrapConn
+	if len(tp.wrapKey) == wrapKeyLen {
+		relayWC, err = newWrapConn(tp.wrapKey, false)
+		if err != nil {
+			cleanup()
+			return nil, nil, fmt.Errorf("wrap init: %w", err)
+		}
+	}
+	dtlsPC := &relayPacketConn{relay: relayConn, peer: peer, wc: relayWC}
 	dtlsConn, err := dtls.ClientWithOptions(dtlsPC, peer,
 		dtls.WithCertificates(certificate),
 		dtls.WithInsecureSkipVerify(true),
@@ -2205,22 +2222,23 @@ func createSmuxSession(ctx context.Context, tp *turnParams, peer *net.UDPAddr, i
 }

 // relayPacketConn wraps a TURN relay PacketConn to direct all writes to the peer.
+// When wc != nil, packets are wrapped/unwrapped with DTLS-mimicry obfuscation.
 type relayPacketConn struct {
-	relay   net.PacketConn
-	peer    net.Addr
-	wrapKey []byte
+	relay net.PacketConn
+	peer  net.Addr
+	wc    *wrapConn
 }

 func (r *relayPacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
-	if len(r.wrapKey) != wrapKeyLen {
+	if r.wc == nil {
 		return r.relay.ReadFrom(b)
 	}
-	buf := make([]byte, len(b)+wrapNonceLen)
+	buf := make([]byte, wrapMaxWire(len(b)))
 	n, addr, err := r.relay.ReadFrom(buf)
 	if err != nil {
 		return 0, addr, err
 	}
-	m, err := unwrapPacket(r.wrapKey, buf[:n], b)
+	m, err := r.wc.unwrapPacket(buf[:n], b)
 	if err != nil {
 		return 0, addr, err
 	}
@@ -2228,14 +2246,15 @@ func (r *relayPacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
 }

 func (r *relayPacketConn) WriteTo(b []byte, _ net.Addr) (int, error) {
-	if len(r.wrapKey) != wrapKeyLen {
+	if r.wc == nil {
 		return r.relay.WriteTo(b, r.peer)
 	}
-	out, err := wrapPacket(r.wrapKey, b)
+	out := make([]byte, wrapMaxWire(len(b)))
+	n, err := r.wc.wrapInto(out, b)
 	if err != nil {
 		return 0, err
 	}
-	if _, err = r.relay.WriteTo(out, r.peer); err != nil {
+	if _, err = r.relay.WriteTo(out[:n], r.peer); err != nil {
 		return 0, err
 	}
 	return len(b), nil
@@ -3,46 +3,189 @@
 package main

 import (
+	"crypto/cipher"
 	"crypto/rand"
+	"crypto/sha256"
 	"encoding/binary"
 	"encoding/hex"
 	"errors"
 	"fmt"
 	"sync/atomic"

-	"golang.org/x/crypto/chacha20"
+	"golang.org/x/crypto/chacha20poly1305"
 )

+// Wire format — DTLS 1.2 ApplicationData mimicry:
+//
+//	[1B content_type=0x17 | 2B version=0xFEFD | 2B epoch | 6B seq_num | 2B length |
+//	 AEAD_ciphertext(padded plaintext) | 16B Poly1305 tag]
+//
+// Plaintext before encryption: [2B payload_length (BE) | payload | random padding]
+// Padded to the nearest bucket in padBuckets to defeat size fingerprinting.
+//
+// Nonce (12B) = XOR(SHA-256(key)[0:12], [00 00 | epoch(2B) | seq(6B) | 00 00])
+// AAD = the 13-byte DTLS record header.
+//
+// Epoch MSB encodes direction: 0 = client, 1 = server.
+// This guarantees no nonce reuse between client and server under the same key.
+
 const (
-	wrapNonceLen = 12
-	wrapKeyLen   = 32
+	wrapKeyLen    = 32
+	wrapHeaderLen = 13 // DTLS record header
+	wrapTagLen    = 16 // Poly1305 tag
+	wrapLenPrefix = 2  // uint16 BE real payload length
+	wrapOverhead  = wrapHeaderLen + wrapTagLen + wrapLenPrefix // 31
 )

-// Nonce layout: [4B session prefix | 8B atomic counter, big-endian].
-// Session prefix is randomised once at process start; counter monotonically
-// increments per-packet. This eliminates the per-packet crypto/rand syscall
-// (which on Linux is a getrandom() round-trip ~hundreds of ns) while keeping
-// (key, nonce) pairs unique across 2^64 packets per session and 2^32 sessions
-// per key. Wire format is byte-identical to the previous random-nonce scheme,
-// so server-side decryption is unchanged.
-var (
-	nonceSessionPrefix [4]byte
-	nonceCounter       atomic.Uint64
-)
+// padBuckets defines plaintext size buckets. The padded plaintext
+// ([2B length | payload | padding]) is rounded up to the nearest bucket.
+var padBuckets = [...]int{128, 256, 512, 768, 1024, 1300}

-func init() {
-	// Best-effort: if rand.Read fails this early, the OS is in trouble and the
-	// binary is going to crash anyway during DTLS cert generation. Leaving the
-	// prefix zero would still keep nonces unique within this process.
-	_, _ = rand.Read(nonceSessionPrefix[:])
+// wrapConn holds per-connection state for DTLS-mimicry obfuscation.
+type wrapConn struct {
+	aead      cipher.AEAD
+	nonceMask [12]byte // SHA-256(key)[0:12]
+	epoch     uint16
+	seq       atomic.Uint64
 }

-func fillWrapNonce(dst []byte) {
-	_ = dst[wrapNonceLen-1] // bounds-check hint
-	copy(dst[:4], nonceSessionPrefix[:])
-	binary.BigEndian.PutUint64(dst[4:wrapNonceLen], nonceCounter.Add(1))
+// newWrapConn creates wrap state. isServer must be true on the server side.
+func newWrapConn(key []byte, isServer bool) (*wrapConn, error) {
+	if len(key) != wrapKeyLen {
+		return nil, fmt.Errorf("wrap: key must be %d bytes (got %d)", wrapKeyLen, len(key))
+	}
+	aead, err := chacha20poly1305.New(key)
+	if err != nil {
+		return nil, fmt.Errorf("wrap: aead init: %w", err)
+	}
+	h := sha256.Sum256(key)
+	var mask [12]byte
+	copy(mask[:], h[:12])
+
+	var eb [2]byte
+	if _, err := rand.Read(eb[:]); err != nil {
+		return nil, fmt.Errorf("wrap: epoch rand: %w", err)
+	}
+	epoch := binary.BigEndian.Uint16(eb[:])
+	if isServer {
+		epoch |= 0x8000
+	} else {
+		epoch &^= 0x8000
+	}
+	return &wrapConn{aead: aead, nonceMask: mask, epoch: epoch}, nil
 }

+// deriveNonce builds the 12-byte AEAD nonce from header fields.
+func (w *wrapConn) deriveNonce(epoch uint16, seqNum uint64) [12]byte {
+	var explicit [12]byte
+	binary.BigEndian.PutUint16(explicit[2:4], epoch)
+	binary.BigEndian.PutUint16(explicit[4:6], uint16(seqNum>>32))
+	binary.BigEndian.PutUint32(explicit[6:10], uint32(seqNum))
+	var nonce [12]byte
+	for i := range 12 {
+		nonce[i] = w.nonceMask[i] ^ explicit[i]
+	}
+	return nonce
+}
+
+// padPlaintextSize returns the padded plaintext size for a given payload.
+func padPlaintextSize(payloadLen int) int {
+	needed := wrapLenPrefix + payloadLen
+	for _, b := range padBuckets {
+		if b >= needed {
+			return b
+		}
+	}
+	return needed
+}
+
+// WrapMaxWire returns the maximum wire length for a given payload size.
+func wrapMaxWire(payloadLen int) int {
+	return wrapHeaderLen + padPlaintextSize(payloadLen) + wrapTagLen
+}
+
+// wrapInto encrypts payload into dst. Returns total wire bytes written.
+// dst must have len >= wrapMaxWire(len(payload)).
+func (w *wrapConn) wrapInto(dst, payload []byte) (int, error) {
+	paddedLen := padPlaintextSize(len(payload))
+	wireLen := wrapHeaderLen + paddedLen + wrapTagLen
+	if len(dst) < wireLen {
+		return 0, errors.New("wrap: dst buffer too small")
+	}
+
+	seq := w.seq.Add(1) - 1
+
+	// --- DTLS record header (13 bytes) ---
+	dst[0] = 0x17 // ApplicationData
+	dst[1] = 0xFE // DTLS 1.2
+	dst[2] = 0xFD
+	binary.BigEndian.PutUint16(dst[3:5], w.epoch)
+	binary.BigEndian.PutUint16(dst[5:7], uint16(seq>>32))
+	binary.BigEndian.PutUint32(dst[7:11], uint32(seq))
+	binary.BigEndian.PutUint16(dst[11:13], uint16(paddedLen+wrapTagLen))
+
+	// --- Build plaintext: [2B len | payload | padding] ---
+	plain := dst[wrapHeaderLen : wrapHeaderLen+paddedLen]
+	binary.BigEndian.PutUint16(plain[0:2], uint16(len(payload)))
+	copy(plain[wrapLenPrefix:], payload)
+	// Random padding for the tail
+	padStart := wrapLenPrefix + len(payload)
+	if padStart < paddedLen {
+		_, _ = rand.Read(plain[padStart:paddedLen])
+	}
+
+	// --- AEAD encrypt in-place ---
+	nonce := w.deriveNonce(w.epoch, seq)
+	aad := dst[:wrapHeaderLen]
+	w.aead.Seal(dst[wrapHeaderLen:wrapHeaderLen], nonce[:], plain, aad)
+
+	return wireLen, nil
+}
+
+// unwrapPacket decrypts a wire-format packet into dst. Returns payload bytes.
+func (w *wrapConn) unwrapPacket(wire, dst []byte) (int, error) {
+	if len(wire) < wrapHeaderLen+wrapTagLen+wrapLenPrefix {
+		return 0, errors.New("wrap: packet too short")
+	}
+	if wire[0] != 0x17 {
+		return 0, fmt.Errorf("wrap: unexpected content type 0x%02X", wire[0])
+	}
+
+	epoch := binary.BigEndian.Uint16(wire[3:5])
+	seqHigh := uint64(binary.BigEndian.Uint16(wire[5:7]))
+	seqLow := uint64(binary.BigEndian.Uint32(wire[7:11]))
+	seqNum := (seqHigh << 32) | seqLow
+	bodyLen := int(binary.BigEndian.Uint16(wire[11:13]))
+
+	if wrapHeaderLen+bodyLen > len(wire) {
+		return 0, errors.New("wrap: truncated packet")
+	}
+
+	nonce := w.deriveNonce(epoch, seqNum)
+	aad := wire[:wrapHeaderLen]
+	ciphertext := wire[wrapHeaderLen : wrapHeaderLen+bodyLen]
+
+	plain, err := w.aead.Open(ciphertext[:0], nonce[:], ciphertext, aad)
+	if err != nil {
+		return 0, fmt.Errorf("wrap: AEAD open: %w", err)
+	}
+
+	if len(plain) < wrapLenPrefix {
+		return 0, errors.New("wrap: plaintext too short for length prefix")
+	}
+	realLen := int(binary.BigEndian.Uint16(plain[0:2]))
+	if realLen > len(plain)-wrapLenPrefix {
+		return 0, fmt.Errorf("wrap: payload length %d exceeds plaintext %d", realLen, len(plain)-wrapLenPrefix)
+	}
+	if realLen > len(dst) {
+		return 0, errors.New("wrap: dst buffer too small")
+	}
+	copy(dst[:realLen], plain[wrapLenPrefix:wrapLenPrefix+realLen])
+	return realLen, nil
+}
+
+// --- Helper functions (unchanged API for flag parsing) ---
+
 func genWrapKeyHex() (string, error) {
 	key := make([]byte, wrapKeyLen)
 	if _, err := rand.Read(key); err != nil {
@@ -67,55 +210,3 @@ func decodeWrapKey(enabled bool, raw string) ([]byte, error) {
 	}
 	return key, nil
 }
-
-// wrapInto encrypts payload into dst using the shared scheme
-// [12B nonce | ChaCha20-XOR(payload)]. dst must have capacity for
-// wrapNonceLen+len(payload). Returns total bytes written. Hot-path callers
-// should reuse a single dst buffer to avoid per-packet allocation.
-func wrapInto(dst, key, payload []byte) (int, error) {
-	if len(key) != wrapKeyLen {
-		return 0, fmt.Errorf("wrap: key must be %d bytes (got %d)", wrapKeyLen, len(key))
-	}
-	total := wrapNonceLen + len(payload)
-	if len(dst) < total {
-		return 0, errors.New("wrap: dst buffer too small")
-	}
-	fillWrapNonce(dst[:wrapNonceLen])
-	cipher, err := chacha20.NewUnauthenticatedCipher(key, dst[:wrapNonceLen])
-	if err != nil {
-		return 0, fmt.Errorf("wrap: cipher init: %w", err)
-	}
-	cipher.XORKeyStream(dst[wrapNonceLen:total], payload)
-	return total, nil
-}
-
-// wrapPacket allocates a fresh wire buffer. Retained for tests and control
-// paths; hot-path code should call wrapInto with a reusable scratch buffer.
-func wrapPacket(key, payload []byte) ([]byte, error) {
-	out := make([]byte, wrapNonceLen+len(payload))
-	n, err := wrapInto(out, key, payload)
-	if err != nil {
-		return nil, err
-	}
-	return out[:n], nil
-}
-
-func unwrapPacket(key, wire, dst []byte) (int, error) {
-	if len(key) != wrapKeyLen {
-		return 0, fmt.Errorf("wrap: key must be %d bytes (got %d)", wrapKeyLen, len(key))
-	}
-	if len(wire) < wrapNonceLen {
-		return 0, errors.New("wrap: short packet (no nonce)")
-	}
-	nonce := wire[:wrapNonceLen]
-	ciphertext := wire[wrapNonceLen:]
-	if len(ciphertext) > len(dst) {
-		return 0, errors.New("wrap: dst buffer too small")
-	}
-	cipher, err := chacha20.NewUnauthenticatedCipher(key, nonce)
-	if err != nil {
-		return 0, fmt.Errorf("wrap: cipher init: %w", err)
-	}
-	cipher.XORKeyStream(dst[:len(ciphertext)], ciphertext)
-	return len(ciphertext), nil
-}
@@ -6,31 +6,94 @@ import (
 	"testing"
 )

-func TestWrapPacketRoundTrip(t *testing.T) {
+func TestWrapConnRoundTrip(t *testing.T) {
 	key := bytes.Repeat([]byte{0x42}, wrapKeyLen)
 	payload := []byte("dtls record bytes")

-	wire, err := wrapPacket(key, payload)
+	client, err := newWrapConn(key, false)
 	if err != nil {
-		t.Fatalf("wrapPacket returned error: %v", err)
+		t.Fatalf("newWrapConn(client): %v", err)
 	}
-	if len(wire) != len(payload)+wrapNonceLen {
-		t.Fatalf("wire len = %d, want %d", len(wire), len(payload)+wrapNonceLen)
+	server, err := newWrapConn(key, true)
+	if err != nil {
+		t.Fatalf("newWrapConn(server): %v", err)
+	}
+
+	// Client → Server
+	wire := make([]byte, wrapMaxWire(len(payload)))
+	n, err := client.wrapInto(wire, payload)
+	if err != nil {
+		t.Fatalf("wrapInto: %v", err)
+	}
+	wire = wire[:n]
+
+	// Check DTLS header mimicry
+	if wire[0] != 0x17 {
+		t.Fatalf("content_type = 0x%02X, want 0x17", wire[0])
+	}
+	if wire[1] != 0xFE || wire[2] != 0xFD {
+		t.Fatalf("version = 0x%02X%02X, want 0xFEFD", wire[1], wire[2])
 	}
 	if bytes.Contains(wire, payload) {
 		t.Fatalf("wrapped packet contains plaintext payload")
 	}

-	dst := make([]byte, len(payload))
-	n, err := unwrapPacket(key, wire, dst)
+	dst := make([]byte, 1600)
+	m, err := server.unwrapPacket(wire, dst)
 	if err != nil {
-		t.Fatalf("unwrapPacket returned error: %v", err)
+		t.Fatalf("unwrapPacket: %v", err)
 	}
-	if n != len(payload) {
-		t.Fatalf("unwrapped len = %d, want %d", n, len(payload))
+	if m != len(payload) {
+		t.Fatalf("unwrapped len = %d, want %d", m, len(payload))
 	}
-	if !bytes.Equal(dst[:n], payload) {
-		t.Fatalf("round trip mismatch: got %q want %q", dst[:n], payload)
+	if !bytes.Equal(dst[:m], payload) {
+		t.Fatalf("round trip mismatch: got %q want %q", dst[:m], payload)
+	}
+
+	// Server → Client
+	wire2 := make([]byte, wrapMaxWire(len(payload)))
+	n2, err := server.wrapInto(wire2, payload)
+	if err != nil {
+		t.Fatalf("server wrapInto: %v", err)
+	}
+	m2, err := client.unwrapPacket(wire2[:n2], dst)
+	if err != nil {
+		t.Fatalf("client unwrapPacket: %v", err)
+	}
+	if !bytes.Equal(dst[:m2], payload) {
+		t.Fatalf("server→client round trip mismatch")
+	}
+}
+
+func TestWrapPaddingHidesSize(t *testing.T) {
+	key := bytes.Repeat([]byte{0x42}, wrapKeyLen)
+	wc, _ := newWrapConn(key, false)
+
+	// Two payloads of different size should produce same wire size if in same bucket
+	p1 := make([]byte, 10)
+	p2 := make([]byte, 100)
+
+	w1 := make([]byte, wrapMaxWire(len(p1)))
+	w2 := make([]byte, wrapMaxWire(len(p2)))
+
+	n1, _ := wc.wrapInto(w1, p1)
+	n2, _ := wc.wrapInto(w2, p2)
+
+	if n1 != n2 {
+		t.Fatalf("different wire sizes for payloads in same bucket: %d vs %d", n1, n2)
+	}
+}
+
+func TestWrapEpochDirection(t *testing.T) {
+	key := bytes.Repeat([]byte{0x42}, wrapKeyLen)
+	client, _ := newWrapConn(key, false)
+	server, _ := newWrapConn(key, true)
+
+	if client.epoch&0x8000 != 0 {
+		t.Fatalf("client epoch MSB should be 0, got 0x%04X", client.epoch)
+	}
+	if server.epoch&0x8000 == 0 {
+		t.Fatalf("server epoch MSB should be 1, got 0x%04X", server.epoch)
 	}
 }

@@ -58,9 +121,29 @@ func TestDecodeWrapKeyRequiresValidKeyWhenEnabled(t *testing.T) {
 	}
 }

-func TestUnwrapPacketRejectsShortPacket(t *testing.T) {
+func TestUnwrapRejectsShortPacket(t *testing.T) {
 	key := bytes.Repeat([]byte{0x42}, wrapKeyLen)
-	if _, err := unwrapPacket(key, []byte("short"), make([]byte, 16)); err == nil {
+	wc, _ := newWrapConn(key, false)
+	if _, err := wc.unwrapPacket([]byte("short"), make([]byte, 16)); err == nil {
 		t.Fatalf("unwrapPacket accepted short packet")
 	}
 }
+
+func TestUnwrapRejectsTamperedPacket(t *testing.T) {
+	key := bytes.Repeat([]byte{0x42}, wrapKeyLen)
+	client, _ := newWrapConn(key, false)
+	server, _ := newWrapConn(key, true)
+
+	payload := []byte("integrity test")
+	wire := make([]byte, wrapMaxWire(len(payload)))
+	n, _ := client.wrapInto(wire, payload)
+	wire = wire[:n]
+
+	// Flip a bit in the ciphertext
+	wire[wrapHeaderLen+5] ^= 0xFF
+
+	dst := make([]byte, 1600)
+	if _, err := server.unwrapPacket(wire, dst); err == nil {
+		t.Fatalf("unwrapPacket accepted tampered packet")
+	}
+}
@@ -3,60 +3,109 @@
 package main

 import (
+	"crypto/cipher"
 	"crypto/rand"
+	"crypto/sha256"
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"net"
+	"sync"
 	"sync/atomic"
 	"time"

 	dtlsnet "github.com/pion/dtls/v3/pkg/net"
 	pionudp "github.com/pion/transport/v4/udp"
-	"golang.org/x/crypto/chacha20"
+	"golang.org/x/crypto/chacha20poly1305"
 )

+// Wire format is identical to the client — see client/wrap.go for the
+// full specification. Server uses epoch MSB=1, client uses MSB=0.
+
 const (
-	wrapNonceLen = 12
-	wrapKeyLen   = 32
+	wrapKeyLen    = 32
+	wrapHeaderLen = 13 // DTLS 1.2 record header
+	wrapTagLen    = 16 // Poly1305 tag
+	wrapLenPrefix = 2  // uint16 BE real payload length
 )

-// Nonce layout matches the client: [4B session prefix | 8B atomic counter].
-// Replaces a per-packet crypto/rand syscall (getrandom on Linux) with a
-// monotonic counter; (key, nonce) uniqueness preserved within the process.
-// Wire format unchanged.
-var (
-	nonceSessionPrefix [4]byte
-	nonceCounter       atomic.Uint64
-)
+var padBuckets = [...]int{128, 256, 512, 768, 1024, 1300}

-func init() {
-	_, _ = rand.Read(nonceSessionPrefix[:])
+// bufPool eliminates per-packet heap allocation on the hot read/write paths.
+var bufPool = sync.Pool{
+	New: func() any {
+		b := make([]byte, 1600)
+		return &b
+	},
 }

-func fillWrapNonce(dst []byte) {
-	_ = dst[wrapNonceLen-1]
-	copy(dst[:4], nonceSessionPrefix[:])
-	binary.BigEndian.PutUint64(dst[4:wrapNonceLen], nonceCounter.Add(1))
+func padPlaintextSize(payloadLen int) int {
+	needed := wrapLenPrefix + payloadLen
+	for _, b := range padBuckets {
+		if b >= needed {
+			return b
+		}
+	}
+	return needed
 }

-func listenWrapped(addr *net.UDPAddr, key []byte) (dtlsnet.PacketListener, error) {
+func wrapMaxWire(payloadLen int) int {
+	return wrapHeaderLen + padPlaintextSize(payloadLen) + wrapTagLen
+}
+
+// wrapState holds the AEAD instance and nonce mask shared by all connections
+// under one key, plus per-connection epoch/seq state.
+type wrapState struct {
+	aead      cipher.AEAD
+	nonceMask [12]byte
+}
+
+func newWrapState(key []byte) (*wrapState, error) {
 	if len(key) != wrapKeyLen {
 		return nil, fmt.Errorf("wrap: key must be %d bytes (got %d)", wrapKeyLen, len(key))
 	}
+	aead, err := chacha20poly1305.New(key)
+	if err != nil {
+		return nil, fmt.Errorf("wrap: aead init: %w", err)
+	}
+	h := sha256.Sum256(key)
+	var mask [12]byte
+	copy(mask[:], h[:12])
+	return &wrapState{aead: aead, nonceMask: mask}, nil
+}
+
+func (s *wrapState) deriveNonce(epoch uint16, seqNum uint64) [12]byte {
+	var explicit [12]byte
+	binary.BigEndian.PutUint16(explicit[2:4], epoch)
+	binary.BigEndian.PutUint16(explicit[4:6], uint16(seqNum>>32))
+	binary.BigEndian.PutUint32(explicit[6:10], uint32(seqNum))
+	var nonce [12]byte
+	for i := range 12 {
+		nonce[i] = s.nonceMask[i] ^ explicit[i]
+	}
+	return nonce
+}
+
+// --- Listener ---
+
+func listenWrapped(addr *net.UDPAddr, key []byte) (dtlsnet.PacketListener, error) {
+	ws, err := newWrapState(key)
+	if err != nil {
+		return nil, err
+	}
 	inner, err := pionudp.Listen("udp", addr)
 	if err != nil {
 		return nil, fmt.Errorf("wrap: udp listen: %w", err)
 	}
 	return &wrapPacketListener{
 		inner: dtlsnet.PacketListenerFromListener(inner),
-		key:   key,
+		ws:    ws,
 	}, nil
 }

 type wrapPacketListener struct {
 	inner dtlsnet.PacketListener
-	key   []byte
+	ws    *wrapState
 }

 func (l *wrapPacketListener) Accept() (net.PacketConn, net.Addr, error) {
@@ -64,47 +113,113 @@ func (l *wrapPacketListener) Accept() (net.PacketConn, net.Addr, error) {
 	if err != nil {
 		return pc, addr, err
 	}
-	return &wrapPacketConn{inner: pc, key: l.key}, addr, nil
+	// Server epoch: random with MSB=1
+	var eb [2]byte
+	_, _ = rand.Read(eb[:])
+	epoch := binary.BigEndian.Uint16(eb[:]) | 0x8000
+
+	return &wrapPacketConn{inner: pc, ws: l.ws, epoch: epoch}, addr, nil
 }

 func (l *wrapPacketListener) Close() error   { return l.inner.Close() }
 func (l *wrapPacketListener) Addr() net.Addr { return l.inner.Addr() }

+// --- Per-peer PacketConn ---
+
 type wrapPacketConn struct {
 	inner net.PacketConn
-	key   []byte
+	ws    *wrapState
+	epoch uint16
+	seq   atomic.Uint64
 }

 func (c *wrapPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
-	buf := make([]byte, len(p)+wrapNonceLen)
-	n, addr, err := c.inner.ReadFrom(buf)
+	bp := bufPool.Get().(*[]byte)
+	buf := *bp
+	if cap(buf) < len(p)+wrapHeaderLen+wrapTagLen {
+		buf = make([]byte, len(p)+wrapHeaderLen+wrapTagLen)
+		*bp = buf
+	}
+	defer bufPool.Put(bp)
+
+	n, addr, err := c.inner.ReadFrom(buf[:cap(buf)])
 	if err != nil {
 		return 0, addr, err
 	}
-	if n < wrapNonceLen {
-		return 0, addr, errors.New("wrap: short packet (no nonce)")
+	wire := buf[:n]
+	if len(wire) < wrapHeaderLen+wrapTagLen+wrapLenPrefix {
+		return 0, addr, errors.New("wrap: packet too short")
 	}
-	nonce := buf[:wrapNonceLen]
-	ciphertext := buf[wrapNonceLen:n]
-	if len(ciphertext) > len(p) {
-		return 0, addr, errors.New("wrap: read buffer too small")
+	if wire[0] != 0x17 {
+		return 0, addr, fmt.Errorf("wrap: unexpected content type 0x%02X", wire[0])
 	}
-	cipher, err := chacha20.NewUnauthenticatedCipher(c.key, nonce)
+
+	epoch := binary.BigEndian.Uint16(wire[3:5])
+	seqHigh := uint64(binary.BigEndian.Uint16(wire[5:7]))
+	seqLow := uint64(binary.BigEndian.Uint32(wire[7:11]))
+	seqNum := (seqHigh << 32) | seqLow
+	bodyLen := int(binary.BigEndian.Uint16(wire[11:13]))
+
+	if wrapHeaderLen+bodyLen > n {
+		return 0, addr, errors.New("wrap: truncated packet")
+	}
+
+	nonce := c.ws.deriveNonce(epoch, seqNum)
+	aad := wire[:wrapHeaderLen]
+	ciphertext := wire[wrapHeaderLen : wrapHeaderLen+bodyLen]
+
+	plain, err := c.ws.aead.Open(ciphertext[:0], nonce[:], ciphertext, aad)
 	if err != nil {
-		return 0, addr, fmt.Errorf("wrap: cipher init: %w", err)
+		return 0, addr, fmt.Errorf("wrap: AEAD open: %w", err)
 	}
-	cipher.XORKeyStream(p[:len(ciphertext)], ciphertext)
-	return len(ciphertext), addr, nil
+	if len(plain) < wrapLenPrefix {
+		return 0, addr, errors.New("wrap: plaintext too short")
+	}
+	realLen := int(binary.BigEndian.Uint16(plain[0:2]))
+	if realLen > len(plain)-wrapLenPrefix || realLen > len(p) {
+		return 0, addr, errors.New("wrap: payload length mismatch")
+	}
+	copy(p[:realLen], plain[wrapLenPrefix:wrapLenPrefix+realLen])
+	return realLen, addr, nil
 }

 func (c *wrapPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
-	out := make([]byte, wrapNonceLen+len(p))
-	fillWrapNonce(out[:wrapNonceLen])
-	cipher, err := chacha20.NewUnauthenticatedCipher(c.key, out[:wrapNonceLen])
-	if err != nil {
-		return 0, fmt.Errorf("wrap: cipher init: %w", err)
+	paddedLen := padPlaintextSize(len(p))
+	wireLen := wrapHeaderLen + paddedLen + wrapTagLen
+
+	bp := bufPool.Get().(*[]byte)
+	out := *bp
+	if cap(out) < wireLen {
+		out = make([]byte, wireLen)
+		*bp = out
 	}
-	cipher.XORKeyStream(out[wrapNonceLen:], p)
+	out = out[:wireLen]
+	defer bufPool.Put(bp)
+
+	seq := c.seq.Add(1) - 1
+
+	// DTLS record header
+	out[0] = 0x17
+	out[1] = 0xFE
+	out[2] = 0xFD
+	binary.BigEndian.PutUint16(out[3:5], c.epoch)
+	binary.BigEndian.PutUint16(out[5:7], uint16(seq>>32))
+	binary.BigEndian.PutUint32(out[7:11], uint32(seq))
+	binary.BigEndian.PutUint16(out[11:13], uint16(paddedLen+wrapTagLen))
+
+	// Plaintext: [2B len | payload | random padding]
+	plain := out[wrapHeaderLen : wrapHeaderLen+paddedLen]
+	binary.BigEndian.PutUint16(plain[0:2], uint16(len(p)))
+	copy(plain[wrapLenPrefix:], p)
+	padStart := wrapLenPrefix + len(p)
+	if padStart < paddedLen {
+		_, _ = rand.Read(plain[padStart:paddedLen])
+	}
+
+	nonce := c.ws.deriveNonce(c.epoch, seq)
+	aad := out[:wrapHeaderLen]
+	c.ws.aead.Seal(out[wrapHeaderLen:wrapHeaderLen], nonce[:], plain, aad)
+
 	if _, err := c.inner.WriteTo(out, addr); err != nil {
 		return 0, err
 	}