| 1 | package discovery |
| 2 | |
| 3 | import ( |
| 4 | "strings" |
| 5 | "sync" |
| 6 | "time" |
| 7 | ) |
| 8 | |
| 9 | // Default per-source-IP rate-limit parameters for the announce endpoint. |
| 10 | // Honest relays announce every 30 seconds per bootstrap. A source IP may |
| 11 | // represent multiple relays behind the same NAT or proxy, so the default |
| 12 | // receiver budget leaves room for shared egress while still bounding abuse. |
| 13 | const ( |
| 14 | defaultAnnounceRatePerMinute = 30 |
| 15 | defaultAnnounceBurst = 60 |
| 16 | announceLimiterPruneInterval = 10 * time.Minute |
| 17 | announceLimiterIdleTTL = 30 * time.Minute |
| 18 | ) |
| 19 | |
| 20 | // AnnounceLimiter is a fixed-window-with-token-refill rate limiter keyed |
| 21 | // by source IP. Buckets that have not been touched for announceLimiterIdleTTL |
| 22 | // are garbage-collected on demand to keep memory bounded under a churning |
| 23 | // IP address space. |
| 24 | // |
| 25 | // AnnounceLimiter is safe for concurrent use. All state lives behind a |
| 26 | // single sync.Mutex; the limiter is on the announce hot path but each |
| 27 | // permit decision is O(1) (bucket lookup + arithmetic), and pruning is |
| 28 | // amortized over normal traffic so the worst-case latency is bounded. |
| 29 | type AnnounceLimiter struct { |
| 30 | mu sync.Mutex |
| 31 | buckets map[string]*announceBucket |
| 32 | ratePerMinute float64 |
| 33 | burst float64 |
| 34 | lastPrune time.Time |
| 35 | pruneInterval time.Duration |
| 36 | bucketIdleTTL time.Duration |
| 37 | clock func() time.Time // overridable for tests |
| 38 | maxBucketCount int |
| 39 | } |
| 40 | |
| 41 | type announceBucket struct { |
| 42 | tokens float64 |
| 43 | updatedAt time.Time |
| 44 | lastUsedAt time.Time |
| 45 | } |
| 46 | |
| 47 | // NewAnnounceLimiter constructs a limiter with the supplied sustained rate |
| 48 | // (requests per minute) and burst capacity. Non-positive values fall back |
| 49 | // to the defaults so the zero-config path is safe. |
| 50 | func NewAnnounceLimiter(ratePerMinute, burst int) *AnnounceLimiter { |
| 51 | if ratePerMinute <= 0 { |
| 52 | ratePerMinute = defaultAnnounceRatePerMinute |
| 53 | } |
| 54 | if burst <= 0 { |
| 55 | burst = defaultAnnounceBurst |
| 56 | } |
| 57 | return &AnnounceLimiter{ |
| 58 | buckets: make(map[string]*announceBucket), |
| 59 | ratePerMinute: float64(ratePerMinute), |
| 60 | burst: float64(burst), |
| 61 | pruneInterval: announceLimiterPruneInterval, |
| 62 | bucketIdleTTL: announceLimiterIdleTTL, |
| 63 | clock: func() time.Time { return time.Now() }, |
| 64 | maxBucketCount: 65536, |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | // Allow returns true if the supplied source IP has remaining capacity in |
| 69 | // its bucket and atomically deducts one token. Empty source IPs share a |
| 70 | // single anonymized bucket so a misconfigured proxy cannot bypass the |
| 71 | // limiter by suppressing client identification. |
| 72 | func (l *AnnounceLimiter) Allow(srcIP string) bool { |
| 73 | if l == nil { |
| 74 | return true |
| 75 | } |
| 76 | key := normalizeAnnounceLimiterKey(srcIP) |
| 77 | |
| 78 | l.mu.Lock() |
| 79 | defer l.mu.Unlock() |
| 80 | |
| 81 | now := l.clock() |
| 82 | l.maybePruneLocked(now) |
| 83 | |
| 84 | bucket, ok := l.buckets[key] |
| 85 | if !ok { |
| 86 | // New buckets start full so a single legitimate announce isn't |
| 87 | // gated by warm-up latency. |
| 88 | bucket = &announceBucket{ |
| 89 | tokens: l.burst, |
| 90 | updatedAt: now, |
| 91 | lastUsedAt: now, |
| 92 | } |
| 93 | // Hard ceiling: if the table is saturated, refuse new IPs rather |
| 94 | // than allow unbounded growth from random source addresses. |
| 95 | if len(l.buckets) >= l.maxBucketCount { |
| 96 | return false |
| 97 | } |
| 98 | l.buckets[key] = bucket |
| 99 | } else { |
| 100 | elapsed := now.Sub(bucket.updatedAt) |
| 101 | if elapsed > 0 { |
| 102 | bucket.tokens += (l.ratePerMinute * float64(elapsed)) / float64(time.Minute) |
| 103 | if bucket.tokens > l.burst { |
| 104 | bucket.tokens = l.burst |
| 105 | } |
| 106 | bucket.updatedAt = now |
| 107 | } |
| 108 | } |
| 109 | |
| 110 | bucket.lastUsedAt = now |
| 111 | if bucket.tokens < 1 { |
| 112 | return false |
| 113 | } |
| 114 | bucket.tokens-- |
| 115 | return true |
| 116 | } |
| 117 | |
| 118 | // maybePruneLocked drops idle buckets so the limiter's memory footprint |
| 119 | // stays proportional to the number of recently-active source IPs. The |
| 120 | // caller MUST already hold l.mu. |
| 121 | func (l *AnnounceLimiter) maybePruneLocked(now time.Time) { |
| 122 | if l.lastPrune.IsZero() { |
| 123 | l.lastPrune = now |
| 124 | return |
| 125 | } |
| 126 | if now.Sub(l.lastPrune) < l.pruneInterval { |
| 127 | return |
| 128 | } |
| 129 | l.lastPrune = now |
| 130 | threshold := now.Add(-l.bucketIdleTTL) |
| 131 | for key, bucket := range l.buckets { |
| 132 | if bucket.lastUsedAt.Before(threshold) { |
| 133 | delete(l.buckets, key) |
| 134 | } |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | func normalizeAnnounceLimiterKey(raw string) string { |
| 139 | key := strings.TrimSpace(raw) |
| 140 | if key == "" { |
| 141 | return "<unknown>" |
| 142 | } |
| 143 | return strings.ToLower(key) |
| 144 | } |