package authcommon import ( "errors" "math" "net" "strings" "sync" "time" ) type AttemptLimiter interface { Allow(string) (bool, time.Duration) Reset(string) } type attemptWindow struct { count int resetAt time.Time } type InMemoryAttemptLimiter struct { mu sync.Mutex maxAttempts int window time.Duration maxKeys int now func() time.Time attempts map[string]attemptWindow } func NewAttemptLimiter( maxAttempts int, window time.Duration, maxKeys int, ) (*InMemoryAttemptLimiter, error) { return newAttemptLimiter(maxAttempts, window, maxKeys, time.Now) } func newAttemptLimiter( maxAttempts int, window time.Duration, maxKeys int, now func() time.Time, ) (*InMemoryAttemptLimiter, error) { if maxAttempts < 1 || window <= 0 || maxKeys < 1 || now == nil { return nil, errors.New("attempt limiter configuration is invalid") } return &InMemoryAttemptLimiter{ maxAttempts: maxAttempts, window: window, maxKeys: maxKeys, now: now, attempts: make(map[string]attemptWindow), }, nil } func (limiter *InMemoryAttemptLimiter) Allow( key string, ) (bool, time.Duration) { now := limiter.now().UTC() key = strings.TrimSpace(key) if key == "" { key = "unknown" } limiter.mu.Lock() defer limiter.mu.Unlock() current, found := limiter.attempts[key] if found && !current.resetAt.After(now) { delete(limiter.attempts, key) found = false } if !found { if len(limiter.attempts) >= limiter.maxKeys { limiter.removeExpired(now) } if len(limiter.attempts) >= limiter.maxKeys { return false, limiter.window } limiter.attempts[key] = attemptWindow{ count: 1, resetAt: now.Add(limiter.window), } return true, 0 } if current.count >= limiter.maxAttempts { return false, current.resetAt.Sub(now) } current.count++ limiter.attempts[key] = current return true, 0 } func (limiter *InMemoryAttemptLimiter) Reset(key string) { limiter.mu.Lock() delete(limiter.attempts, strings.TrimSpace(key)) limiter.mu.Unlock() } func (limiter *InMemoryAttemptLimiter) removeExpired(now time.Time) { for key, current := range limiter.attempts { if !current.resetAt.After(now) { delete(limiter.attempts, key) } } } func LoginAttemptKey(scope, remoteAddress string) string { host, _, err := net.SplitHostPort(strings.TrimSpace(remoteAddress)) if err != nil { host = strings.TrimSpace(remoteAddress) } if parsed := net.ParseIP(host); parsed != nil { host = parsed.String() } if host == "" { host = "unknown" } return scope + ":" + host } func RetryAfterSeconds(wait time.Duration) int { seconds := int(math.Ceil(wait.Seconds())) if seconds < 1 { return 1 } return seconds } var _ AttemptLimiter = (*InMemoryAttemptLimiter)(nil)