| 1 | package relay |
| 2 | |
| 3 | import ( |
| 4 | "net" |
| 5 | "sync" |
| 6 | "testing" |
| 7 | "time" |
| 8 | ) |
| 9 | |
| 10 | func TestRelayConnGatesEIPPorts(t *testing.T) { |
| 11 | inner := &fakePacketConn{} |
| 12 | ports := &portSet{m: map[int]bool{50000: true}} |
| 13 | eip := net.ParseIP("203.0.113.5") |
| 14 | now := time.Unix(1_800_000_000, 0) |
| 15 | clock := func() time.Time { return now } |
| 16 | var in int |
| 17 | c := newRelayConn(inner, &Binding{ID: Identity{"abc", "meet"}}, eip, ports, NewBucket(1<<40, clock), NewBucket(1<<40, clock), |
| 18 | func(_ Identity, n int) { in += n }, nil, func(*relayConn) {}) |
| 19 | ok := &net.UDPAddr{IP: eip, Port: 50000} |
| 20 | bad := &net.UDPAddr{IP: eip, Port: 3478} |
| 21 | other := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444} |
| 22 | c.WriteTo([]byte("abc"), ok) |
| 23 | c.WriteTo([]byte("abc"), bad) |
| 24 | c.WriteTo([]byte("abcd"), other) |
| 25 | if len(inner.written) != 2 || in != 7 { |
| 26 | t.Fatalf("written=%d in=%d, want 2 packets and 7 bytes", len(inner.written), in) |
| 27 | } |
| 28 | } |
| 29 | |
| 30 | func TestRelayConnBucketDrops(t *testing.T) { |
| 31 | inner := &fakePacketConn{} |
| 32 | now := time.Unix(1_800_000_000, 0) |
| 33 | clock := func() time.Time { return now } |
| 34 | c := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, NewBucket(80, clock), NewBucket(1<<40, clock), func(Identity, int) {}, nil, func(*relayConn) {}) |
| 35 | dst := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444} |
| 36 | c.WriteTo(make([]byte, 10), dst) // 10 bytes = the whole second |
| 37 | c.WriteTo(make([]byte, 1), dst) |
| 38 | if len(inner.written) != 1 { |
| 39 | t.Fatalf("written = %d, want 1 (second dropped by bucket)", len(inner.written)) |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | func TestRelayConnBoxBucketRefusalDoesNotDebitApp(t *testing.T) { |
| 44 | now := time.Unix(1_800_000_000, 0) |
| 45 | clock := func() time.Time { return now } |
| 46 | appBucket := NewBucket(88, clock) // 11 bytes/s: exactly 1 byte of headroom after a 10-byte write |
| 47 | boxBucket := NewBucket(80, clock) // 10 bytes/s |
| 48 | inner := &fakePacketConn{} |
| 49 | c := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, appBucket, boxBucket, func(Identity, int) {}, nil, func(*relayConn) {}) |
| 50 | dst := &net.UDPAddr{IP: net.ParseIP("198.51.100.9"), Port: 4444} |
| 51 | |
| 52 | c.WriteTo(make([]byte, 10), dst) // spends the whole second on both buckets: app 11->1, box 10->0 |
| 53 | if len(inner.written) != 1 { |
| 54 | t.Fatalf("written = %d, want 1 after the first write", len(inner.written)) |
| 55 | } |
| 56 | |
| 57 | c.WriteTo(make([]byte, 1), dst) // box has 0 left: refused. With the bug, appBucket.Allow already spent app's last token here. |
| 58 | if len(inner.written) != 1 { |
| 59 | t.Fatalf("written = %d, want 1 (dropped by the box bucket)", len(inner.written)) |
| 60 | } |
| 61 | |
| 62 | // A fresh, effectively unlimited box bucket, but the SAME app bucket. |
| 63 | // If the earlier box refusal wrongly debited app (the bug), app is now |
| 64 | // at 0 and this write is refused too. Fixed, app still has its 1 token. |
| 65 | freshBox := NewBucket(1<<40, clock) |
| 66 | c2 := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, appBucket, freshBox, func(Identity, int) {}, nil, func(*relayConn) {}) |
| 67 | c2.WriteTo(make([]byte, 1), dst) |
| 68 | if len(inner.written) != 2 { |
| 69 | t.Fatalf("written = %d, want 2 (app bucket still had its last token after the box refusal)", len(inner.written)) |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | func TestRelayConnCloseOnce(t *testing.T) { |
| 74 | n := 0 |
| 75 | c := newRelayConn(&fakePacketConn{}, &Binding{}, nil, &portSet{m: map[int]bool{}}, nil, nil, nil, nil, func(*relayConn) { n++ }) |
| 76 | c.Close() |
| 77 | c.Close() |
| 78 | if n != 1 { |
| 79 | t.Fatalf("onClose ran %d times", n) |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | func TestTypedSrc(t *testing.T) { |
| 84 | u := &net.UDPAddr{IP: net.ParseIP("192.0.2.1"), Port: 5000} |
| 85 | tc := &net.TCPAddr{IP: net.ParseIP("192.0.2.1"), Port: 5000} |
| 86 | if typedSrc(u) == typedSrc(tc) { |
| 87 | t.Fatal("udp and tcp with the same address must not collide") |
| 88 | } |
| 89 | } |
| 90 | |
| 91 | // fakePacketConn records writes; ReadFrom blocks until Close. |
| 92 | type fakePacketConn struct { |
| 93 | mu sync.Mutex |
| 94 | written [][]byte |
| 95 | closed chan struct{} |
| 96 | once sync.Once |
| 97 | } |
| 98 | |
| 99 | func (f *fakePacketConn) init() { f.once.Do(func() { f.closed = make(chan struct{}) }) } |
| 100 | func (f *fakePacketConn) ReadFrom(b []byte) (int, net.Addr, error) { |
| 101 | f.init() |
| 102 | <-f.closed |
| 103 | return 0, nil, net.ErrClosed |
| 104 | } |
| 105 | func (f *fakePacketConn) WriteTo(b []byte, a net.Addr) (int, error) { |
| 106 | f.mu.Lock() |
| 107 | defer f.mu.Unlock() |
| 108 | f.written = append(f.written, append([]byte(nil), b...)) |
| 109 | return len(b), nil |
| 110 | } |
| 111 | func (f *fakePacketConn) Close() error { |
| 112 | f.init() |
| 113 | select { |
| 114 | case <-f.closed: |
| 115 | default: |
| 116 | close(f.closed) |
| 117 | } |
| 118 | return nil |
| 119 | } |
| 120 | func (f *fakePacketConn) LocalAddr() net.Addr { return &net.UDPAddr{IP: net.IPv4zero, Port: 50000} } |
| 121 | func (f *fakePacketConn) SetDeadline(time.Time) error { return nil } |
| 122 | func (f *fakePacketConn) SetReadDeadline(time.Time) error { return nil } |
| 123 | func (f *fakePacketConn) SetWriteDeadline(time.Time) error { return nil } |
| 124 | |
amadan