OneAuth/api/middleware/ratelimit.go

// Copyright (C) 2024 veypi <i@veypi.com>
// 2025-03-04 16:08:06
// Distributed under terms of the MIT license.

package middleware

import (
	"fmt"
	"time"

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

// RateLimiter 限流中间件
func RateLimiter(requests int, window time.Duration) func(*vigo.X) error {
	return func(x *vigo.X) error {
		// 获取标识（优先用户ID，其次IP）
		identifier := ""
		if uid, ok := x.Get("user_id").(string); ok && uid != "" {
			identifier = uid
		} else {
			identifier = x.GetRemoteIP()
		}

		key := fmt.Sprintf("ratelimit:%s:%s", identifier, x.Request.URL.Path)

		count, err := cache.Incr(key)
		if err != nil {
			// Redis未启用，跳过限流
			return nil
		}

		// 第一次请求设置过期时间
		if count == 1 {
			cache.Expire(key, window)
		}

		if count > int64(requests) {
			return vigo.ErrForbidden.WithString("rate limit exceeded")
		}

		return nil
	}
}