OneAuth/internal/api/middleware/ratelimit.go

package middleware

import (
	"fmt"
	"net/http"
	"time"

	"github.com/veypi/vbase/internal/cache"
	"github.com/veypi/vigo"
)

// RateLimiter 限流中间件
func RateLimiter(maxRequests int, window time.Duration) func(*vigo.X) (any, error) {
	return func(x *vigo.X) (any, error) {
		if !cache.IsEnabled() {
			x.Next()
			return nil, nil
		}

		// 使用IP+路径作为标识
		identifier := x.GetRemoteIP()
		path := x.Request.URL.Path

		count, err := cache.IncrRateLimit(identifier, path, window)
		if err != nil {
			x.Next() // 缓存失败时放行
			return nil, nil
		}

		if count > int64(maxRequests) {
			x.ResponseWriter().Header().Set("Retry-After", fmt.Sprintf("%d", int(window.Seconds())))
			return nil, vigo.NewError("rate limit exceeded").WithCode(http.StatusTooManyRequests)
		}

		x.Next()
		return nil, nil
	}
}

// RateLimiterByUser 基于用户的限流
func RateLimiterByUser(maxRequests int, window time.Duration) func(*vigo.X) (any, error) {
	return func(x *vigo.X) (any, error) {
		if !cache.IsEnabled() {
			x.Next()
			return nil, nil
		}

		userID := CurrentUser(x)
		if userID == "" {
			// 未登录用户使用IP限流
			_, err := RateLimiter(maxRequests, window)(x)
			return nil, err
		}

		path := x.Request.URL.Path
		count, err := cache.IncrRateLimit("user:"+userID, path, window)
		if err != nil {
			x.Next()
			return nil, nil
		}

		if count > int64(maxRequests) {
			x.ResponseWriter().Header().Set("Retry-After", fmt.Sprintf("%d", int(window.Seconds())))
			return nil, vigo.NewError("rate limit exceeded").WithCode(http.StatusTooManyRequests)
		}

		x.Next()
		return nil, nil
	}
}

// LoginRateLimit 登录限流（更严格）
func LoginRateLimit() func(*vigo.X) (any, error) {
	return func(x *vigo.X) (any, error) {
		if !cache.IsEnabled() {
			x.Next()
			return nil, nil
		}

		identifier := x.GetRemoteIP()
		key := "login_attempt:" + identifier

		count, _ := cache.Incr(key)
		if count == 1 {
			cache.Expire(key, 15*time.Minute)
		}

		// 5分钟内超过5次尝试，需要验证码
		if count >= 5 {
			x.Set("require_captcha", true)
		}

		// 超过10次直接拒绝
		if count >= 10 {
			return nil, vigo.NewError("too many login attempts, please try again later").WithCode(http.StatusTooManyRequests)
		}

		x.Next()
		return nil, nil
	}
}

// ResetLoginAttempts 重置登录尝试次数
func ResetLoginAttempts(x *vigo.X) {
	if !cache.IsEnabled() {
		return
	}
	identifier := x.GetRemoteIP()
	key := "login_attempt:" + identifier
	cache.Delete(key)
}