geo/server/internal/shared/auth/login_guard.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"math"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/redis/go-redis/v9"
)

type LoginGuardBlockReason string

const (
	LoginGuardBlockedRateLimited LoginGuardBlockReason = "login_rate_limited"
	LoginGuardBlockedLocked      LoginGuardBlockReason = "login_locked"
	LoginGuardBlockedInProgress  LoginGuardBlockReason = "login_in_progress"
)

var ErrLoginGuardUnavailable = errors.New("login guard unavailable")

type LoginGuardBlockedError struct {
	Reason     LoginGuardBlockReason
	RetryAfter time.Duration
}

func (e *LoginGuardBlockedError) Error() string {
	return string(e.Reason)
}

func (e *LoginGuardBlockedError) RetryAfterSeconds() int {
	if e.RetryAfter <= 0 {
		return 1
	}
	return max(1, int(math.Ceil(e.RetryAfter.Seconds())))
}

type LoginGuardConfig struct {
	Scope string

	Enabled bool

	RequestIPLimit           int
	RequestIPWindow          time.Duration
	RequestIdentifierLimit   int
	RequestIdentifierWindow  time.Duration
	RequestPairLimit         int
	RequestPairWindow        time.Duration
	FailureIdentifierLimit   int
	FailurePairLimit         int
	FailureWindow            time.Duration
	FailureLockBaseTTL       time.Duration
	FailureLockMaxTTL        time.Duration
	ConcurrentAttemptLockTTL time.Duration
}

func DefaultLoginGuardConfig(scope string) LoginGuardConfig {
	cfg := LoginGuardConfig{
		Scope:                    scope,
		Enabled:                  true,
		RequestIPLimit:           60,
		RequestIPWindow:          time.Minute,
		RequestIdentifierLimit:   12,
		RequestIdentifierWindow:  5 * time.Minute,
		RequestPairLimit:         6,
		RequestPairWindow:        time.Minute,
		FailureIdentifierLimit:   8,
		FailurePairLimit:         5,
		FailureWindow:            5 * time.Minute,
		FailureLockBaseTTL:       2 * time.Minute,
		FailureLockMaxTTL:        time.Hour,
		ConcurrentAttemptLockTTL: 15 * time.Second,
	}

	if strings.EqualFold(strings.TrimSpace(scope), "ops") {
		cfg.RequestIPLimit = 30
		cfg.RequestIdentifierLimit = 8
		cfg.RequestPairLimit = 6
		cfg.FailureIdentifierLimit = 10
		cfg.FailurePairLimit = 5
		cfg.FailureLockBaseTTL = 2 * time.Minute
		cfg.FailureLockMaxTTL = 30 * time.Minute
	}

	return cfg
}

type LoginGuard struct {
	rdb      *redis.Client
	cfg      LoginGuardConfig
	fallback *memoryLoginGuard
}

func NewLoginGuard(rdb *redis.Client, cfg LoginGuardConfig) *LoginGuard {
	cfg = normalizeLoginGuardConfig(cfg)
	return &LoginGuard{
		rdb:      rdb,
		cfg:      cfg,
		fallback: newMemoryLoginGuard(cfg),
	}
}

type LoginGuardSubject struct {
	Identifier string
	IP         string
}

type LoginPermit struct {
	guard          *LoginGuard
	keys           loginGuardKeys
	token          string
	fallbackPermit *memoryLoginPermit
	done           atomic.Bool
}

func (g *LoginGuard) Acquire(ctx context.Context, subject LoginGuardSubject) (*LoginPermit, error) {
	if g == nil || !g.cfg.Enabled {
		return nil, nil
	}

	keys := g.keysFor(subject)
	token, err := randomLoginGuardToken()
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrLoginGuardUnavailable, err)
	}
	if g.rdb == nil {
		return g.acquireFallback(keys, token)
	}

	result, err := loginGuardAcquireScript.Run(
		ctx,
		g.rdb,
		[]string{
			keys.lockIdentifier,
			keys.lockPair,
			keys.inflightIdentifier,
			keys.requestIP,
			keys.requestIdentifier,
			keys.requestPair,
		},
		token,
		g.cfg.ConcurrentAttemptLockTTL.Milliseconds(),
		g.cfg.RequestIPLimit,
		g.cfg.RequestIPWindow.Milliseconds(),
		g.cfg.RequestIdentifierLimit,
		g.cfg.RequestIdentifierWindow.Milliseconds(),
		g.cfg.RequestPairLimit,
		g.cfg.RequestPairWindow.Milliseconds(),
	).Slice()
	if err != nil {
		return g.acquireFallback(keys, token)
	}
	if len(result) < 3 {
		return nil, fmt.Errorf("%w: unexpected acquire result", ErrLoginGuardUnavailable)
	}

	allowed, _ := result[0].(int64)
	reason, _ := result[1].(string)
	retryAfter := redisMillisToDuration(result[2])
	if allowed == 1 {
		return &LoginPermit{guard: g, keys: keys, token: token}, nil
	}

	switch LoginGuardBlockReason(reason) {
	case LoginGuardBlockedLocked, LoginGuardBlockedInProgress, LoginGuardBlockedRateLimited:
		return nil, &LoginGuardBlockedError{
			Reason:     LoginGuardBlockReason(reason),
			RetryAfter: retryAfter,
		}
	default:
		return nil, fmt.Errorf("%w: unexpected block reason %q", ErrLoginGuardUnavailable, reason)
	}
}

func (g *LoginGuard) acquireFallback(keys loginGuardKeys, token string) (*LoginPermit, error) {
	if g.fallback == nil {
		return nil, fmt.Errorf("%w: fallback guard missing", ErrLoginGuardUnavailable)
	}
	permit, err := g.fallback.acquire(keys, token)
	if err != nil {
		return nil, err
	}
	return &LoginPermit{
		guard:          g,
		keys:           keys,
		token:          token,
		fallbackPermit: permit,
	}, nil
}

func (p *LoginPermit) RecordFailure(ctx context.Context) {
	if p == nil || p.guard == nil || !p.done.CompareAndSwap(false, true) {
		return
	}
	if p.fallbackPermit != nil {
		p.guard.fallback.recordFailure(p.fallbackPermit)
		return
	}
	_, _ = loginGuardFailureScript.Run(
		ctx,
		p.guard.rdb,
		[]string{
			p.keys.failureIdentifier,
			p.keys.failurePair,
			p.keys.lockIdentifier,
			p.keys.lockPair,
			p.keys.inflightIdentifier,
		},
		p.token,
		p.guard.cfg.FailureWindow.Milliseconds(),
		p.guard.cfg.FailureIdentifierLimit,
		p.guard.cfg.FailurePairLimit,
		p.guard.cfg.FailureLockBaseTTL.Milliseconds(),
		p.guard.cfg.FailureLockMaxTTL.Milliseconds(),
	).Result()
}

func (p *LoginPermit) RecordSuccess(ctx context.Context) {
	if p == nil || p.guard == nil || !p.done.CompareAndSwap(false, true) {
		return
	}
	if p.fallbackPermit != nil {
		p.guard.fallback.recordSuccess(p.fallbackPermit)
		return
	}
	_, _ = loginGuardSuccessScript.Run(
		ctx,
		p.guard.rdb,
		[]string{
			p.keys.failureIdentifier,
			p.keys.failurePair,
			p.keys.lockIdentifier,
			p.keys.lockPair,
			p.keys.inflightIdentifier,
		},
		p.token,
	).Result()
}

func (p *LoginPermit) Release(ctx context.Context) {
	if p == nil || p.guard == nil || !p.done.CompareAndSwap(false, true) {
		return
	}
	if p.fallbackPermit != nil {
		p.guard.fallback.release(p.fallbackPermit)
		return
	}
	_, _ = loginGuardReleaseScript.Run(ctx, p.guard.rdb, []string{p.keys.inflightIdentifier}, p.token).Result()
}

func (g *LoginGuard) Unlock(ctx context.Context, identifier string) error {
	if g == nil || !g.cfg.Enabled {
		return nil
	}
	keys := g.keysFor(LoginGuardSubject{Identifier: identifier})
	if g.rdb == nil {
		g.fallback.unlock(keys)
		return nil
	}
	return g.rdb.Del(ctx, keys.lockIdentifier, keys.failureIdentifier).Err()
}

type loginGuardKeys struct {
	requestIP          string
	requestIdentifier  string
	requestPair        string
	failureIdentifier  string
	failurePair        string
	lockIdentifier     string
	lockPair           string
	inflightIdentifier string
}

func (g *LoginGuard) keysFor(subject LoginGuardSubject) loginGuardKeys {
	scope := normalizeLoginGuardScope(g.cfg.Scope)
	identifier := normalizeLoginIdentifier(subject.Identifier)
	ip := normalizeLoginIP(subject.IP)
	identifierHash := hashLoginGuardKey(identifier)
	ipHash := hashLoginGuardKey(ip)
	pairHash := hashLoginGuardKey(ip + "|" + identifier)
	prefix := "auth:login:" + scope + ":"

	return loginGuardKeys{
		requestIP:          prefix + "req:ip:" + ipHash,
		requestIdentifier:  prefix + "req:id:" + identifierHash,
		requestPair:        prefix + "req:pair:" + pairHash,
		failureIdentifier:  prefix + "fail:id:" + identifierHash,
		failurePair:        prefix + "fail:pair:" + pairHash,
		lockIdentifier:     prefix + "lock:id:" + identifierHash,
		lockPair:           prefix + "lock:pair:" + pairHash,
		inflightIdentifier: prefix + "inflight:id:" + identifierHash,
	}
}

func normalizeLoginGuardConfig(cfg LoginGuardConfig) LoginGuardConfig {
	defaults := DefaultLoginGuardConfig(cfg.Scope)
	if strings.TrimSpace(cfg.Scope) == "" {
		cfg.Scope = defaults.Scope
	}
	if cfg.RequestIPLimit <= 0 {
		cfg.RequestIPLimit = defaults.RequestIPLimit
	}
	if cfg.RequestIPWindow <= 0 {
		cfg.RequestIPWindow = defaults.RequestIPWindow
	}
	if cfg.RequestIdentifierLimit <= 0 {
		cfg.RequestIdentifierLimit = defaults.RequestIdentifierLimit
	}
	if cfg.RequestIdentifierWindow <= 0 {
		cfg.RequestIdentifierWindow = defaults.RequestIdentifierWindow
	}
	if cfg.RequestPairLimit <= 0 {
		cfg.RequestPairLimit = defaults.RequestPairLimit
	}
	if cfg.RequestPairWindow <= 0 {
		cfg.RequestPairWindow = defaults.RequestPairWindow
	}
	if cfg.FailureIdentifierLimit <= 0 {
		cfg.FailureIdentifierLimit = defaults.FailureIdentifierLimit
	}
	if cfg.FailurePairLimit <= 0 {
		cfg.FailurePairLimit = defaults.FailurePairLimit
	}
	if cfg.FailureWindow <= 0 {
		cfg.FailureWindow = defaults.FailureWindow
	}
	if cfg.FailureLockBaseTTL <= 0 {
		cfg.FailureLockBaseTTL = defaults.FailureLockBaseTTL
	}
	if cfg.FailureLockMaxTTL <= 0 {
		cfg.FailureLockMaxTTL = defaults.FailureLockMaxTTL
	}
	if cfg.ConcurrentAttemptLockTTL <= 0 {
		cfg.ConcurrentAttemptLockTTL = defaults.ConcurrentAttemptLockTTL
	}
	if cfg.FailureLockMaxTTL < cfg.FailureLockBaseTTL {
		cfg.FailureLockMaxTTL = cfg.FailureLockBaseTTL
	}
	return cfg
}

func normalizeLoginGuardScope(scope string) string {
	scope = strings.ToLower(strings.TrimSpace(scope))
	if scope == "" {
		return "default"
	}
	var b strings.Builder
	for _, r := range scope {
		switch {
		case r >= 'a' && r <= 'z':
			b.WriteRune(r)
		case r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == '-' || r == '_':
			b.WriteRune(r)
		}
	}
	if b.Len() == 0 {
		return "default"
	}
	return b.String()
}

func normalizeLoginIdentifier(value string) string {
	normalized := strings.ToLower(strings.TrimSpace(value))
	if normalized == "" {
		return "unknown"
	}
	return normalized
}

func normalizeLoginIP(value string) string {
	trimmed := strings.TrimSpace(value)
	if trimmed == "" {
		return "unknown"
	}
	if host, _, err := net.SplitHostPort(trimmed); err == nil {
		trimmed = host
	}
	if ip := net.ParseIP(trimmed); ip != nil {
		return ip.String()
	}
	return strings.ToLower(trimmed)
}

func hashLoginGuardKey(value string) string {
	sum := sha256.Sum256([]byte(value))
	return hex.EncodeToString(sum[:])[:32]
}

func randomLoginGuardToken() (string, error) {
	buf := make([]byte, 16)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf), nil
}

func redisMillisToDuration(value any) time.Duration {
	switch v := value.(type) {
	case int64:
		if v <= 0 {
			return time.Second
		}
		return time.Duration(v) * time.Millisecond
	case string:
		if v == "" {
			return time.Second
		}
		var n int64
		if _, err := fmt.Sscan(v, &n); err == nil && n > 0 {
			return time.Duration(n) * time.Millisecond
		}
	}
	return time.Second
}

type memoryLoginGuard struct {
	mu       sync.Mutex
	cfg      LoginGuardConfig
	requests map[string]memoryLoginCounter
	failures map[string]memoryLoginCounter
	locks    map[string]time.Time
	inflight map[string]memoryLoginInflight
}

type memoryLoginCounter struct {
	count     int
	expiresAt time.Time
}

type memoryLoginInflight struct {
	token     string
	expiresAt time.Time
}

type memoryLoginPermit struct {
	keys  loginGuardKeys
	token string
}

func newMemoryLoginGuard(cfg LoginGuardConfig) *memoryLoginGuard {
	return &memoryLoginGuard{
		cfg:      cfg,
		requests: make(map[string]memoryLoginCounter),
		failures: make(map[string]memoryLoginCounter),
		locks:    make(map[string]time.Time),
		inflight: make(map[string]memoryLoginInflight),
	}
}

func (g *memoryLoginGuard) acquire(keys loginGuardKeys, token string) (*memoryLoginPermit, error) {
	now := time.Now()
	g.mu.Lock()
	defer g.mu.Unlock()

	if retryAfter := g.lockRetryAfterLocked(keys.lockIdentifier, now); retryAfter > 0 {
		return nil, &LoginGuardBlockedError{Reason: LoginGuardBlockedLocked, RetryAfter: retryAfter}
	}
	if retryAfter := g.lockRetryAfterLocked(keys.lockPair, now); retryAfter > 0 {
		return nil, &LoginGuardBlockedError{Reason: LoginGuardBlockedLocked, RetryAfter: retryAfter}
	}
	if retryAfter := g.inflightRetryAfterLocked(keys.inflightIdentifier, now); retryAfter > 0 {
		return nil, &LoginGuardBlockedError{Reason: LoginGuardBlockedInProgress, RetryAfter: retryAfter}
	}

	g.inflight[keys.inflightIdentifier] = memoryLoginInflight{
		token:     token,
		expiresAt: now.Add(g.cfg.ConcurrentAttemptLockTTL),
	}

	if retryAfter := g.consumeRequestLocked(keys.requestIP, g.cfg.RequestIPLimit, g.cfg.RequestIPWindow, now); retryAfter > 0 {
		delete(g.inflight, keys.inflightIdentifier)
		return nil, &LoginGuardBlockedError{Reason: LoginGuardBlockedRateLimited, RetryAfter: retryAfter}
	}
	if retryAfter := g.consumeRequestLocked(keys.requestIdentifier, g.cfg.RequestIdentifierLimit, g.cfg.RequestIdentifierWindow, now); retryAfter > 0 {
		delete(g.inflight, keys.inflightIdentifier)
		return nil, &LoginGuardBlockedError{Reason: LoginGuardBlockedRateLimited, RetryAfter: retryAfter}
	}
	if retryAfter := g.consumeRequestLocked(keys.requestPair, g.cfg.RequestPairLimit, g.cfg.RequestPairWindow, now); retryAfter > 0 {
		delete(g.inflight, keys.inflightIdentifier)
		return nil, &LoginGuardBlockedError{Reason: LoginGuardBlockedRateLimited, RetryAfter: retryAfter}
	}

	return &memoryLoginPermit{keys: keys, token: token}, nil
}

func (g *memoryLoginGuard) recordFailure(permit *memoryLoginPermit) {
	if permit == nil {
		return
	}
	now := time.Now()
	g.mu.Lock()
	defer g.mu.Unlock()

	g.bumpFailureLocked(permit.keys.failureIdentifier, permit.keys.lockIdentifier, g.cfg.FailureIdentifierLimit, now)
	g.bumpFailureLocked(permit.keys.failurePair, permit.keys.lockPair, g.cfg.FailurePairLimit, now)
	g.releaseLocked(permit.keys.inflightIdentifier, permit.token)
}

func (g *memoryLoginGuard) recordSuccess(permit *memoryLoginPermit) {
	if permit == nil {
		return
	}
	g.mu.Lock()
	defer g.mu.Unlock()

	delete(g.failures, permit.keys.failureIdentifier)
	delete(g.failures, permit.keys.failurePair)
	delete(g.locks, permit.keys.lockIdentifier)
	delete(g.locks, permit.keys.lockPair)
	g.releaseLocked(permit.keys.inflightIdentifier, permit.token)
}

func (g *memoryLoginGuard) release(permit *memoryLoginPermit) {
	if permit == nil {
		return
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	g.releaseLocked(permit.keys.inflightIdentifier, permit.token)
}

func (g *memoryLoginGuard) lockRetryAfterLocked(key string, now time.Time) time.Duration {
	expiresAt, ok := g.locks[key]
	if !ok {
		return 0
	}
	if !expiresAt.After(now) {
		delete(g.locks, key)
		return 0
	}
	return expiresAt.Sub(now)
}

func (g *memoryLoginGuard) inflightRetryAfterLocked(key string, now time.Time) time.Duration {
	current, ok := g.inflight[key]
	if !ok {
		return 0
	}
	if !current.expiresAt.After(now) {
		delete(g.inflight, key)
		return 0
	}
	return current.expiresAt.Sub(now)
}

func (g *memoryLoginGuard) consumeRequestLocked(key string, limit int, window time.Duration, now time.Time) time.Duration {
	if limit <= 0 {
		return 0
	}
	current := g.requests[key]
	if !current.expiresAt.After(now) {
		current = memoryLoginCounter{expiresAt: now.Add(window)}
	}
	current.count++
	g.requests[key] = current
	if current.count <= limit {
		return 0
	}
	return current.expiresAt.Sub(now)
}

func (g *memoryLoginGuard) bumpFailureLocked(failureKey, lockKey string, threshold int, now time.Time) {
	if threshold <= 0 {
		return
	}
	current := g.failures[failureKey]
	if !current.expiresAt.After(now) {
		current = memoryLoginCounter{expiresAt: now.Add(g.cfg.FailureWindow)}
	}
	current.count++
	g.failures[failureKey] = current
	if current.count < threshold {
		return
	}
	steps := (current.count - threshold) / threshold
	lockTTL := g.cfg.FailureLockBaseTTL
	for i := 0; i < steps; i++ {
		lockTTL *= 2
		if lockTTL >= g.cfg.FailureLockMaxTTL {
			lockTTL = g.cfg.FailureLockMaxTTL
			break
		}
	}
	g.locks[lockKey] = now.Add(lockTTL)
}

func (g *memoryLoginGuard) unlock(keys loginGuardKeys) {
	g.mu.Lock()
	defer g.mu.Unlock()
	delete(g.locks, keys.lockIdentifier)
	delete(g.failures, keys.failureIdentifier)
}

func (g *memoryLoginGuard) releaseLocked(key, token string) {
	current, ok := g.inflight[key]
	if !ok || current.token != token {
		return
	}
	delete(g.inflight, key)
}

var loginGuardAcquireScript = redis.NewScript(`
local token = ARGV[1]
local inflight_ttl = tonumber(ARGV[2])
local ip_limit = tonumber(ARGV[3])
local ip_window = tonumber(ARGV[4])
local id_limit = tonumber(ARGV[5])
local id_window = tonumber(ARGV[6])
local pair_limit = tonumber(ARGV[7])
local pair_window = tonumber(ARGV[8])

local function positive_ttl(key, fallback)
  local ttl = redis.call("PTTL", key)
  if ttl <= 0 then
    return fallback
  end
  return ttl
end

if redis.call("EXISTS", KEYS[1]) == 1 then
  return {0, "login_locked", positive_ttl(KEYS[1], inflight_ttl)}
end
if redis.call("EXISTS", KEYS[2]) == 1 then
  return {0, "login_locked", positive_ttl(KEYS[2], inflight_ttl)}
end

if not redis.call("SET", KEYS[3], token, "PX", inflight_ttl, "NX") then
  return {0, "login_in_progress", positive_ttl(KEYS[3], inflight_ttl)}
end

local function consume(key, limit, window)
  if limit <= 0 then
    return 0
  end
  local count = redis.call("INCR", key)
  if count == 1 then
    redis.call("PEXPIRE", key, window)
  end
  if count > limit then
    return positive_ttl(key, window)
  end
  return 0
end

local ttl = consume(KEYS[4], ip_limit, ip_window)
if ttl <= 0 then
  ttl = consume(KEYS[5], id_limit, id_window)
end
if ttl <= 0 then
  ttl = consume(KEYS[6], pair_limit, pair_window)
end
if ttl > 0 then
  redis.call("DEL", KEYS[3])
  return {0, "login_rate_limited", ttl}
end

return {1, "ok", 0}
`)

var loginGuardFailureScript = redis.NewScript(`
local token = ARGV[1]
local failure_window = tonumber(ARGV[2])
local id_threshold = tonumber(ARGV[3])
local pair_threshold = tonumber(ARGV[4])
local base_ttl = tonumber(ARGV[5])
local max_ttl = tonumber(ARGV[6])

local function bump(fail_key, lock_key, threshold)
  if threshold <= 0 then
    return 0
  end
  local count = redis.call("INCR", fail_key)
  redis.call("PEXPIRE", fail_key, failure_window)
  if count < threshold then
    return 0
  end
  local steps = math.floor((count - threshold) / threshold)
  local ttl = base_ttl * (2 ^ steps)
  if ttl > max_ttl then
    ttl = max_ttl
  end
  redis.call("PSETEX", lock_key, ttl, tostring(count))
  return ttl
end

local id_ttl = bump(KEYS[1], KEYS[3], id_threshold)
local pair_ttl = bump(KEYS[2], KEYS[4], pair_threshold)
if redis.call("GET", KEYS[5]) == token then
  redis.call("DEL", KEYS[5])
end
if id_ttl > pair_ttl then
  return id_ttl
end
return pair_ttl
`)

var loginGuardSuccessScript = redis.NewScript(`
redis.call("DEL", KEYS[1], KEYS[2], KEYS[3], KEYS[4])
if redis.call("GET", KEYS[5]) == ARGV[1] then
  redis.call("DEL", KEYS[5])
end
return 1
`)

var loginGuardReleaseScript = redis.NewScript(`
if redis.call("GET", KEYS[1]) == ARGV[1] then
  redis.call("DEL", KEYS[1])
  return 1
end
return 0
`)