MetaLab/internal/middleware/ratelimit_core.go

package middleware

import (
	"fmt"
	"time"
)

// RateLimitResult 限流检查结果
type RateLimitResult struct {
	Blocked    bool
	RetryAfter int
	Message    string
}

const (
	accountWindow   = 1 * time.Minute
	accountMaxFails = 3
	accountBlockDur = 1 * time.Minute
	ipWindow        = 1 * time.Minute
	ipMaxFails      = 10
	ipBlockDur      = 1 * time.Minute
	cleanupInterval = 2 * time.Minute
	maxEntries      = 10000
)

// windowState 单个维度的限流状态
type windowState struct {
	count        int
	windowStart  time.Time
	blockedUntil time.Time
}

// check 核心检查逻辑
func (rl *RateLimiter) check(m map[string]*windowState, key string, window, blockDur time.Duration, maxFails int, lowKey bool) (RateLimitResult, func()) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	state, exists := m[key]
	if !exists {
		state = &windowState{}
		if len(m) < maxEntries {
			m[key] = state
		}
	}

	if !state.blockedUntil.IsZero() && now.Before(state.blockedUntil) {
		retry := int(state.blockedUntil.Sub(now).Seconds()) + 1
		msg := "请求过于频繁，请稍后重试"
		if lowKey {
			msg = fmt.Sprintf("该账号登录尝试过于频繁，请 %d 秒后重试", retry)
		}
		return RateLimitResult{Blocked: true, RetryAfter: retry, Message: msg}, nil
	}

	if now.Sub(state.windowStart) > window {
		state.count = 0
		state.windowStart = now
		state.blockedUntil = time.Time{}
	}

	recordFail := func() {
		rl.mu.Lock()
		defer rl.mu.Unlock()
		s := m[key]
		if s == nil {
			return
		}
		if now.Sub(s.windowStart) > window {
			s.count = 1
			s.windowStart = now
			return
		}
		s.count++
		if s.count >= maxFails {
			s.blockedUntil = now.Add(blockDur)
		}
	}

	return RateLimitResult{Blocked: false}, recordFail
}