cmroubao/backend-api/internal/transport/authcommon/attempt_limiter.go

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)