From 7c3334ed52769e54637ca6096194c2a311066c74 Mon Sep 17 00:00:00 2001 From: samosvalishe Date: Fri, 22 May 2026 17:09:40 +0700 Subject: [PATCH] refactor(wrap): mimic SRTP packets Port SRTP-like AEAD wrap from samosvalishe/vk-turn-proxy 0.11.0 and adapt it to the current client relay paths. Co-authored-by: samosvalishe --- client/main.go | 69 ++++++++++++++---- client/wrap.go | 166 +++++++++++++++++++++++++++++++++----------- client/wrap_test.go | 136 +++++++++++++++++++++++++++++++++--- server/main.go | 2 +- server/wrap.go | 152 ++++++++++++++++++++++++++++++++-------- 5 files changed, 432 insertions(+), 93 deletions(-) diff --git a/client/main.go b/client/main.go index 1310559..3f82a43 100644 --- a/client/main.go +++ b/client/main.go @@ -1945,10 +1945,25 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD }) var internalPipeAddr atomic.Value useWrap := len(turnParams.wrapKey) == wrapKeyLen + var wrapTX, wrapRX *wrapConn + if useWrap { + var wrapErr error + wrapTX, wrapErr = newWrapConn(turnParams.wrapKey, false) + if wrapErr != nil { + log.Printf("[STREAM %d] WRAP init failed: %v", streamID, wrapErr) + return + } + wrapRX, wrapErr = newWrapConn(turnParams.wrapKey, false) + if wrapErr != nil { + log.Printf("[STREAM %d] UNWRAP init failed: %v", streamID, wrapErr) + return + } + } go func() { defer turncancel() buf := make([]byte, 1600) + wrapBuf := make([]byte, wrapMaxWire(len(buf))) for { if turnctx.Err() != nil { return @@ -1965,12 +1980,12 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD out := buf[:n] if useWrap { - wrapped, wrapErr := wrapPacket(turnParams.wrapKey, out) + m, wrapErr := wrapTX.wrapInto(wrapBuf, out) if wrapErr != nil { log.Printf("[STREAM %d] WRAP failed: %v", streamID, wrapErr) return } - out = wrapped + out = wrapBuf[:m] } written, err1 := relayConn.WriteTo(out, peer) @@ -1986,7 +2001,7 @@ func oneTurnConnection(ctx context.Context, turnParams *turnParams, peer *net.UD defer turncancel() readBufLen := 1600 if useWrap { - readBufLen += wrapNonceLen + readBufLen = wrapMaxWire(readBufLen) } buf := make([]byte, readBufLen) plain := make([]byte, 1600) @@ -2003,7 +2018,7 @@ 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) + m, wrapErr := wrapRX.unwrapPacket(payload, plain) if wrapErr != nil { log.Printf("[STREAM %d] UNWRAP failed: %v (n=%d)", streamID, wrapErr, n) continue @@ -2148,7 +2163,7 @@ func main() { direct := flag.Bool("no-dtls", false, "connect without obfuscation. DO NOT USE") vlessMode := flag.Bool("vless", false, "VLESS mode: forward TCP connections (for VLESS) instead of UDP packets") vlessBond := flag.Bool("vless-bond", false, "bond one VLESS TCP connection across all active smux sessions") - wrapMode := flag.Bool("wrap", false, "WRAP mode: ChaCha20-XOR obfuscate DTLS packets before they reach TURN ChannelData") + wrapMode := flag.Bool("wrap", false, "WRAP mode: SRTP-like AEAD obfuscation for DTLS packets before they reach TURN ChannelData") wrapKeyHex := flag.String("wrap-key", "", "32-byte hex-encoded shared key for -wrap (64 hex chars)") genWrapKey := flag.Bool("gen-wrap-key", false, "print a fresh 64-character hex key for -wrap-key and exit") streamsPerCredFlag := flag.Int("streams-per-cred", streamsPerCache, "number of TURN streams sharing one VK credential cache") @@ -3017,7 +3032,11 @@ 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} + dtlsPC, err := newRelayPacketConn(relayConn, peer, tp.wrapKey) + if err != nil { + cleanup() + return nil, nil, err + } dtlsConn, err := dtls.ClientWithOptions(dtlsPC, peer, dtls.WithCertificates(certificate), dtls.WithInsecureSkipVerify(true), @@ -3066,22 +3085,41 @@ 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 wrapTX/wrapRX are set, packets are wrapped/unwrapped with SRTP-mimicry AEAD. type relayPacketConn struct { - relay net.PacketConn - peer net.Addr - wrapKey []byte + relay net.PacketConn + peer net.Addr + wrapTX *wrapConn + wrapRX *wrapConn +} + +func newRelayPacketConn(relay net.PacketConn, peer net.Addr, wrapKey []byte) (*relayPacketConn, error) { + r := &relayPacketConn{relay: relay, peer: peer} + if len(wrapKey) != wrapKeyLen { + return r, nil + } + var err error + r.wrapTX, err = newWrapConn(wrapKey, false) + if err != nil { + return nil, fmt.Errorf("wrap tx init: %w", err) + } + r.wrapRX, err = newWrapConn(wrapKey, false) + if err != nil { + return nil, fmt.Errorf("wrap rx init: %w", err) + } + return r, nil } func (r *relayPacketConn) ReadFrom(b []byte) (int, net.Addr, error) { - if len(r.wrapKey) != wrapKeyLen { + if r.wrapRX == 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.wrapRX.unwrapPacket(buf[:n], b) if err != nil { return 0, addr, err } @@ -3089,14 +3127,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.wrapTX == nil { return r.relay.WriteTo(b, r.peer) } - out, err := wrapPacket(r.wrapKey, b) + out := make([]byte, wrapMaxWire(len(b))) + n, err := r.wrapTX.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 diff --git a/client/wrap.go b/client/wrap.go index f271e13..996d4a4 100644 --- a/client/wrap.go +++ b/client/wrap.go @@ -3,19 +3,143 @@ package main import ( + "crypto/cipher" "crypto/rand" + "encoding/binary" "encoding/hex" "errors" "fmt" + "sync/atomic" - "golang.org/x/crypto/chacha20" + "golang.org/x/crypto/chacha20poly1305" ) +// Wire format - SRTP-like mimicry: +// +// [12B RTP header | 12B explicit nonce | AEAD ciphertext | 16B tag] +// +// RTP header (RFC 3550): +// +// byte 0: 0x80 V=2, P=0, X=0, CC=0 +// byte 1: 0x6F M=0, PT=111 (opus, typical voice PT) +// byte 2-3: seq16 BE monotonic, init random +// byte 4-7: ts32 BE monotonic, init random, increments by 960 (20ms @ 48kHz) +// byte 8-11: SSRC random per conn, MSB encodes direction +// +// 12B explicit nonce = 4B sessionID || 8B counter (BE). sessionID MSB +// matches SSRC MSB (direction bit). counter starts at a random uint64. +// AAD = first 24 bytes (RTP header || nonce). +// +// VK TURN appears to forward SRTP-shaped ChannelData on a fast path and +// drop anomalous payloads. AEAD ciphertext + 16B tag is plausible as +// AES-GCM SRTP per RFC 7714. + const ( - wrapNonceLen = 12 - wrapKeyLen = 32 + wrapKeyLen = 32 + wrapRTPHdrLen = 12 + wrapNonceLen = 12 + wrapTagLen = 16 + wrapHeaderLen = wrapRTPHdrLen + wrapNonceLen + wrapOverhead = wrapHeaderLen + wrapTagLen + wrapRTPVersion = 0x80 + wrapRTPPT = 0x6F + wrapTSStep = 960 ) +type wrapConn struct { + aead cipher.AEAD + sessionID [4]byte + ssrc [4]byte + counter atomic.Uint64 + seq atomic.Uint32 + timestamp atomic.Uint32 +} + +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) + } + w := &wrapConn{aead: aead} + + var rnd [16]byte + if _, err := rand.Read(rnd[:]); err != nil { + return nil, fmt.Errorf("wrap: rand init: %w", err) + } + copy(w.sessionID[:], rnd[0:4]) + copy(w.ssrc[:], rnd[4:8]) + if isServer { + w.sessionID[0] |= 0x80 + w.ssrc[0] |= 0x80 + } else { + w.sessionID[0] &^= 0x80 + w.ssrc[0] &^= 0x80 + } + w.seq.Store(uint32(binary.BigEndian.Uint16(rnd[8:10]))) + w.timestamp.Store(binary.BigEndian.Uint32(rnd[10:14])) + + var cb [8]byte + if _, err := rand.Read(cb[:]); err != nil { + return nil, fmt.Errorf("wrap: counter rand: %w", err) + } + w.counter.Store(binary.BigEndian.Uint64(cb[:])) + return w, nil +} + +func wrapMaxWire(payloadLen int) int { + return wrapOverhead + payloadLen +} + +func (w *wrapConn) wrapInto(dst, payload []byte) (int, error) { + wireLen := wrapOverhead + len(payload) + if len(dst) < wireLen { + return 0, errors.New("wrap: dst buffer too small") + } + + dst[0] = wrapRTPVersion + dst[1] = wrapRTPPT + seq := uint16(w.seq.Add(1) - 1) + binary.BigEndian.PutUint16(dst[2:4], seq) + ts := w.timestamp.Add(wrapTSStep) - wrapTSStep + binary.BigEndian.PutUint32(dst[4:8], ts) + copy(dst[8:12], w.ssrc[:]) + + noncePos := wrapRTPHdrLen + copy(dst[noncePos:noncePos+4], w.sessionID[:]) + ctr := w.counter.Add(1) - 1 + binary.BigEndian.PutUint64(dst[noncePos+4:noncePos+wrapNonceLen], ctr) + + nonce := dst[noncePos : noncePos+wrapNonceLen] + aad := dst[:wrapHeaderLen] + ctPos := wrapHeaderLen + copy(dst[ctPos:], payload) + w.aead.Seal(dst[ctPos:ctPos], nonce, dst[ctPos:ctPos+len(payload)], aad) + + return wireLen, nil +} + +func (w *wrapConn) unwrapPacket(wire, dst []byte) (int, error) { + if len(wire) < wrapOverhead { + return 0, errors.New("wrap: packet too short") + } + nonce := wire[wrapRTPHdrLen : wrapRTPHdrLen+wrapNonceLen] + aad := wire[:wrapHeaderLen] + ct := wire[wrapHeaderLen:] + + plain, err := w.aead.Open(ct[:0], nonce, ct, aad) + if err != nil { + return 0, fmt.Errorf("wrap: AEAD open: %w", err) + } + if len(plain) > len(dst) { + return 0, errors.New("wrap: dst buffer too small") + } + copy(dst[:len(plain)], plain) + return len(plain), nil +} + func genWrapKeyHex() (string, error) { key := make([]byte, wrapKeyLen) if _, err := rand.Read(key); err != nil { @@ -40,39 +164,3 @@ func decodeWrapKey(enabled bool, raw string) ([]byte, error) { } return key, nil } - -func wrapPacket(key, payload []byte) ([]byte, error) { - if len(key) != wrapKeyLen { - return nil, fmt.Errorf("wrap: key must be %d bytes (got %d)", wrapKeyLen, len(key)) - } - out := make([]byte, wrapNonceLen+len(payload)) - if _, err := rand.Read(out[:wrapNonceLen]); err != nil { - return nil, fmt.Errorf("wrap: nonce gen: %w", err) - } - cipher, err := chacha20.NewUnauthenticatedCipher(key, out[:wrapNonceLen]) - if err != nil { - return nil, fmt.Errorf("wrap: cipher init: %w", err) - } - cipher.XORKeyStream(out[wrapNonceLen:], payload) - return out, 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 -} diff --git a/client/wrap_test.go b/client/wrap_test.go index 8a0162d..7f90f7b 100644 --- a/client/wrap_test.go +++ b/client/wrap_test.go @@ -2,27 +2,43 @@ package main import ( "bytes" + "encoding/binary" "strings" "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) + } + + wire := make([]byte, wrapMaxWire(len(payload))) + n, err := client.wrapInto(wire, payload) + if err != nil { + t.Fatalf("wrapInto returned error: %v", err) + } + wire = wire[:n] + + if wire[0] != wrapRTPVersion { + t.Fatalf("RTP byte0 = 0x%02X, want 0x%02X", wire[0], wrapRTPVersion) + } + if wire[1] != wrapRTPPT { + t.Fatalf("RTP byte1 (PT) = 0x%02X, want 0x%02X", wire[1], wrapRTPPT) } if bytes.Contains(wire, payload) { t.Fatalf("wrapped packet contains plaintext payload") } dst := make([]byte, len(payload)) - n, err := unwrapPacket(key, wire, dst) + n, err = server.unwrapPacket(wire, dst) if err != nil { t.Fatalf("unwrapPacket returned error: %v", err) } @@ -34,6 +50,70 @@ func TestWrapPacketRoundTrip(t *testing.T) { } } +func TestWrapRTPHeaderProgression(t *testing.T) { + key := bytes.Repeat([]byte{0x42}, wrapKeyLen) + wc, err := newWrapConn(key, false) + if err != nil { + t.Fatalf("newWrapConn: %v", err) + } + payload := []byte("x") + + wire1 := make([]byte, wrapMaxWire(len(payload))) + n1, err := wc.wrapInto(wire1, payload) + if err != nil { + t.Fatalf("wrapInto 1: %v", err) + } + wire2 := make([]byte, wrapMaxWire(len(payload))) + n2, err := wc.wrapInto(wire2, payload) + if err != nil { + t.Fatalf("wrapInto 2: %v", err) + } + if n1 != n2 { + t.Fatalf("wire size variance: %d vs %d", n1, n2) + } + + seq1 := binary.BigEndian.Uint16(wire1[2:4]) + seq2 := binary.BigEndian.Uint16(wire2[2:4]) + if seq2 != seq1+1 { + t.Fatalf("seq did not increment: %d -> %d", seq1, seq2) + } + + ts1 := binary.BigEndian.Uint32(wire1[4:8]) + ts2 := binary.BigEndian.Uint32(wire2[4:8]) + if ts2-ts1 != wrapTSStep { + t.Fatalf("timestamp step = %d, want %d", ts2-ts1, wrapTSStep) + } + + if !bytes.Equal(wire1[8:12], wire2[8:12]) { + t.Fatalf("SSRC changed between packets") + } +} + +func TestWrapDirectionBit(t *testing.T) { + key := bytes.Repeat([]byte{0x42}, wrapKeyLen) + client, err := newWrapConn(key, false) + if err != nil { + t.Fatalf("newWrapConn(client): %v", err) + } + server, err := newWrapConn(key, true) + if err != nil { + t.Fatalf("newWrapConn(server): %v", err) + } + + if client.sessionID[0]&0x80 != 0 { + t.Fatalf("client sessionID MSB should be 0, got 0x%02X", client.sessionID[0]) + } + if server.sessionID[0]&0x80 == 0 { + t.Fatalf("server sessionID MSB should be 1, got 0x%02X", server.sessionID[0]) + } + if client.ssrc[0]&0x80 != 0 { + t.Fatalf("client SSRC MSB should be 0, got 0x%02X", client.ssrc[0]) + } + if server.ssrc[0]&0x80 == 0 { + t.Fatalf("server SSRC MSB should be 1, got 0x%02X", server.ssrc[0]) + } +} + func TestDecodeWrapKeyRequiresValidKeyWhenEnabled(t *testing.T) { if key, err := decodeWrapKey(false, ""); err != nil || key != nil { t.Fatalf("disabled decodeWrapKey = (%v, %v), want (nil, nil)", key, err) @@ -58,9 +138,49 @@ 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, err := newWrapConn(key, false) + if err != nil { + t.Fatalf("newWrapConn: %v", err) + } + 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, err := newWrapConn(key, false) + if err != nil { + t.Fatalf("newWrapConn(client): %v", err) + } + server, err := newWrapConn(key, true) + if err != nil { + t.Fatalf("newWrapConn(server): %v", err) + } + + payload := []byte("integrity test") + wire := make([]byte, wrapMaxWire(len(payload))) + n, err := client.wrapInto(wire, payload) + if err != nil { + t.Fatalf("wrapInto: %v", err) + } + wire = wire[:n] + wire[wrapHeaderLen+1] ^= 0xFF + + dst := make([]byte, 1600) + if _, err := server.unwrapPacket(wire, dst); err == nil { + t.Fatalf("unwrapPacket accepted tampered ciphertext") + } + + n, err = client.wrapInto(wire, payload) + if err != nil { + t.Fatalf("wrapInto: %v", err) + } + wire = wire[:n] + wire[8] ^= 0x01 + if _, err := server.unwrapPacket(wire, dst); err == nil { + t.Fatalf("unwrapPacket accepted tampered AAD") + } +} diff --git a/server/main.go b/server/main.go index 1ff83f4..e4f2e0e 100644 --- a/server/main.go +++ b/server/main.go @@ -36,7 +36,7 @@ func main() { connect := flag.String("connect", "", "connect to ip:port") vlessMode := flag.Bool("vless", false, "VLESS mode: forward TCP connections (for VLESS) instead of UDP packets") vlessBond := flag.Bool("vless-bond", false, "bond one VLESS TCP connection across all active smux sessions") - wrapMode := flag.Bool("wrap", false, "WRAP mode: ChaCha20-XOR obfuscate DTLS packets before they reach TURN ChannelData") + wrapMode := flag.Bool("wrap", false, "WRAP mode: SRTP-like AEAD obfuscation for DTLS packets before they reach TURN ChannelData") wrapKeyHex := flag.String("wrap-key", "", "32-byte hex-encoded shared key for -wrap (64 hex chars)") genWrapKey := flag.Bool("gen-wrap-key", false, "print a fresh 64-character hex key for -wrap-key and exit") debugFlag := flag.Bool("debug", false, "enable debug logging") diff --git a/server/wrap.go b/server/wrap.go index be3d62a..1a0917b 100644 --- a/server/wrap.go +++ b/server/wrap.go @@ -3,39 +3,76 @@ package main import ( + "crypto/cipher" "crypto/rand" + "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 client. Server sets the MSB of sessionID/SSRC; +// client clears it. RTP header fields are per-conn. + const ( - wrapNonceLen = 12 - wrapKeyLen = 32 + wrapKeyLen = 32 + wrapRTPHdrLen = 12 + wrapNonceLen = 12 + wrapTagLen = 16 + wrapHeaderLen = wrapRTPHdrLen + wrapNonceLen + wrapOverhead = wrapHeaderLen + wrapTagLen + wrapRTPVersion = 0x80 + wrapRTPPT = 0x6F + wrapTSStep = 960 ) -func listenWrapped(addr *net.UDPAddr, key []byte) (dtlsnet.PacketListener, error) { +var bufPool = sync.Pool{ + New: func() any { + b := make([]byte, 1600+wrapOverhead) + return &b + }, +} + +type wrapState struct { + aead cipher.AEAD +} + +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) + } + return &wrapState{aead: aead}, nil +} + +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) { @@ -43,49 +80,104 @@ 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 + c := &wrapPacketConn{inner: pc, ws: l.ws} + + var rnd [16]byte + if _, err := rand.Read(rnd[:]); err != nil { + return nil, addr, fmt.Errorf("wrap: rand init: %w", err) + } + copy(c.sessionID[:], rnd[0:4]) + copy(c.ssrc[:], rnd[4:8]) + c.sessionID[0] |= 0x80 + c.ssrc[0] |= 0x80 + c.seq.Store(uint32(binary.BigEndian.Uint16(rnd[8:10]))) + c.timestamp.Store(binary.BigEndian.Uint32(rnd[10:14])) + + var cb [8]byte + if _, err := rand.Read(cb[:]); err != nil { + return nil, addr, fmt.Errorf("wrap: counter rand: %w", err) + } + c.counter.Store(binary.BigEndian.Uint64(cb[:])) + return c, addr, nil } func (l *wrapPacketListener) Close() error { return l.inner.Close() } func (l *wrapPacketListener) Addr() net.Addr { return l.inner.Addr() } type wrapPacketConn struct { - inner net.PacketConn - key []byte + inner net.PacketConn + ws *wrapState + sessionID [4]byte + ssrc [4]byte + counter atomic.Uint64 + seq atomic.Uint32 + timestamp atomic.Uint32 } 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 + need := len(p) + wrapOverhead + if cap(buf) < need { + buf = make([]byte, need) + *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) < wrapOverhead { + 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") - } - cipher, err := chacha20.NewUnauthenticatedCipher(c.key, nonce) + nonce := wire[wrapRTPHdrLen : wrapRTPHdrLen+wrapNonceLen] + aad := wire[:wrapHeaderLen] + ct := wire[wrapHeaderLen:] + + plain, err := c.ws.aead.Open(ct[:0], nonce, ct, aad) if err != nil { - return 0, addr, fmt.Errorf("wrap: cipher init: %w", err) + return 0, addr, fmt.Errorf("wrap: AEAD open: %w", err) + } + if len(plain) > len(p) { + return 0, addr, errors.New("wrap: dst buffer too small") } - cipher.XORKeyStream(p[:len(ciphertext)], ciphertext) - return len(ciphertext), addr, nil + copy(p[:len(plain)], plain) + return len(plain), addr, nil } func (c *wrapPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) { - out := make([]byte, wrapNonceLen+len(p)) - if _, err := rand.Read(out[:wrapNonceLen]); err != nil { - return 0, fmt.Errorf("wrap: nonce gen: %w", err) - } - cipher, err := chacha20.NewUnauthenticatedCipher(c.key, out[:wrapNonceLen]) - if err != nil { - return 0, fmt.Errorf("wrap: cipher init: %w", err) + wireLen := wrapOverhead + len(p) + + 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) + + out[0] = wrapRTPVersion + out[1] = wrapRTPPT + seq := uint16(c.seq.Add(1) - 1) + binary.BigEndian.PutUint16(out[2:4], seq) + ts := c.timestamp.Add(wrapTSStep) - wrapTSStep + binary.BigEndian.PutUint32(out[4:8], ts) + copy(out[8:12], c.ssrc[:]) + + noncePos := wrapRTPHdrLen + copy(out[noncePos:noncePos+4], c.sessionID[:]) + ctr := c.counter.Add(1) - 1 + binary.BigEndian.PutUint64(out[noncePos+4:noncePos+wrapNonceLen], ctr) + + nonce := out[noncePos : noncePos+wrapNonceLen] + aad := out[:wrapHeaderLen] + ctPos := wrapHeaderLen + copy(out[ctPos:], p) + c.ws.aead.Seal(out[ctPos:ctPos], nonce, out[ctPos:ctPos+len(p)], aad) + if _, err := c.inner.WriteTo(out, addr); err != nil { return 0, err }