ai-gateway-go/internal/identity/ratelimit.go

package identity

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"strings"
	"time"

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

// LoginLimiter bounds login attempts per client IP with a Redis-backed
// sliding-window counter. It is defense-in-depth layered on top of the
// per-account lockout (repository.RecordFailure):
//
//   - IP limiter: throttles credential-stuffing spread across many accounts
//     from one source (returns HTTP 429).
//   - Account lockout: stops repeated attempts on a single account.
//
// If Redis is unavailable or not configured the limiter fails OPEN (login
// proceeds, account lockout still applies) rather than locking every user out.
type LoginLimiter struct {
	client  *redis.Client
	max     int
	window  time.Duration
	trusted []netip.Prefix
}

func NewLoginLimiter(client *redis.Client, max int, window time.Duration, trustedProxies []netip.Prefix) *LoginLimiter {
	return &LoginLimiter{client: client, max: max, window: window, trusted: trustedProxies}
}

// ClientIP 提取用于登录限流的客户端 IP。仅当直连对端(RemoteAddr)属于可信
// 代理网段时才采信 X-Forwarded-For;否则任何公网客户端都可以伪造该头,把
// 每 IP 滑动窗口的键旋转掉,彻底绕过登录限流。
func (l *LoginLimiter) ClientIP(r *http.Request) string {
	if l != nil && len(l.trusted) > 0 {
		peer, err := netip.ParseAddr(peerHost(r.RemoteAddr))
		if err == nil {
			peer = peer.Unmap()
			trustedPeer := false
			for _, prefix := range l.trusted {
				if prefix.Contains(peer) {
					trustedPeer = true
					break
				}
			}
			if trustedPeer {
				if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
					if first := strings.TrimSpace(strings.Split(fwd, ",")[0]); first != "" {
						return first
					}
				}
			}
		}
	}
	return peerHost(r.RemoteAddr)
}

func peerHost(remoteAddr string) string {
	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		host = remoteAddr
	}
	return host
}

// ClientIP 兼容旧签名:未配置可信代理时退化为"仅信任内网对端"。
// 保留给直接调用方;HTTP 处理路径统一走 LoginLimiter.ClientIP。
func ClientIP(r *http.Request) string {
	return NewLoginLimiter(nil, 0, 0, nil).ClientIP(r)
}

// Allow reports whether a login attempt from ip may proceed.
func (l *LoginLimiter) Allow(ctx context.Context, ip string) bool {
	if l == nil || l.client == nil || l.max <= 0 || l.window <= 0 || ip == "" {
		return true // not configured -> fail open
	}
	now := time.Now().UnixMilli()
	// Member must be unique per attempt so ZADD appends instead of overwriting
	// the score of an identical timestamp.
	res, err := allowLoginScript.Run(ctx, l.client,
		[]string{loginLimitKey(ip)},
		now, l.window.Milliseconds(), l.max, uniqueMember(now),
		int(l.window.Seconds())+60,
	).Int64Slice()
	if err != nil {
		return true // Redis hiccup -> fail open
	}
	// Script returns {1, count} when limited, {0, count+1} when admitted.
	return len(res) == 2 && res[0] == 0
}

func loginLimitKey(ip string) string {
	return "gateway:login-limit:" + ip
}

func uniqueMember(now int64) string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return fmt.Sprintf("%d:%d", now, time.Now().UnixNano())
	}
	return fmt.Sprintf("%d:%s", now, hex.EncodeToString(b[:]))
}

// allowLoginScript atomically trims the window, counts entries, and (when
// under the limit) records the attempt and refreshes the key TTL.
//
//	KEYS[1] = key
//	ARGV[1] = now (ms)   ARGV[2] = window (ms)   ARGV[3] = max
//	ARGV[4] = unique member   ARGV[5] = TTL (s)
var allowLoginScript = redis.NewScript(`
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local max = tonumber(ARGV[3])
redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window)
local count = redis.call('ZCARD', key)
if count >= max then
  return {1, count}
end
redis.call('ZADD', key, now, ARGV[4])
redis.call('EXPIRE', key, ARGV[5])
return {0, count + 1}
`)