80 lines
1.3 KiB
Go
80 lines
1.3 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type RateLimiter struct {
|
|
visitors map[string]*Visitor
|
|
mu sync.Mutex
|
|
}
|
|
|
|
type Visitor struct {
|
|
lastSeen time.Time
|
|
tokens int
|
|
}
|
|
|
|
func NewRateLimiter() *RateLimiter {
|
|
rl := &RateLimiter{
|
|
visitors: make(map[string]*Visitor),
|
|
}
|
|
// 定期清理过期访问者
|
|
go rl.cleanupVisitors()
|
|
return rl
|
|
}
|
|
|
|
func (rl *RateLimiter) cleanupVisitors() {
|
|
for {
|
|
time.Sleep(time.Minute)
|
|
rl.mu.Lock()
|
|
for ip, v := range rl.visitors {
|
|
if time.Since(v.lastSeen) > time.Minute {
|
|
delete(rl.visitors, ip)
|
|
}
|
|
}
|
|
rl.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (rl *RateLimiter) LimitMiddleware(maxRequests int, per time.Duration) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
ip := c.ClientIP()
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
v, exists := rl.visitors[ip]
|
|
if !exists {
|
|
rl.visitors[ip] = &Visitor{
|
|
lastSeen: time.Now(),
|
|
tokens: maxRequests - 1,
|
|
}
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
// 重置 token
|
|
if time.Since(v.lastSeen) > per {
|
|
v.tokens = maxRequests - 1
|
|
v.lastSeen = time.Now()
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
if v.tokens > 0 {
|
|
v.tokens--
|
|
v.lastSeen = time.Now()
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
// 超过限流
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
|
"error": "请求过于频繁，请稍后再试",
|
|
})
|
|
}
|
|
}
|