9148d2f6da
Some checks failed
CI / Build & Test (push) Failing after 33s
- 37 个端点全部添加 godoc Swagger 注解（RBAC 11 / Org 5 / Audit 1 / Settings 7 / Notification 7） - 重新生成 docs/（swagger.json/yaml/docs.go） - M8 质量横切：IP 限流、安全头、Body 限制、结构化日志增强、优雅关闭 - 新增 Dockerfile 多阶段构建 + .dockerignore - 新增 testutil 测试工具包 - 修复 testutil.go 编译错误
90 lines
1.5 KiB
Go
90 lines
1.5 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"mengstack/internal/kernel/response"
|
|
"mengstack/internal/kernel/errors"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
type visitor struct {
|
|
count int
|
|
lastSeen time.Time
|
|
}
|
|
|
|
type rateLimiter struct {
|
|
mu sync.Mutex
|
|
visitors map[string]*visitor
|
|
limit int
|
|
window time.Duration
|
|
}
|
|
|
|
func newRateLimiter(limit int, window time.Duration) *rateLimiter {
|
|
rl := &rateLimiter{
|
|
visitors: make(map[string]*visitor),
|
|
limit: limit,
|
|
window: window,
|
|
}
|
|
go rl.cleanup()
|
|
return rl
|
|
}
|
|
|
|
func (rl *rateLimiter) cleanup() {
|
|
ticker := time.NewTicker(rl.window)
|
|
for range ticker.C {
|
|
rl.mu.Lock()
|
|
now := time.Now()
|
|
for key, v := range rl.visitors {
|
|
if now.Sub(v.lastSeen) > rl.window {
|
|
delete(rl.visitors, key)
|
|
}
|
|
}
|
|
rl.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (rl *rateLimiter) allow(key string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
v, exists := rl.visitors[key]
|
|
if !exists {
|
|
rl.visitors[key] = &visitor{count: 1, lastSeen: time.Now()}
|
|
return true
|
|
}
|
|
|
|
if time.Since(v.lastSeen) > rl.window {
|
|
v.count = 1
|
|
v.lastSeen = time.Now()
|
|
return true
|
|
}
|
|
|
|
if v.count >= rl.limit {
|
|
return false
|
|
}
|
|
|
|
v.count++
|
|
v.lastSeen = time.Now()
|
|
return true
|
|
}
|
|
|
|
// RateLimit godoc
|
|
// Limits requests per IP address.
|
|
func RateLimit(limit int, window time.Duration) gin.HandlerFunc {
|
|
rl := newRateLimiter(limit, window)
|
|
|
|
return func(c *gin.Context) {
|
|
ip := c.ClientIP()
|
|
if !rl.allow(ip) {
|
|
response.Fail(c, errors.New("RATE_LIMITED", "too many requests", http.StatusTooManyRequests))
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|