kra-oa/app/system/internal/server/middleware/rate_limit.go

46 lines
1.2 KiB
Go

package middleware
import (
"kra/app/system/internal/service"
"net/http"
"strings"
"time"
"github.com/gin-gonic/gin"
)
func SecurityRateLimit(settings *service.SecurityService) gin.HandlerFunc {
return func(c *gin.Context) {
path := strings.TrimSuffix(c.Request.URL.Path, "/")
if !strings.HasSuffix(path, "/base/login") && !strings.HasSuffix(path, "/base/captcha") {
c.Next()
return
}
config, err := settings.CurrentSecurity(c.Request.Context())
if err != nil || config == nil {
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"code": CodeError, "msg": "安全服务暂不可用"})
return
}
if !config.LimitEnable {
c.Next()
return
}
window := config.LimitWindow
if window < 1 {
window = 60
}
key := "KRA_SecLimit" + c.ClientIP() + c.FullPath()
count, cacheErr := settings.IncrementRateLimit(c.Request.Context(), key, time.Duration(window)*time.Second)
if cacheErr != nil {
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"code": CodeError, "msg": "安全服务暂不可用"})
return
}
if int(count) > config.LimitCount {
c.JSON(200, gin.H{"code": CodeError, "msg": "请求太过频繁,请稍后再试"})
c.Abort()
return
}
c.Next()
}
}