amadan

carlos / platform Public

relay: check both buckets before spending either (review fix)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Paul Campbell pushed by paul@keymail.dev dca6dbceef6eec3fccf7f2fbf9c79c72e57bd13c parent f2a5c75
3 files changed, +61 −2
  • internal/relay/conn.go +15 −2
  • internal/relay/conn_test.go +21 −0
  • internal/relay/limits.go +25 −0
diff --git a/internal/relay/conn.go b/internal/relay/conn.go
index a857c6789..14c1c4931 100644
--- a/internal/relay/conn.go
+++ b/internal/relay/conn.go
@@ -75,8 +75,21 @@ func (c *relayConn) WriteTo(p []byte, addr net.Addr) (int, error) {
}
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
+ // Check both buckets before spending either: if we spent on appBucket
+ // and boxBucket then refused, the packet is dropped but appBucket is
+ // left debited, throttling the app harder than its rate whenever the
+ // box is under pressure.
+ if c.appBucket != nil && !c.appBucket.Peek(len(p)) {
+ return len(p), nil // dropped by the app bucket
+ }
+ if c.boxBucket != nil && !c.boxBucket.Peek(len(p)) {
+ return len(p), nil // dropped by the box bucket
+ }
+ if c.appBucket != nil {
+ c.appBucket.Spend(len(p))
+ }
+ if c.boxBucket != nil {
+ c.boxBucket.Spend(len(p))
}
n, err := c.PacketConn.WriteTo(p, addr)
if n > 0 && c.onIn != nil {
diff --git a/internal/relay/conn_test.go b/internal/relay/conn_test.go
index 080f31730..6418bfbb4 100644
--- a/internal/relay/conn_test.go
+++ b/internal/relay/conn_test.go
@@ -40,6 +40,27 @@ 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)
+ 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
+ if len(inner.written) != 1 {
+ t.Fatalf("written = %d, want 1 (second dropped by the box bucket)", len(inner.written))
+ }
+ 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)
+ if len(inner.written) != 2 {
+ t.Fatalf("written = %d, want 2 (app bucket still had tokens after the box refusal)", 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++ })
diff --git a/internal/relay/limits.go b/internal/relay/limits.go
index e0e208132..661b71bbd 100644
--- a/internal/relay/limits.go
+++ b/internal/relay/limits.go
@@ -214,3 +214,28 @@ func (b *Bucket) Allow(nbytes int) bool {
b.tokens -= float64(nbytes)
return true
}
+
+// Peek refills the bucket and reports whether nbytes would fit, without
+// spending anything. Pair it with Spend so multiple buckets can be checked
+// before any of them is debited.
+func (b *Bucket) Peek(nbytes int) bool {
+ b.mu.Lock()
+ defer b.mu.Unlock()
+ t := b.now()
+ b.tokens += t.Sub(b.last).Seconds() * b.rate
+ if b.tokens > b.rate {
+ b.tokens = b.rate
+ }
+ b.last = t
+ return float64(nbytes) <= b.tokens
+}
+
+// Spend debits nbytes unconditionally, clamping at zero.
+func (b *Bucket) Spend(nbytes int) {
+ b.mu.Lock()
+ defer b.mu.Unlock()
+ b.tokens -= float64(nbytes)
+ if b.tokens < 0 {
+ b.tokens = 0
+ }
+}