amadan

carlos / platform Public

relay: wrapped relay socket with counting, buckets and the EIP port gate (spec §3)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Paul Campbell pushed by paul@keymail.dev f2a5c75ecc480c248f55a867f81a85ab3ad27e17 parent 0853fd9
2 files changed, +190 −0
  • internal/relay/conn.go +97 −0
  • internal/relay/conn_test.go +93 −0
diff --git a/internal/relay/conn.go b/internal/relay/conn.go
new file mode 100644
index 000000000..a857c6789
--- /dev/null
+++ b/internal/relay/conn.go
@@ -0,0 +1,97 @@
+package relay
+
+import (
+ "net"
+ "sync"
+ "time"
+)
+
+// Binding is the identity an allocation was created under. Src is the
+// typed source key (typedSrc) pion will present on later requests.
+type Binding struct {
+ ID Identity
+ Expiry time.Time
+ Kid string
+ Incarnation string
+ Src string
+}
+
+// typedSrc keys the owner binding on transport AND address, because pion
+// tags every allocation's tuple as UDP even for TCP clients and a UDP and a
+// TCP client can share an address string.
+func typedSrc(a net.Addr) string {
+ switch a.(type) {
+ case *net.TCPAddr:
+ return "tcp:" + a.String()
+ default:
+ return "udp:" + a.String()
+ }
+}
+
+type portSet struct {
+ mu sync.Mutex
+ m map[int]bool
+}
+
+func (p *portSet) add(port int) { p.mu.Lock(); p.m[port] = true; p.mu.Unlock() }
+func (p *portSet) remove(port int) { p.mu.Lock(); delete(p.m, port); p.mu.Unlock() }
+func (p *portSet) has(port int) bool {
+ p.mu.Lock()
+ defer p.mu.Unlock()
+ return p.m[port]
+}
+
+// relayConn is the PacketConn the relay hands pion for one allocation.
+// Closing it is how the relay ends the allocation: pion's read loop deletes
+// the allocation when ReadFrom errors (allocation.go:301 in pion/turn
+// v4.1.4). WriteTo is the client→peer leg and is where forwarded bytes are
+// counted, the per-allocation and box buckets applied, and the EIP port
+// gate enforced. ReadFrom is not overridden: peer→relay traffic is counted
+// by the server only after pion forwards it (see Task 6), so this wrapper
+// counts nothing on read.
+type relayConn struct {
+ net.PacketConn
+ b *Binding
+ eip net.IP
+ ports *portSet
+ appBucket *Bucket
+ boxBucket *Bucket
+ onIn func(Identity, int)
+ onRejected func(Identity, int)
+ onClose func(*relayConn)
+ once sync.Once
+}
+
+func newRelayConn(inner net.PacketConn, b *Binding, eip net.IP, ports *portSet, appBucket, boxBucket *Bucket,
+ onIn func(Identity, int), onRejected func(Identity, int), onClose func(*relayConn)) *relayConn {
+ return &relayConn{PacketConn: inner, b: b, eip: eip, ports: ports, appBucket: appBucket, boxBucket: boxBucket,
+ onIn: onIn, onRejected: onRejected, onClose: onClose}
+}
+
+func (c *relayConn) WriteTo(p []byte, addr net.Addr) (int, error) {
+ if ua, ok := addr.(*net.UDPAddr); ok && c.eip != nil && ua.IP.Equal(c.eip) && !c.ports.has(ua.Port) {
+ if c.onRejected != nil {
+ c.onRejected(c.b.ID, len(p))
+ }
+ return len(p), nil // dropped: only live relay ports are reachable on our own EIP
+ }
+ if (c.appBucket != nil && !c.appBucket.Allow(len(p))) || (c.boxBucket != nil && !c.boxBucket.Allow(len(p))) {
+ return len(p), nil // dropped by a token bucket
+ }
+ n, err := c.PacketConn.WriteTo(p, addr)
+ if n > 0 && c.onIn != nil {
+ c.onIn(c.b.ID, n)
+ }
+ return n, err
+}
+
+func (c *relayConn) Close() error {
+ var err error
+ c.once.Do(func() {
+ err = c.PacketConn.Close()
+ if c.onClose != nil {
+ c.onClose(c)
+ }
+ })
+ return err
+}
diff --git a/internal/relay/conn_test.go b/internal/relay/conn_test.go
new file mode 100644
index 000000000..080f31730
--- /dev/null
+++ b/internal/relay/conn_test.go
@@ -0,0 +1,93 @@
+package relay
+
+import (
+ "net"
+ "sync"
+ "testing"
+ "time"
+)
+
+func TestRelayConnGatesEIPPorts(t *testing.T) {
+ inner := &fakePacketConn{}
+ ports := &portSet{m: map[int]bool{50000: true}}
+ eip := net.ParseIP("203.0.113.5")
+ now := time.Unix(1_800_000_000, 0)
+ clock := func() time.Time { return now }
+ var in int
+ c := newRelayConn(inner, &Binding{ID: Identity{"abc", "meet"}}, eip, ports, NewBucket(1<<40, clock), NewBucket(1<<40, clock),
+ func(_ Identity, n int) { in += n }, nil, func(*relayConn) {})
+ ok := &net.UDPAddr{IP: eip, Port: 50000}
+ bad := &net.UDPAddr{IP: eip, Port: 3478}
+ other := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444}
+ c.WriteTo([]byte("abc"), ok)
+ c.WriteTo([]byte("abc"), bad)
+ c.WriteTo([]byte("abcd"), other)
+ if len(inner.written) != 2 || in != 7 {
+ t.Fatalf("written=%d in=%d, want 2 packets and 7 bytes", len(inner.written), in)
+ }
+}
+
+func TestRelayConnBucketDrops(t *testing.T) {
+ inner := &fakePacketConn{}
+ now := time.Unix(1_800_000_000, 0)
+ clock := func() time.Time { return now }
+ c := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, NewBucket(80, clock), NewBucket(1<<40, clock), func(Identity, int) {}, nil, func(*relayConn) {})
+ dst := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444}
+ c.WriteTo(make([]byte, 10), dst) // 10 bytes = the whole second
+ c.WriteTo(make([]byte, 1), dst)
+ if len(inner.written) != 1 {
+ t.Fatalf("written = %d, want 1 (second dropped by bucket)", len(inner.written))
+ }
+}
+
+func TestRelayConnCloseOnce(t *testing.T) {
+ n := 0
+ c := newRelayConn(&fakePacketConn{}, &Binding{}, nil, &portSet{m: map[int]bool{}}, nil, nil, nil, nil, func(*relayConn) { n++ })
+ c.Close()
+ c.Close()
+ if n != 1 {
+ t.Fatalf("onClose ran %d times", n)
+ }
+}
+
+func TestTypedSrc(t *testing.T) {
+ u := &net.UDPAddr{IP: net.ParseIP("192.0.2.1"), Port: 5000}
+ tc := &net.TCPAddr{IP: net.ParseIP("192.0.2.1"), Port: 5000}
+ if typedSrc(u) == typedSrc(tc) {
+ t.Fatal("udp and tcp with the same address must not collide")
+ }
+}
+
+// fakePacketConn records writes; ReadFrom blocks until Close.
+type fakePacketConn struct {
+ mu sync.Mutex
+ written [][]byte
+ closed chan struct{}
+ once sync.Once
+}
+
+func (f *fakePacketConn) init() { f.once.Do(func() { f.closed = make(chan struct{}) }) }
+func (f *fakePacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
+ f.init()
+ <-f.closed
+ return 0, nil, net.ErrClosed
+}
+func (f *fakePacketConn) WriteTo(b []byte, a net.Addr) (int, error) {
+ f.mu.Lock()
+ defer f.mu.Unlock()
+ f.written = append(f.written, append([]byte(nil), b...))
+ return len(b), nil
+}
+func (f *fakePacketConn) Close() error {
+ f.init()
+ select {
+ case <-f.closed:
+ default:
+ close(f.closed)
+ }
+ return nil
+}
+func (f *fakePacketConn) LocalAddr() net.Addr { return &net.UDPAddr{IP: net.IPv4zero, Port: 50000} }
+func (f *fakePacketConn) SetDeadline(time.Time) error { return nil }
+func (f *fakePacketConn) SetReadDeadline(time.Time) error { return nil }
+func (f *fakePacketConn) SetWriteDeadline(time.Time) error { return nil }