Edu/services/api-gateway/internal/middleware/ratelimit.go

package middleware

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
)

// bucket 是单个客户端的令牌桶
type bucket struct {
	mu       sync.Mutex
	tokens   float64
	lastTime time.Time
}

// rateLimiter 持有所有客户端 IP 的令牌桶
type rateLimiter struct {
	rps     float64
	burst   int
	buckets sync.Map // map[string]*bucket
}

// RateLimit 返回基于令牌桶的限流中间件（内存版，不依赖 Redis）。
// 每个客户端 IP 独立桶，按 rps（每秒令牌数）补充，最大容量为 burst。
// 超限返回 429 + JSON {"error":"rate_limited","retry_after":60}。
// 内部启动 cleanup goroutine，每 5 分钟清理 10 分钟未访问的桶。
func RateLimit(rps float64, burst int) gin.HandlerFunc {
	rl := &rateLimiter{rps: rps, burst: burst}
	// 启动清理 goroutine
	go rl.cleanup(5*time.Minute, 10*time.Minute)

	return func(c *gin.Context) {
		ip := c.ClientIP()
		b := rl.getBucket(ip)

		b.mu.Lock()
		now := time.Now()
		elapsed := now.Sub(b.lastTime).Seconds()
		b.lastTime = now
		// 按经过时间补充令牌
		b.tokens += elapsed * rl.rps
		if b.tokens > float64(rl.burst) {
			b.tokens = float64(rl.burst)
		}
		// 令牌不足：拒绝
		if b.tokens < 1 {
			b.mu.Unlock()
			c.Header("Retry-After", "60")
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"error":       "rate_limited",
				"retry_after": 60,
			})
			return
		}
		b.tokens--
		b.mu.Unlock()

		c.Next()
	}
}

// getBucket 获取指定 IP 的令牌桶，不存在则新建（初始满桶）
func (rl *rateLimiter) getBucket(ip string) *bucket {
	if v, ok := rl.buckets.Load(ip); ok {
		return v.(*bucket)
	}
	b := &bucket{
		tokens:   float64(rl.burst),
		lastTime: time.Now(),
	}
	actual, _ := rl.buckets.LoadOrStore(ip, b)
	return actual.(*bucket)
}

// cleanup 周期性清理长时间未访问的桶，避免内存无限增长
func (rl *rateLimiter) cleanup(interval, maxIdle time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for range ticker.C {
		now := time.Now()
		rl.buckets.Range(func(key, value any) bool {
			b := value.(*bucket)
			b.mu.Lock()
			idle := now.Sub(b.lastTime)
			b.mu.Unlock()
			if idle > maxIdle {
				rl.buckets.Delete(key)
			}
			return true
		})
	}
}