70 lines
1.4 KiB
Go
70 lines
1.4 KiB
Go
package middleware
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
|
|
"github.com/geo-platform/tenant-api/internal/shared/response"
|
|
)
|
|
|
|
type rateLimitEntry struct {
|
|
windowStart time.Time
|
|
count int
|
|
}
|
|
|
|
type inMemoryRateLimiter struct {
|
|
mu sync.Mutex
|
|
entries map[string]rateLimitEntry
|
|
max int
|
|
window time.Duration
|
|
lastSweep time.Time
|
|
}
|
|
|
|
func RateLimitByClientIP(max int, window time.Duration, code int, message, detail string) gin.HandlerFunc {
|
|
limiter := &inMemoryRateLimiter{
|
|
entries: make(map[string]rateLimitEntry),
|
|
max: max,
|
|
window: window,
|
|
}
|
|
return func(c *gin.Context) {
|
|
if !limiter.allow(c.ClientIP(), time.Now()) {
|
|
response.Error(c, response.ErrTooManyRequests(code, message, detail))
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func (l *inMemoryRateLimiter) allow(key string, now time.Time) bool {
|
|
if l == nil || l.max <= 0 || l.window <= 0 || key == "" {
|
|
return true
|
|
}
|
|
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
if l.lastSweep.IsZero() || now.Sub(l.lastSweep) >= l.window {
|
|
l.lastSweep = now
|
|
for entryKey, entry := range l.entries {
|
|
if now.Sub(entry.windowStart) >= l.window {
|
|
delete(l.entries, entryKey)
|
|
}
|
|
}
|
|
}
|
|
|
|
entry, ok := l.entries[key]
|
|
if !ok || now.Sub(entry.windowStart) >= l.window {
|
|
l.entries[key] = rateLimitEntry{windowStart: now, count: 1}
|
|
return true
|
|
}
|
|
if entry.count >= l.max {
|
|
return false
|
|
}
|
|
entry.count++
|
|
l.entries[key] = entry
|
|
return true
|
|
}
|