relay: make the bucket double-spend test discriminating (review fix 2)
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 file changed,
+15
−6
internal/relay/conn_test.go+15 −6
diff --git a/internal/relay/conn_test.go b/internal/relay/conn_test.go| index 6418bfbb4..8029ac70f 100644 |
| --- a/internal/relay/conn_test.go |
| +++ b/internal/relay/conn_test.go |
| @@ -43,21 +43,30 @@ func TestRelayConnBucketDrops(t *testing.T) { |
| func TestRelayConnBoxBucketRefusalDoesNotDebitApp(t *testing.T) { |
| now := time.Unix(1_800_000_000, 0) |
| clock := func() time.Time { return now } |
| - appBucket := NewBucket(1<<40, clock) |
| + appBucket := NewBucket(88, clock) // 11 bytes/s: exactly 1 byte of headroom after a 10-byte write |
| boxBucket := NewBucket(80, clock) // 10 bytes/s |
| inner := &fakePacketConn{} |
| c := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, appBucket, boxBucket, 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, forwards |
| - c.WriteTo(make([]byte, 1), dst) // box bucket is empty: dropped |
| + |
| + c.WriteTo(make([]byte, 10), dst) // spends the whole second on both buckets: app 11->1, box 10->0 |
| if len(inner.written) != 1 { |
| - t.Fatalf("written = %d, want 1 (second dropped by the box bucket)", len(inner.written)) |
| + t.Fatalf("written = %d, want 1 after the first write", len(inner.written)) |
| } |
| + |
| + c.WriteTo(make([]byte, 1), dst) // box has 0 left: refused. With the bug, appBucket.Allow already spent app's last token here. |
| + if len(inner.written) != 1 { |
| + t.Fatalf("written = %d, want 1 (dropped by the box bucket)", len(inner.written)) |
| + } |
| + |
| + // A fresh, effectively unlimited box bucket, but the SAME app bucket. |
| + // If the earlier box refusal wrongly debited app (the bug), app is now |
| + // at 0 and this write is refused too. Fixed, app still has its 1 token. |
| freshBox := NewBucket(1<<40, clock) |
| c2 := newRelayConn(inner, &Binding{}, nil, &portSet{m: map[int]bool{}}, appBucket, freshBox, func(Identity, int) {}, nil, func(*relayConn) {}) |
| - c2.WriteTo(make([]byte, 10), dst) |
| + c2.WriteTo(make([]byte, 1), dst) |
| if len(inner.written) != 2 { |
| - t.Fatalf("written = %d, want 2 (app bucket still had tokens after the box refusal)", len(inner.written)) |
| + t.Fatalf("written = %d, want 2 (app bucket still had its last token after the box refusal)", len(inner.written)) |
| } |
| } |