Prism/internal/auth/loginlimiter.go

package auth

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/prism/proxy/internal/store"
)

const (
	defaultMaxAttempts = 5
	defaultLockMinutes = 15
)

type LoginConfig struct {
	MaxAttempts  int
	Window       time.Duration
	LockDuration time.Duration
}

type LoginResult struct {
	Allowed    bool
	Locked     bool
	Remaining  int
	RetryAfter time.Duration
}

type LoginLimiter struct {
	store *store.Store
}

func NewLoginLimiter(st *store.Store) *LoginLimiter {
	return &LoginLimiter{store: st}
}

func ParseLoginConfig(settings map[string]string) LoginConfig {
	maxAttempts := parsePositiveInt(settings["login_max_attempts"], defaultMaxAttempts)
	lockMin := parsePositiveInt(settings["login_lock_minutes"], defaultLockMinutes)
	lock := time.Duration(lockMin) * time.Minute
	return LoginConfig{
		MaxAttempts:  maxAttempts,
		Window:       lock,
		LockDuration: lock,
	}
}

func (l *LoginLimiter) Check(ctx context.Context, ip string, cfg LoginConfig) (LoginResult, error) {
	attempt, err := l.store.GetLoginAttempt(ctx, ip)
	if err != nil {
		return LoginResult{}, err
	}
	if attempt == nil {
		return LoginResult{Allowed: true, Remaining: cfg.MaxAttempts}, nil
	}
	now := time.Now()
	if attempt.LockedUntil != nil && now.Before(*attempt.LockedUntil) {
		return LoginResult{
			Allowed:    false,
			Locked:     true,
			RetryAfter: attempt.LockedUntil.Sub(now),
		}, nil
	}
	remaining := cfg.MaxAttempts - attempt.FailCount
	if remaining < 0 {
		remaining = 0
	}
	if attempt.FailCount > 0 && now.Sub(attempt.WindowStart) > cfg.Window {
		remaining = cfg.MaxAttempts
	}
	return LoginResult{Allowed: true, Remaining: remaining}, nil
}

func (l *LoginLimiter) RecordFailure(ctx context.Context, ip string, cfg LoginConfig) (LoginResult, error) {
	now := time.Now()
	attempt, err := l.store.GetLoginAttempt(ctx, ip)
	if err != nil {
		return LoginResult{}, err
	}
	if attempt != nil && attempt.LockedUntil != nil && now.Before(*attempt.LockedUntil) {
		return LoginResult{
			Allowed:    false,
			Locked:     true,
			RetryAfter: attempt.LockedUntil.Sub(now),
		}, nil
	}

	failCount := 1
	windowStart := now
	if attempt != nil {
		if attempt.LockedUntil != nil || now.Sub(attempt.WindowStart) <= cfg.Window {
			if now.Sub(attempt.WindowStart) <= cfg.Window {
				failCount = attempt.FailCount + 1
				windowStart = attempt.WindowStart
			}
		}
	}

	var lockedUntil *time.Time
	locked := false
	if failCount >= cfg.MaxAttempts {
		t := now.Add(cfg.LockDuration)
		lockedUntil = &t
		locked = true
		_ = l.store.AddAuditLog(ctx, "login_lock", "auth", fmt.Sprintf("ip=%s locked %v", ip, cfg.LockDuration))
	}
	if err := l.store.SetLoginAttempt(ctx, ip, store.LoginAttempt{
		FailCount:   failCount,
		WindowStart: windowStart,
		LockedUntil: lockedUntil,
	}); err != nil {
		return LoginResult{}, err
	}

	remaining := cfg.MaxAttempts - failCount
	if remaining < 0 {
		remaining = 0
	}
	res := LoginResult{Allowed: !locked, Remaining: remaining, Locked: locked}
	if locked && lockedUntil != nil {
		res.RetryAfter = lockedUntil.Sub(now)
	}
	return res, nil
}

func (l *LoginLimiter) RecordSuccess(ctx context.Context, ip string) error {
	return l.store.ClearLoginAttempt(ctx, ip)
}

func parsePositiveInt(raw string, fallback int) int {
	n, err := strconv.Atoi(raw)
	if err != nil || n <= 0 {
		return fallback
	}
	return n
}