55 lines
1.4 KiB
Go
55 lines
1.4 KiB
Go
package middleware
|
|
|
|
import (
|
|
httpx "kra/internal/server/httpx"
|
|
systemservice "kra/internal/service/system"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func SecurityRateLimit(settings *systemservice.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
|
|
}
|
|
if settings == nil {
|
|
c.Next()
|
|
return
|
|
}
|
|
config, err := settings.CurrentSecurity(c.Request.Context())
|
|
if err != nil || config == nil {
|
|
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"code": httpx.CodeError, "msg": "安全服务暂不可用"})
|
|
return
|
|
}
|
|
if !config.LimitEnable {
|
|
c.Next()
|
|
return
|
|
}
|
|
window := config.LimitWindow
|
|
if window < 1 {
|
|
window = 60
|
|
}
|
|
route := c.FullPath()
|
|
if route == "" {
|
|
route = path
|
|
}
|
|
key := "KRA_SecLimit:" + c.ClientIP() + ":" + route
|
|
count, cacheErr := settings.IncrementRateLimit(c.Request.Context(), key, time.Duration(window)*time.Second)
|
|
if cacheErr != nil {
|
|
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"code": httpx.CodeError, "msg": "安全服务暂不可用"})
|
|
return
|
|
}
|
|
if int(count) > config.LimitCount {
|
|
c.JSON(http.StatusOK, gin.H{"code": httpx.CodeError, "msg": "请求太过频繁，请稍后再试"})
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|