| 1 | package relay |
| 2 | |
| 3 | import ( |
| 4 | "net" |
| 5 | "sync" |
| 6 | "time" |
| 7 | ) |
| 8 | |
| 9 | // Binding is the identity an allocation was created under. Src is the |
| 10 | // typed source key (typedSrc) pion will present on later requests. |
| 11 | type Binding struct { |
| 12 | ID Identity |
| 13 | Expiry time.Time |
| 14 | Kid string |
| 15 | Incarnation string |
| 16 | Src string |
| 17 | } |
| 18 | |
| 19 | // typedSrc keys the owner binding on transport AND address, because pion |
| 20 | // tags every allocation's tuple as UDP even for TCP clients and a UDP and a |
| 21 | // TCP client can share an address string. |
| 22 | func typedSrc(a net.Addr) string { |
| 23 | switch a.(type) { |
| 24 | case *net.TCPAddr: |
| 25 | return "tcp:" + a.String() |
| 26 | default: |
| 27 | return "udp:" + a.String() |
| 28 | } |
| 29 | } |
| 30 | |
| 31 | type portSet struct { |
| 32 | mu sync.Mutex |
| 33 | m map[int]bool |
| 34 | } |
| 35 | |
| 36 | func (p *portSet) add(port int) { p.mu.Lock(); p.m[port] = true; p.mu.Unlock() } |
| 37 | func (p *portSet) remove(port int) { p.mu.Lock(); delete(p.m, port); p.mu.Unlock() } |
| 38 | func (p *portSet) has(port int) bool { |
| 39 | p.mu.Lock() |
| 40 | defer p.mu.Unlock() |
| 41 | return p.m[port] |
| 42 | } |
| 43 | |
| 44 | // relayConn is the PacketConn the relay hands pion for one allocation. |
| 45 | // Closing it is how the relay ends the allocation: pion's read loop deletes |
| 46 | // the allocation when ReadFrom errors (allocation.go:301 in pion/turn |
| 47 | // v4.1.4). WriteTo is the client→peer leg and is where forwarded bytes are |
| 48 | // counted, the per-allocation and box buckets applied, and the EIP port |
| 49 | // gate enforced. ReadFrom is not overridden: peer→relay traffic is counted |
| 50 | // by the server only after pion forwards it (see Task 6), so this wrapper |
| 51 | // counts nothing on read. |
| 52 | type relayConn struct { |
| 53 | net.PacketConn |
| 54 | b *Binding |
| 55 | eip net.IP |
| 56 | ports *portSet |
| 57 | appBucket *Bucket |
| 58 | boxBucket *Bucket |
| 59 | onIn func(Identity, int) |
| 60 | onRejected func(Identity, int) |
| 61 | onClose func(*relayConn) |
| 62 | once sync.Once |
| 63 | } |
| 64 | |
| 65 | func newRelayConn(inner net.PacketConn, b *Binding, eip net.IP, ports *portSet, appBucket, boxBucket *Bucket, |
| 66 | onIn func(Identity, int), onRejected func(Identity, int), onClose func(*relayConn)) *relayConn { |
| 67 | return &relayConn{PacketConn: inner, b: b, eip: eip, ports: ports, appBucket: appBucket, boxBucket: boxBucket, |
| 68 | onIn: onIn, onRejected: onRejected, onClose: onClose} |
| 69 | } |
| 70 | |
| 71 | func (c *relayConn) WriteTo(p []byte, addr net.Addr) (int, error) { |
| 72 | if ua, ok := addr.(*net.UDPAddr); ok && c.eip != nil && ua.IP.Equal(c.eip) && !c.ports.has(ua.Port) { |
| 73 | if c.onRejected != nil { |
| 74 | c.onRejected(c.b.ID, len(p)) |
| 75 | } |
| 76 | return len(p), nil // dropped: only live relay ports are reachable on our own EIP |
| 77 | } |
| 78 | // Check both buckets before spending either: if we spent on appBucket |
| 79 | // and boxBucket then refused, the packet is dropped but appBucket is |
| 80 | // left debited, throttling the app harder than its rate whenever the |
| 81 | // box is under pressure. |
| 82 | if c.appBucket != nil && !c.appBucket.Peek(len(p)) { |
| 83 | return len(p), nil // dropped by the app bucket |
| 84 | } |
| 85 | if c.boxBucket != nil && !c.boxBucket.Peek(len(p)) { |
| 86 | return len(p), nil // dropped by the box bucket |
| 87 | } |
| 88 | if c.appBucket != nil { |
| 89 | c.appBucket.Spend(len(p)) |
| 90 | } |
| 91 | if c.boxBucket != nil { |
| 92 | c.boxBucket.Spend(len(p)) |
| 93 | } |
| 94 | n, err := c.PacketConn.WriteTo(p, addr) |
| 95 | if n > 0 && c.onIn != nil { |
| 96 | c.onIn(c.b.ID, n) |
| 97 | } |
| 98 | return n, err |
| 99 | } |
| 100 | |
| 101 | func (c *relayConn) Close() error { |
| 102 | var err error |
| 103 | c.once.Do(func() { |
| 104 | err = c.PacketConn.Close() |
| 105 | if c.onClose != nil { |
| 106 | c.onClose(c) |
| 107 | } |
| 108 | }) |
| 109 | return err |
| 110 | } |
| 111 | |
amadan