53 lines
1.0 KiB
Go
53 lines
1.0 KiB
Go
package middleware
|
|
|
|
import (
|
|
"FileRelay/internal/config"
|
|
"FileRelay/internal/model"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
var (
|
|
pickupFailures = make(map[string]int)
|
|
failureMutex sync.Mutex
|
|
)
|
|
|
|
func PickupRateLimit() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
ip := c.ClientIP()
|
|
code := c.Param("pickup_code")
|
|
key := ip + ":" + code
|
|
|
|
failureMutex.Lock()
|
|
count, exists := pickupFailures[key]
|
|
failureMutex.Unlock()
|
|
|
|
if exists && count >= config.GlobalConfig.Security.PickupFailLimit {
|
|
c.JSON(http.StatusTooManyRequests, model.ErrorResponse(model.CodeTooManyRequests, "Too many failed attempts. Please try again later."))
|
|
c.Abort()
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func RecordPickupFailure(ip, code string) {
|
|
key := ip + ":" + code
|
|
failureMutex.Lock()
|
|
pickupFailures[key]++
|
|
|
|
// 设置 1 小时后清除记录 (简单实现)
|
|
go func() {
|
|
time.Sleep(1 * time.Hour)
|
|
failureMutex.Lock()
|
|
delete(pickupFailures, key)
|
|
failureMutex.Unlock()
|
|
}()
|
|
|
|
failureMutex.Unlock()
|
|
}
|