Wormhole/internal/auth/limiter.go

package auth

import (
	"sync"
	"time"

	"github.com/wormhole/wormhole/internal/config"
)

type LoginBlockedError struct {
	RetryAfter time.Duration
}

func (e *LoginBlockedError) Error() string {
	return "too many login attempts"
}

type loginRecord struct {
	failures    int
	windowStart time.Time
	lockedUntil time.Time
}

type LoginLimiter struct {
	mu          sync.Mutex
	maxAttempts int
	window      time.Duration
	lockout     time.Duration
	records     map[string]*loginRecord
}

func NewLoginLimiter(cfg *config.AuthConfig) *LoginLimiter {
	if cfg.LoginMaxAttempts <= 0 {
		return nil
	}
	max := cfg.LoginMaxAttempts
	window := cfg.LoginWindow
	lockout := cfg.LoginLockout
	if window <= 0 {
		window = 5 * time.Minute
	}
	if lockout <= 0 {
		lockout = 15 * time.Minute
	}
	return &LoginLimiter{
		maxAttempts: max,
		window:      window,
		lockout:     lockout,
		records:     make(map[string]*loginRecord),
	}
}

func (l *LoginLimiter) keys(ip, username string) []string {
	keys := []string{"ip:" + ip}
	if username != "" {
		keys = append(keys, "user:"+username)
	}
	return keys
}

func (l *LoginLimiter) IsBlocked(ip, username string) (time.Duration, bool) {
	if l == nil {
		return 0, false
	}
	now := time.Now()
	l.mu.Lock()
	defer l.mu.Unlock()

	var maxRetry time.Duration
	for _, key := range l.keys(ip, username) {
		rec := l.records[key]
		if rec == nil {
			continue
		}
		if now.Before(rec.lockedUntil) {
			retry := time.Until(rec.lockedUntil)
			if retry > maxRetry {
				maxRetry = retry
			}
		}
	}
	return maxRetry, maxRetry > 0
}

func (l *LoginLimiter) RecordFailure(ip, username string) {
	if l == nil {
		return
	}
	now := time.Now()
	l.mu.Lock()
	defer l.mu.Unlock()

	for _, key := range l.keys(ip, username) {
		rec := l.records[key]
		if rec == nil {
			rec = &loginRecord{windowStart: now}
			l.records[key] = rec
		}
		if now.Before(rec.lockedUntil) {
			continue
		}
		if now.Sub(rec.windowStart) > l.window {
			rec.failures = 0
			rec.windowStart = now
		}
		rec.failures++
		if rec.failures >= l.maxAttempts {
			rec.lockedUntil = now.Add(l.lockout)
			rec.failures = 0
		}
	}
}

func (l *LoginLimiter) Reset(ip, username string) {
	if l == nil {
		return
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	for _, key := range l.keys(ip, username) {
		delete(l.records, key)
	}
}