mirror of
https://github.com/phishingclub/phishingclub.git
synced 2026-02-12 16:12:44 +00:00
84 lines
2.4 KiB
Go
84 lines
2.4 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// NewIPRateLimiterMiddleware creates a middleware that limits the number of requests per IP
|
|
// limit is the number of requests per second
|
|
// burst is the maximum burst size, the maximum number of requests that can be made in a burst without being limited
|
|
func NewIPRateLimiterMiddleware(limit float64, burst int) gin.HandlerFunc {
|
|
ipLimiter := NewKeyRateLimiter(rate.Limit(limit), burst, 10*time.Minute)
|
|
return func(c *gin.Context) {
|
|
limiter := ipLimiter.GetLimiter(c.ClientIP())
|
|
if !limiter.Allow() {
|
|
c.AbortWithStatus(http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
//const cleanupInterval = 1 * time.Minute
|
|
//const entryExpiration = 10 * time.Minute
|
|
|
|
// KeyRateLimiter is a rate limiter for key such as username, email or IP
|
|
type KeyRateLimiter struct {
|
|
// ips is a map of key to rate limit
|
|
key sync.Map
|
|
// limiter is the rate limit, e.g. 1 request per seconds
|
|
limiter rate.Limit
|
|
// burst is the maximum burst size, the maximum number of requests that can be made in a burst without being limited
|
|
burst int
|
|
// cleanupInterval is the interval at which the expired keys are cleaned up
|
|
cleanupInterval time.Duration
|
|
}
|
|
|
|
// NewKeyRateLimiter creates a new key rate limiter
|
|
// limiter is the rate limit, e.g. 1 request per seconds
|
|
// burst is the maximum burst size, the maximum number of requests that can be made in a burst without being limited
|
|
func NewKeyRateLimiter(
|
|
limiter rate.Limit,
|
|
burst int,
|
|
cleanupInterval time.Duration,
|
|
) *KeyRateLimiter {
|
|
rl := &KeyRateLimiter{
|
|
limiter: limiter,
|
|
burst: burst,
|
|
}
|
|
go rl.cleanup()
|
|
return rl
|
|
}
|
|
|
|
// cleanup cleans up the expired keys, this is to avoid
|
|
// memory leaking through the sync.Map when the key is not used anymore
|
|
func (r *KeyRateLimiter) cleanup() {
|
|
for range time.Tick(r.cleanupInterval) {
|
|
now := time.Now()
|
|
r.key.Range(func(key, value interface{}) bool {
|
|
expirationTime := value.(time.Time)
|
|
if now.After(expirationTime) {
|
|
r.key.Delete(key)
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
}
|
|
|
|
// GetLimiter gets the limiter for an key or creates one if it does not exist
|
|
func (r *KeyRateLimiter) GetLimiter(key string) *rate.Limiter {
|
|
value, exists := r.key.Load(key)
|
|
if exists {
|
|
return value.(*rate.Limiter)
|
|
}
|
|
|
|
limiter := rate.NewLimiter(r.limiter, r.burst)
|
|
r.key.Store(key, limiter)
|
|
return limiter
|
|
}
|