diff --git a/ratelimit_shard.go b/ratelimit_shard.go new file mode 100644 index 0000000..321fdf9 --- /dev/null +++ b/ratelimit_shard.go @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package beacon + +import ( + "sync" + "time" +) + +// rateLimitShards is the shard count for the per-packet rate limiters. A +// power of two so the shard index is a mask. Sized well above the reader +// goroutine count (2*NumCPU) so contention on any one shard stays rare +// even under a full-fleet reconverge flood. +const rateLimitShards = 256 + +func shardU32(k uint32) uint32 { + return (k * 2654435761) & (rateLimitShards - 1) +} + +func shardStr(s string) uint32 { + var h uint32 = 2166136261 + for i := 0; i < len(s); i++ { + h ^= uint32(s[i]) + h *= 16777619 + } + return h & (rateLimitShards - 1) +} + +// discoverRateLimiter is a sharded nodeID -> last-allowed-time map. Each +// shard carries its own mutex, so discovers for different node ids proceed +// in parallel instead of serialising on one global lock. +type discoverRateLimiter struct { + shards [rateLimitShards]struct { + mu sync.Mutex + m map[uint32]time.Time + } +} + +func newDiscoverRateLimiter() *discoverRateLimiter { + rl := &discoverRateLimiter{} + for i := range rl.shards { + rl.shards[i].m = make(map[uint32]time.Time) + } + return rl +} + +// allow reports whether a discover endpoint update for nodeID is permitted +// now, recording the timestamp when it is. +func (rl *discoverRateLimiter) allow(nodeID uint32, minInterval time.Duration) bool { + sh := &rl.shards[shardU32(nodeID)] + sh.mu.Lock() + defer sh.mu.Unlock() + if last, ok := sh.m[nodeID]; ok && time.Since(last) < minInterval { + return false + } + sh.m[nodeID] = time.Now() + return true +} + +func (rl *discoverRateLimiter) sweep(cutoff time.Time) { + for i := range rl.shards { + sh := &rl.shards[i] + sh.mu.Lock() + for id, last := range sh.m { + if last.Before(cutoff) { + delete(sh.m, id) + } + } + sh.mu.Unlock() + } +} + +// punchRateLimiter is a sharded source-IP -> last-allowed-time map. +type punchRateLimiter struct { + shards [rateLimitShards]struct { + mu sync.Mutex + m map[string]time.Time + } +} + +func newPunchRateLimiter() *punchRateLimiter { + rl := &punchRateLimiter{} + for i := range rl.shards { + rl.shards[i].m = make(map[string]time.Time) + } + return rl +} + +func (rl *punchRateLimiter) allow(sourceKey string, minInterval time.Duration) bool { + sh := &rl.shards[shardStr(sourceKey)] + sh.mu.Lock() + defer sh.mu.Unlock() + if last, ok := sh.m[sourceKey]; ok && time.Since(last) < minInterval { + return false + } + sh.m[sourceKey] = time.Now() + return true +} + +func (rl *punchRateLimiter) sweep(cutoff time.Time) { + for i := range rl.shards { + sh := &rl.shards[i] + sh.mu.Lock() + for ip, last := range sh.m { + if last.Before(cutoff) { + delete(sh.m, ip) + } + } + sh.mu.Unlock() + } +} + +// relayRateLimiter is a sharded senderID -> sliding-window map. +type relayRateLimiter struct { + shards [rateLimitShards]struct { + mu sync.Mutex + m map[uint32]*relaySourceWindow + } +} + +func newRelayRateLimiter() *relayRateLimiter { + rl := &relayRateLimiter{} + for i := range rl.shards { + rl.shards[i].m = make(map[uint32]*relaySourceWindow) + } + return rl +} + +// allow reports whether a relay from senderID is within its per-second +// budget, advancing the window as needed. +func (rl *relayRateLimiter) allow(senderID uint32, nowNano int64, maxPerSecond uint32) bool { + sh := &rl.shards[shardU32(senderID)] + sh.mu.Lock() + defer sh.mu.Unlock() + w, ok := sh.m[senderID] + if !ok || nowNano-w.windowStart >= int64(time.Second) { + sh.m[senderID] = &relaySourceWindow{windowStart: nowNano, count: 1} + return true + } + if w.count >= maxPerSecond { + return false + } + w.count++ + return true +} + +func (rl *relayRateLimiter) sweep(cutoffNano int64) { + for i := range rl.shards { + sh := &rl.shards[i] + sh.mu.Lock() + for id, w := range sh.m { + if w.windowStart < cutoffNano { + delete(sh.m, id) + } + } + sh.mu.Unlock() + } +} diff --git a/server.go b/server.go index c410147..484f6f6 100644 --- a/server.go +++ b/server.go @@ -56,10 +56,10 @@ type Server struct { lastDropLog atomic.Int64 // UnixNano of last drop warning (rate limit) lastNotFoundLog atomic.Int64 // UnixNano of last not-found warning (rate limit) - // Punch-request rate limiters (SEC-026). - punchRateMu sync.Mutex // protects punchSourceLast - punchSourceLast map[string]time.Time // source IP → last allowed punch time - lastPunchTime atomic.Int64 // UnixNano of last global punch (rate limit) + // Punch-request rate limiters (SEC-026). Sharded per source IP so the + // per-source cap does not serialise all readers on one lock. + punchRL *punchRateLimiter + lastPunchTime atomic.Int64 // UnixNano of last global punch (rate limit) // PILOT-342: source-IP whitelist. Sources in this set bypass BOTH the // per-source and the global punch-rate caps. Written once at startup // via SetPunchWhitelist, read on every handlePunchRequest under @@ -83,13 +83,12 @@ type Server struct { // before Serve; race-free reads via atomic.Pointer. breakerAllow atomic.Pointer[func(name string) (bool, string)] - // Per-source relay rate limiters (SEC-037). - relayRateMu sync.Mutex // protects relaySourceCount - relaySourceCount map[uint32]*relaySourceWindow // senderID → sliding window state + // Per-source relay rate limiters (SEC-037). Sharded per senderID. + relayRL *relayRateLimiter - // Per-nodeID discover rate limiter (PILOT-334) — prevents endpoint flapping. - discoverRateMu sync.Mutex - discoverRateLast map[uint32]time.Time // nodeID → last allowed discover endpoint update + // Per-nodeID discover rate limiter (PILOT-334) — prevents endpoint + // flapping. Sharded per nodeID. + discoverRL *discoverRateLimiter // Peer mesh (gossip) beaconID uint32 @@ -184,9 +183,9 @@ func NewWithPeers(beaconID uint32, peers []string) *Server { relayCh: make(chan relayJob, relayQueueSize), beaconID: beaconID, done: make(chan struct{}), - punchSourceLast: make(map[string]time.Time), - relaySourceCount: make(map[uint32]*relaySourceWindow), - discoverRateLast: make(map[uint32]time.Time), + punchRL: newPunchRateLimiter(), + relayRL: newRelayRateLimiter(), + discoverRL: newDiscoverRateLimiter(), } emptyPeers := make(map[uint32]*net.UDPAddr) s.peerNodes.Store(&emptyPeers) @@ -627,16 +626,10 @@ func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr, wantDest bool) } // Per-nodeID endpoint update rate limit (PILOT-334) — prevents a single - // nodeID from flapping its endpoint via rapid Discover messages. - s.discoverRateMu.Lock() - if last, ok := s.discoverRateLast[nodeID]; ok && time.Since(last) < discoverMinInterval { - s.discoverRateMu.Unlock() - // Rate-limited: skip the Upsert but still reply with the - // observed address so the node learns its public endpoint. - } else { - s.discoverRateLast[nodeID] = time.Now() - s.discoverRateMu.Unlock() - // Record this node's observed public endpoint. Sharded — no global lock. + // nodeID from flapping its endpoint via rapid Discover messages. When + // rate-limited we skip the Upsert but still reply with the observed + // address so the node learns its public endpoint. + if s.discoverRL.allow(nodeID, discoverMinInterval) { if _, atCap := s.nodes.Upsert(nodeID, remote, time.Now(), maxBeaconNodes); atCap { return // shard at capacity — drop silently } @@ -743,14 +736,10 @@ func (s *Server) handlePunchRequest(data []byte, remote *net.UDPAddr) { return } // 2. Per-source cap: at most one punch per punchPerSourceInterval. - s.punchRateMu.Lock() - if last, ok := s.punchSourceLast[sourceKey]; ok && time.Since(last) < punchPerSourceInterval { - s.punchRateMu.Unlock() + if !s.punchRL.allow(sourceKey, punchPerSourceInterval) { return } - s.punchSourceLast[sourceKey] = time.Now() s.lastPunchTime.Store(now) - s.punchRateMu.Unlock() } rateLimitBypass: @@ -876,22 +865,12 @@ func (s *Server) dispatchRelay(data []byte) { // a known destination can saturate the 524288-deep relayCh at rates // far above normal — this cap gives each source a fixed share. now := time.Now().UnixNano() - s.relayRateMu.Lock() - w, ok := s.relaySourceCount[senderID] - if !ok || now-w.windowStart >= int64(time.Second) { - // New 1-second window. - s.relaySourceCount[senderID] = &relaySourceWindow{windowStart: now, count: 1} - s.relayRateMu.Unlock() - } else if w.count >= maxRelaysPerSourcePerSecond { - // Source exceeded per-second budget — silently drop. - // The sender's daemon retries (3-attempt path in - // pkg/daemon/daemon.go relay branch), so a drop here - // is eventually self-healing for honest senders. - s.relayRateMu.Unlock() + // Source exceeded per-second budget — silently drop. The sender's + // daemon retries (3-attempt path in pkg/daemon/daemon.go relay + // branch), so a drop here is eventually self-healing for honest + // senders. + if !s.relayRL.allow(senderID, now, maxRelaysPerSourcePerSecond) { return - } else { - w.count++ - s.relayRateMu.Unlock() } // Copy payload into a pooled buffer so we don't hold the read buffer @@ -1274,35 +1253,10 @@ func (s *Server) reapStaleNodes() { threshold := time.Now().Add(-beaconNodeTTL) s.nodes.ReapStale(threshold) - // Sweep stale punch-rate entries to prevent unbounded map growth. - s.punchRateMu.Lock() - cutoff := time.Now().Add(-punchRateCleanupInterval) - for ip, last := range s.punchSourceLast { - if last.Before(cutoff) { - delete(s.punchSourceLast, ip) - } - } - s.punchRateMu.Unlock() - - // Sweep stale relay-source entries. - s.relayRateMu.Lock() - cutoffNs := time.Now().Add(-relaySourceCleanupInterval).UnixNano() - for id, w := range s.relaySourceCount { - if w.windowStart < cutoffNs { - delete(s.relaySourceCount, id) - } - } - s.relayRateMu.Unlock() - - // Sweep stale discover-rate entries (PILOT-334). - s.discoverRateMu.Lock() - discoverCutoff := time.Now().Add(-discoverMinInterval * 2) - for id, last := range s.discoverRateLast { - if last.Before(discoverCutoff) { - delete(s.discoverRateLast, id) - } - } - s.discoverRateMu.Unlock() + // Sweep stale rate-limiter entries to prevent unbounded map growth. + s.punchRL.sweep(time.Now().Add(-punchRateCleanupInterval)) + s.relayRL.sweep(time.Now().Add(-relaySourceCleanupInterval).UnixNano()) + s.discoverRL.sweep(time.Now().Add(-discoverMinInterval * 2)) s.nodePubKeys.Range(func(k, _ interface{}) bool { id, ok := k.(uint32)