Files
ai-gateway-go/internal/identity/ratelimit.go
T
superidou 5759c1862e AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。
含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 11:45:54 +08:00

105 lines
3.2 KiB
Go

package identity
import (
"context"
"crypto/rand"
"encoding/hex"
"fmt"
"net"
"net/http"
"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
}
func NewLoginLimiter(client *redis.Client, max int, window time.Duration) *LoginLimiter {
return &LoginLimiter{client: client, max: max, window: window}
}
// 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}
`)
// ClientIP extracts the caller's IP for login rate limiting. X-Forwarded-For
// is trusted here because nginx is the only ingress and overwrites the header
// on every proxy hop; the first value is the client address. Falls back to
// RemoteAddr for direct connections.
func ClientIP(r *http.Request) string {
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
if first := strings.TrimSpace(strings.Split(fwd, ",")[0]); first != "" {
return first
}
}
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
host = r.RemoteAddr
}
return host
}