342 lines
9.5 KiB
Go
342 lines
9.5 KiB
Go
package http
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"hyapp/services/gateway-service/internal/transport/http/httpkit"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/redis/go-redis/v9"
|
||
)
|
||
|
||
const (
|
||
authOperationPassword = "password"
|
||
authOperationThirdParty = "third_party"
|
||
authOperationRefresh = "refresh"
|
||
authOperationLogout = "logout"
|
||
)
|
||
|
||
// AuthRateLimitConfig 描述 gateway public auth 入口的固定窗口频控策略。
|
||
type AuthRateLimitConfig struct {
|
||
// Enabled 为 false 时关闭 gateway 频控。
|
||
Enabled bool
|
||
// KeyPrefix 是 Redis key 前缀，用于区分环境和服务。
|
||
KeyPrefix string
|
||
// Window 是固定计数窗口。
|
||
Window time.Duration
|
||
// IPLimit 限制同一入口 IP 对所有 public auth API 的请求数。
|
||
IPLimit int
|
||
// DeviceIDLimit 限制同一 device_id 对登录、注册和 refresh 的请求数。
|
||
DeviceIDLimit int
|
||
// ProviderIPLimit 限制同一 IP 对同一三方 provider 的注册/登录请求数。
|
||
ProviderIPLimit int
|
||
// DisplayUserIDIPLimit 限制同一 IP 对同一 display_user_id 的密码尝试。
|
||
DisplayUserIDIPLimit int
|
||
// DisplayUserIDLimit 限制同一 display_user_id 被跨 IP 喷射尝试。
|
||
DisplayUserIDLimit int
|
||
// SessionIDIPLimit 限制同一 IP 对同一 session_id 的 logout 请求数。
|
||
SessionIDIPLimit int
|
||
}
|
||
|
||
type authRateInput struct {
|
||
operation string
|
||
deviceID string
|
||
provider string
|
||
displayUserID string
|
||
sessionID string
|
||
}
|
||
|
||
type authRateKey struct {
|
||
key string
|
||
limit int
|
||
}
|
||
|
||
type authRateWindow struct {
|
||
startedAt time.Time
|
||
count int
|
||
}
|
||
|
||
type publicAuthRateLimiter interface {
|
||
allow(ctx context.Context, keys []authRateKey) (bool, error)
|
||
}
|
||
|
||
type memoryAuthRateLimiter struct {
|
||
mu sync.Mutex
|
||
now func() time.Time
|
||
window time.Duration
|
||
windows map[string]authRateWindow
|
||
}
|
||
|
||
type redisAuthRateLimiter struct {
|
||
runner authRateRedisRunner
|
||
window time.Duration
|
||
}
|
||
|
||
type authRateRedisRunner interface {
|
||
evalInt(ctx context.Context, script string, keys []string, args ...any) (int64, error)
|
||
}
|
||
|
||
type goRedisAuthRateRunner struct {
|
||
client *redis.Client
|
||
}
|
||
|
||
const redisAuthRateLimitScript = `
|
||
for i = 1, #KEYS do
|
||
local current = tonumber(redis.call("GET", KEYS[i]) or "0")
|
||
local limit = tonumber(ARGV[i + 1])
|
||
if current >= limit then
|
||
return 0
|
||
end
|
||
end
|
||
|
||
for i = 1, #KEYS do
|
||
local next = redis.call("INCR", KEYS[i])
|
||
if next == 1 then
|
||
redis.call("PEXPIRE", KEYS[i], ARGV[1])
|
||
end
|
||
end
|
||
|
||
return 1
|
||
`
|
||
|
||
func newMemoryAuthRateLimiter(config AuthRateLimitConfig) *memoryAuthRateLimiter {
|
||
config = normalizeAuthRateLimitConfig(config)
|
||
return &memoryAuthRateLimiter{
|
||
now: time.Now,
|
||
window: config.Window,
|
||
windows: make(map[string]authRateWindow),
|
||
}
|
||
}
|
||
|
||
func NewAuthRateLimitRedisClient(ctx context.Context, addr string, password string, db int) (*redis.Client, error) {
|
||
client := redis.NewClient(&redis.Options{
|
||
Addr: strings.TrimSpace(addr),
|
||
Password: password,
|
||
DB: db,
|
||
})
|
||
if err := client.Ping(ctx).Err(); err != nil {
|
||
_ = client.Close()
|
||
return nil, err
|
||
}
|
||
|
||
return client, nil
|
||
}
|
||
|
||
func NewRedisAuthRateLimiter(client *redis.Client, config AuthRateLimitConfig) publicAuthRateLimiter {
|
||
return newRedisAuthRateLimiter(goRedisAuthRateRunner{client: client}, config)
|
||
}
|
||
|
||
func newRedisAuthRateLimiter(runner authRateRedisRunner, config AuthRateLimitConfig) *redisAuthRateLimiter {
|
||
config = normalizeAuthRateLimitConfig(config)
|
||
return &redisAuthRateLimiter{runner: runner, window: config.Window}
|
||
}
|
||
|
||
func defaultAuthRateLimitConfig() AuthRateLimitConfig {
|
||
return AuthRateLimitConfig{
|
||
Enabled: true,
|
||
KeyPrefix: "gateway:auth:rate:",
|
||
Window: time.Minute,
|
||
IPLimit: 120,
|
||
DeviceIDLimit: 30,
|
||
ProviderIPLimit: 20,
|
||
DisplayUserIDIPLimit: 10,
|
||
DisplayUserIDLimit: 50,
|
||
SessionIDIPLimit: 30,
|
||
}
|
||
}
|
||
|
||
func normalizeAuthRateLimitConfig(config AuthRateLimitConfig) AuthRateLimitConfig {
|
||
defaults := defaultAuthRateLimitConfig()
|
||
if strings.TrimSpace(config.KeyPrefix) == "" {
|
||
config.KeyPrefix = defaults.KeyPrefix
|
||
}
|
||
if config.Window <= 0 {
|
||
config.Window = defaults.Window
|
||
}
|
||
if config.IPLimit <= 0 {
|
||
config.IPLimit = defaults.IPLimit
|
||
}
|
||
if config.DeviceIDLimit <= 0 {
|
||
config.DeviceIDLimit = defaults.DeviceIDLimit
|
||
}
|
||
if config.ProviderIPLimit <= 0 {
|
||
config.ProviderIPLimit = defaults.ProviderIPLimit
|
||
}
|
||
if config.DisplayUserIDIPLimit <= 0 {
|
||
config.DisplayUserIDIPLimit = defaults.DisplayUserIDIPLimit
|
||
}
|
||
if config.DisplayUserIDLimit <= 0 {
|
||
config.DisplayUserIDLimit = defaults.DisplayUserIDLimit
|
||
}
|
||
if config.SessionIDIPLimit <= 0 {
|
||
config.SessionIDIPLimit = defaults.SessionIDIPLimit
|
||
}
|
||
|
||
return config
|
||
}
|
||
|
||
func (l *memoryAuthRateLimiter) allow(_ context.Context, keys []authRateKey) (bool, error) {
|
||
now := l.now()
|
||
|
||
l.mu.Lock()
|
||
defer l.mu.Unlock()
|
||
|
||
l.pruneLocked(now)
|
||
for _, key := range keys {
|
||
if key.key == "" || key.limit <= 0 {
|
||
continue
|
||
}
|
||
window := l.currentWindowLocked(now, key.key)
|
||
if window.count >= key.limit {
|
||
return false, nil
|
||
}
|
||
}
|
||
for _, key := range keys {
|
||
if key.key == "" || key.limit <= 0 {
|
||
continue
|
||
}
|
||
window := l.currentWindowLocked(now, key.key)
|
||
window.count++
|
||
l.windows[key.key] = window
|
||
}
|
||
|
||
return true, nil
|
||
}
|
||
|
||
func (l *memoryAuthRateLimiter) currentWindowLocked(now time.Time, key string) authRateWindow {
|
||
window, exists := l.windows[key]
|
||
if !exists || now.Sub(window.startedAt) >= l.window {
|
||
return authRateWindow{startedAt: now}
|
||
}
|
||
|
||
return window
|
||
}
|
||
|
||
func (l *memoryAuthRateLimiter) pruneLocked(now time.Time) {
|
||
if len(l.windows) < 4096 {
|
||
return
|
||
}
|
||
for key, window := range l.windows {
|
||
if now.Sub(window.startedAt) >= 2*l.window {
|
||
delete(l.windows, key)
|
||
}
|
||
}
|
||
}
|
||
|
||
func (l *redisAuthRateLimiter) allow(ctx context.Context, keys []authRateKey) (bool, error) {
|
||
if len(keys) == 0 {
|
||
return true, nil
|
||
}
|
||
|
||
redisKeys := make([]string, 0, len(keys))
|
||
args := make([]any, 0, len(keys)+1)
|
||
args = append(args, l.window.Milliseconds())
|
||
for _, key := range keys {
|
||
if key.key == "" || key.limit <= 0 {
|
||
continue
|
||
}
|
||
redisKeys = append(redisKeys, key.key)
|
||
args = append(args, key.limit)
|
||
}
|
||
if len(redisKeys) == 0 {
|
||
return true, nil
|
||
}
|
||
|
||
result, err := l.runner.evalInt(ctx, redisAuthRateLimitScript, redisKeys, args...)
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
|
||
return result == 1, nil
|
||
}
|
||
|
||
func (r goRedisAuthRateRunner) evalInt(ctx context.Context, script string, keys []string, args ...any) (int64, error) {
|
||
return r.client.Eval(ctx, script, keys, args...).Int64()
|
||
}
|
||
|
||
// SetAuthRateLimit 允许测试或本地构造覆盖 gateway public auth 入口频控策略，使用内存 backend。
|
||
func (h *Handler) SetAuthRateLimit(config AuthRateLimitConfig) {
|
||
config = normalizeAuthRateLimitConfig(config)
|
||
h.authRateLimitConfig = config
|
||
h.authRateLimiter = newMemoryAuthRateLimiter(config)
|
||
}
|
||
|
||
// SetAuthRateLimitBackend 允许 app 层注入 Redis backend，避免多 gateway 节点各自计数。
|
||
func (h *Handler) SetAuthRateLimitBackend(config AuthRateLimitConfig, limiter publicAuthRateLimiter) {
|
||
config = normalizeAuthRateLimitConfig(config)
|
||
h.authRateLimitConfig = config
|
||
if limiter == nil {
|
||
h.authRateLimiter = newMemoryAuthRateLimiter(config)
|
||
return
|
||
}
|
||
h.authRateLimiter = limiter
|
||
}
|
||
|
||
func (h *Handler) allowPublicAuthRequest(writer http.ResponseWriter, request *http.Request, input authRateInput) bool {
|
||
config := normalizeAuthRateLimitConfig(h.authRateLimitConfig)
|
||
if !config.Enabled {
|
||
return true
|
||
}
|
||
if h.authRateLimiter == nil {
|
||
h.SetAuthRateLimit(config)
|
||
}
|
||
allowed, err := h.authRateLimiter.allow(request.Context(), h.authRateKeys(request, input))
|
||
if err != nil {
|
||
httpkit.WriteError(writer, request, http.StatusBadGateway, httpkit.CodeUpstreamError, "upstream service error")
|
||
return false
|
||
}
|
||
if allowed {
|
||
return true
|
||
}
|
||
|
||
httpkit.WriteError(writer, request, http.StatusTooManyRequests, httpkit.CodeRateLimited, "rate limited")
|
||
return false
|
||
}
|
||
|
||
func (h *Handler) authRateKeys(request *http.Request, input authRateInput) []authRateKey {
|
||
config := normalizeAuthRateLimitConfig(h.authRateLimitConfig)
|
||
ip := clientIP(request)
|
||
ipPart := authRateKeyPart(ip)
|
||
keyPrefix := strings.TrimSpace(config.KeyPrefix)
|
||
keys := []authRateKey{{key: keyPrefix + "ip:" + ipPart, limit: config.IPLimit}}
|
||
|
||
deviceID := strings.TrimSpace(input.deviceID)
|
||
if deviceID != "" {
|
||
keys = append(keys, authRateKey{key: keyPrefix + "device:" + authRateKeyPart(deviceID), limit: config.DeviceIDLimit})
|
||
}
|
||
|
||
switch input.operation {
|
||
case authOperationThirdParty:
|
||
provider := strings.ToLower(strings.TrimSpace(input.provider))
|
||
if provider != "" {
|
||
keys = append(keys, authRateKey{key: keyPrefix + "third_party:provider_ip:" + authRateKeyPart(provider) + ":" + ipPart, limit: config.ProviderIPLimit})
|
||
}
|
||
case authOperationPassword:
|
||
displayUserID := strings.TrimSpace(input.displayUserID)
|
||
if displayUserID != "" {
|
||
displayPart := authRateKeyPart(displayUserID)
|
||
keys = append(keys,
|
||
authRateKey{key: keyPrefix + "password:display_ip:" + displayPart + ":" + ipPart, limit: config.DisplayUserIDIPLimit},
|
||
authRateKey{key: keyPrefix + "password:display:" + displayPart, limit: config.DisplayUserIDLimit},
|
||
)
|
||
}
|
||
case authOperationLogout:
|
||
sessionID := strings.TrimSpace(input.sessionID)
|
||
if sessionID != "" {
|
||
keys = append(keys, authRateKey{key: keyPrefix + "logout:session_ip:" + authRateKeyPart(sessionID) + ":" + ipPart, limit: config.SessionIDIPLimit})
|
||
}
|
||
}
|
||
|
||
return keys
|
||
}
|
||
|
||
func authRateKeyPart(value string) string {
|
||
sum := sha256.Sum256([]byte(strings.TrimSpace(value)))
|
||
// Public auth 字段和 X-Forwarded-For 都来自入口流量，Redis key 只保留固定长度摘要。
|
||
return hex.EncodeToString(sum[:16])
|
||
}
|