55 lines
831 B
Go
55 lines
831 B
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
type bucket struct {
|
||
|
|
tokens float64
|
||
|
|
lastCheck time.Time
|
||
|
|
mu sync.Mutex
|
||
|
|
}
|
||
|
|
|
||
|
|
var buckets sync.Map
|
||
|
|
|
||
|
|
func RateLimit(rps int) gin.HandlerFunc {
|
||
|
|
rate := float64(rps)
|
||
|
|
|
||
|
|
return func(c *gin.Context) {
|
||
|
|
ip := c.ClientIP()
|
||
|
|
|
||
|
|
val, _ := buckets.LoadOrStore(ip, &bucket{
|
||
|
|
tokens: rate,
|
||
|
|
lastCheck: time.Now(),
|
||
|
|
})
|
||
|
|
b := val.(*bucket)
|
||
|
|
|
||
|
|
b.mu.Lock()
|
||
|
|
now := time.Now()
|
||
|
|
elapsed := now.Sub(b.lastCheck).Seconds()
|
||
|
|
b.lastCheck = now
|
||
|
|
b.tokens += elapsed * rate
|
||
|
|
if b.tokens > rate {
|
||
|
|
b.tokens = rate
|
||
|
|
}
|
||
|
|
|
||
|
|
if b.tokens < 1 {
|
||
|
|
b.mu.Unlock()
|
||
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
||
|
|
"code": 42900,
|
||
|
|
"message": "rate limit exceeded",
|
||
|
|
})
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
b.tokens--
|
||
|
|
b.mu.Unlock()
|
||
|
|
|
||
|
|
c.Next()
|
||
|
|
}
|
||
|
|
}
|