67 lines
1.6 KiB
Go
67 lines
1.6 KiB
Go
package middleware
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
|
|
"git.echol.cn/loser/ai_proxy/server/global"
|
|
"git.echol.cn/loser/ai_proxy/server/model/common/response"
|
|
"github.com/gin-gonic/gin"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
type LimitConfig struct {
|
|
GenerationDuration time.Duration
|
|
Limit int
|
|
}
|
|
|
|
func (l LimitConfig) LimitKey(c *gin.Context) string {
|
|
return "GVA_Limit" + c.ClientIP()
|
|
}
|
|
|
|
func (l LimitConfig) GetLimit(c *gin.Context) int {
|
|
return l.Limit
|
|
}
|
|
|
|
func (l LimitConfig) Reached(c *gin.Context) response.Response {
|
|
return response.Response{Code: response.ERROR, Data: nil, Msg: "操作过于频繁，请稍后再试"}
|
|
}
|
|
|
|
// IPLimit IP 限流中间件
|
|
func IPLimit() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if global.GVA_CONFIG.System.UseRedis {
|
|
key := "GVA_Limit" + c.ClientIP()
|
|
limit := global.GVA_CONFIG.System.IpLimitCount
|
|
limitTime := global.GVA_CONFIG.System.IpLimitTime
|
|
|
|
ctx := context.Background()
|
|
count, err := global.GVA_REDIS.Get(ctx, key).Int()
|
|
if err != nil && !errors.Is(err, context.DeadlineExceeded) {
|
|
global.GVA_LOG.Error("get redis key error", zap.Error(err))
|
|
}
|
|
|
|
if count >= limit {
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"code": response.ERROR,
|
|
"msg": fmt.Sprintf("操作过于频繁，请在 %d 秒后再试", limitTime),
|
|
})
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
pipe := global.GVA_REDIS.Pipeline()
|
|
pipe.Incr(ctx, key)
|
|
pipe.Expire(ctx, key, time.Second*time.Duration(limitTime))
|
|
_, err = pipe.Exec(ctx)
|
|
if err != nil {
|
|
global.GVA_LOG.Error("redis pipeline error", zap.Error(err))
|
|
}
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|