74 lines
2.1 KiB
Go
74 lines
2.1 KiB
Go
|
package am_ratelimit
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"github.com/gin-gonic/gin"
|
|||
|
|
"net/http"
|
|||
|
|
"sync"
|
|||
|
|
"time"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// 储存用户信息，包括最后一次访问时间和访问计数
|
|||
|
|
type RequestInfo struct {
|
|||
|
|
LastAccessTime time.Time
|
|||
|
|
RequestCount int
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 互斥锁，用于同步对共享资源（如requestInfo映射）的访问，确保在多线程环境下的数据一致性
|
|||
|
|
var mutex = &sync.Mutex{}
|
|||
|
|
|
|||
|
|
// 速率限制配置项结构体
|
|||
|
|
type RateLimitConfig struct {
|
|||
|
|
maxRequests int // 最大访问次数
|
|||
|
|
timeWindow time.Duration // 窗口时间
|
|||
|
|
requestInfo map[string]*RequestInfo // 用户信息
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/*
|
|||
|
|
* 初始化速率限制配置
|
|||
|
|
* 参数：
|
|||
|
|
* 1. maxRequests int：窗口时间内最大访问数量限制
|
|||
|
|
* 2. timeWindow int：限制的窗口时间，单位为秒
|
|||
|
|
*/
|
|||
|
|
func NewRateLimitConfig(maxRequests int, timeWindow int) *RateLimitConfig {
|
|||
|
|
return &RateLimitConfig{
|
|||
|
|
maxRequests: maxRequests,
|
|||
|
|
timeWindow: time.Duration(timeWindow) * time.Second,
|
|||
|
|
requestInfo: make(map[string]*RequestInfo),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/*
|
|||
|
|
* 速率限制中间件
|
|||
|
|
*/
|
|||
|
|
func (r1 *RateLimitConfig) RateLimitMiddleware(c *gin.Context) {
|
|||
|
|
ip := c.ClientIP() // 获取用户IP
|
|||
|
|
mutex.Lock() // 在访问共享资源前加锁，确保只有同一线程才能访问
|
|||
|
|
defer mutex.Unlock() // 在数据返回前解锁
|
|||
|
|
|
|||
|
|
info, exists := r1.requestInfo[ip] // 检查requestInfo映射中是否已经存在该IP信息
|
|||
|
|
|
|||
|
|
// 如果该IP的请求信息不存在，则初始化一个新的RequestInfo实例，并将其添加到requestInfo映射中
|
|||
|
|
if !exists {
|
|||
|
|
r1.requestInfo[ip] = &RequestInfo{LastAccessTime: time.Now(), RequestCount: 1}
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 如果最后一次请求的时间已经超出了时间窗口，重置请求计数并更新最后一次访问时间
|
|||
|
|
if time.Since(info.LastAccessTime) > r1.timeWindow {
|
|||
|
|
info.RequestCount = 1
|
|||
|
|
info.LastAccessTime = time.Now()
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 增加请求数并检查速率限制
|
|||
|
|
info.RequestCount++
|
|||
|
|
if info.RequestCount > r1.maxRequests {
|
|||
|
|
c.JSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests"})
|
|||
|
|
c.Abort()
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 更新最后一次访问时间
|
|||
|
|
info.LastAccessTime = time.Now()
|
|||
|
|
}
|