54 lines
1.1 KiB
Go
54 lines
1.1 KiB
Go
package main
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type loginEntry struct {
|
|
window time.Time
|
|
count int
|
|
seen time.Time
|
|
}
|
|
|
|
type LoginLimiter struct {
|
|
mu sync.Mutex
|
|
entries map[string]loginEntry
|
|
limit int
|
|
window time.Duration
|
|
maxKeys int
|
|
}
|
|
|
|
func NewLoginLimiter(limit int, window time.Duration) *LoginLimiter {
|
|
return &LoginLimiter{entries: make(map[string]loginEntry), limit: limit, window: window, maxKeys: 10000}
|
|
}
|
|
|
|
func (limiter *LoginLimiter) Allow(key string, now time.Time) bool {
|
|
limiter.mu.Lock()
|
|
defer limiter.mu.Unlock()
|
|
_, known := limiter.entries[key]
|
|
if !known && len(limiter.entries) >= limiter.maxKeys {
|
|
cutoff := now.Add(-2 * limiter.window)
|
|
for item, entry := range limiter.entries {
|
|
if entry.seen.Before(cutoff) {
|
|
delete(limiter.entries, item)
|
|
}
|
|
}
|
|
if len(limiter.entries) >= limiter.maxKeys {
|
|
return false
|
|
}
|
|
}
|
|
entry := limiter.entries[key]
|
|
if entry.window.IsZero() || now.Sub(entry.window) >= limiter.window {
|
|
entry.window, entry.count = now, 0
|
|
}
|
|
entry.seen = now
|
|
if entry.count >= limiter.limit {
|
|
limiter.entries[key] = entry
|
|
return false
|
|
}
|
|
entry.count++
|
|
limiter.entries[key] = entry
|
|
return true
|
|
}
|