Kyle Isom
a60e5cb86a
- Fix 314 errcheck violations (blank identifier for unrecoverable errors) - Fix errorlint violation (errors.Is for io.EOF) - Remove unused serveL7Route test helper - Simplify Duration.Seconds() selectors in tests - Remove unnecessary fmt.Sprintf in test - Migrate exclusion rules from issues.exclusions to linters.exclusions (v2 schema) - Add gosec test exclusions (G115, G304, G402, G705) - Disable fieldalignment govet analyzer (optimization, not correctness) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
82 lines
1.7 KiB
Go
82 lines
1.7 KiB
Go
package firewall
|
|
|
|
import (
|
|
"net/netip"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type rateLimitEntry struct {
|
|
count atomic.Int64
|
|
start atomic.Int64 // UnixNano
|
|
}
|
|
|
|
type rateLimiter struct {
|
|
limit int64
|
|
window time.Duration
|
|
entries sync.Map // netip.Addr → *rateLimitEntry
|
|
now func() time.Time
|
|
done chan struct{}
|
|
}
|
|
|
|
func newRateLimiter(limit int64, window time.Duration) *rateLimiter {
|
|
rl := &rateLimiter{
|
|
limit: limit,
|
|
window: window,
|
|
now: time.Now,
|
|
done: make(chan struct{}),
|
|
}
|
|
|
|
go rl.cleanup()
|
|
return rl
|
|
}
|
|
|
|
// Allow checks whether the given address is within its rate limit window.
|
|
// Returns false if the address has exceeded the limit.
|
|
func (rl *rateLimiter) Allow(addr netip.Addr) bool {
|
|
now := rl.now().UnixNano()
|
|
|
|
val, _ := rl.entries.LoadOrStore(addr, &rateLimitEntry{})
|
|
entry, _ := val.(*rateLimitEntry)
|
|
|
|
windowStart := entry.start.Load()
|
|
if now-windowStart >= rl.window.Nanoseconds() {
|
|
// Window expired — reset. Intentionally non-atomic across the two
|
|
// stores; worst case a few extra connections slip through at the
|
|
// boundary, which is acceptable for a connection rate limiter.
|
|
entry.start.Store(now)
|
|
entry.count.Store(1)
|
|
return true
|
|
}
|
|
|
|
n := entry.count.Add(1)
|
|
return n <= rl.limit
|
|
}
|
|
|
|
// Stop terminates the cleanup goroutine.
|
|
func (rl *rateLimiter) Stop() {
|
|
close(rl.done)
|
|
}
|
|
|
|
func (rl *rateLimiter) cleanup() {
|
|
ticker := time.NewTicker(rl.window)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-rl.done:
|
|
return
|
|
case <-ticker.C:
|
|
cutoff := rl.now().Add(-2 * rl.window).UnixNano()
|
|
rl.entries.Range(func(key, value any) bool {
|
|
entry, _ := value.(*rateLimitEntry)
|
|
if entry.start.Load() < cutoff {
|
|
rl.entries.Delete(key)
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
}
|
|
}
|