CamTalk/backend/internal/ratelimit/middleware.go

package ratelimit

import (
	"fmt"
	"net/http"

	"github.com/gin-gonic/gin"

	"github.com/hhs/camtalk/internal/trace"
)

// Middleware 返回 Gin 中间件，按 key 维度限流。
// keyFunc 从请求中提取限流 key（如 IP、用户 ID）。
func Middleware(limiter Limiter, keyFunc func(*gin.Context) string) gin.HandlerFunc {
	return func(c *gin.Context) {
		if limiter == nil {
			c.Next()
			return
		}

		key := keyFunc(c)
		if key == "" {
			// key 为空时跳过限流
			c.Next()
			return
		}

		allowed, retryAfter := limiter.Allow(c.Request.Context(), key)

		if !allowed {
			log := trace.FromContext(c.Request.Context())
			log.Warnw("rate limited",
				"client_ip", c.ClientIP(),
				"path", c.Request.URL.Path,
				"limit_key", key,
				"retry_after_sec", int(retryAfter.Seconds()+0.5))

			// 设置 Retry-After header（秒）
			c.Header("Retry-After", fmt.Sprintf("%d", int(retryAfter.Seconds()+0.5)))

			c.JSON(http.StatusTooManyRequests, gin.H{
				"code":    "RATE_LIMITED",
				"message": fmt.Sprintf("too many requests, retry after %s", retryAfter.Round(1)),
			})
			c.Abort()
			return
		}

		c.Next()
	}
}