98 lines
3.0 KiB
Go
98 lines
3.0 KiB
Go
// Package ratelimit — брутфорс-лимиты «оконный счётчик» по ключу
|
||
// (обычно IP). Два хранилища за общим интерфейсом Allower:
|
||
// - DBLimiter (db.go) — в проде: счётчик в БД переживает рестарт
|
||
// и общий для нескольких инстансов;
|
||
// - Limiter (этот файл) — in-memory, для юнит-тестов без базы.
|
||
package ratelimit
|
||
|
||
import (
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
type bucket struct {
|
||
count int
|
||
start time.Time
|
||
}
|
||
|
||
// Allower — то, что умеет Middleware: засчитать попытку по ключу
|
||
// и сказать, разрешена ли она. *Limiter и *DBLimiter оба его реализуют.
|
||
type Allower interface {
|
||
Allow(key string) bool
|
||
}
|
||
|
||
// Limiter считает обращения по ключу в пределах скользящего окна.
|
||
// Значение ограничений задаётся при создании — там же, где оно
|
||
// наиболее очевидно читателю.
|
||
type Limiter struct {
|
||
mu sync.Mutex
|
||
limit int
|
||
window time.Duration
|
||
buckets map[string]*bucket
|
||
lastCleanup time.Time
|
||
}
|
||
|
||
// New создаёт лимитер на limit запросов за каждые window по одному ключу.
|
||
func New(limit int, window time.Duration) *Limiter {
|
||
return &Limiter{
|
||
limit: limit,
|
||
window: window,
|
||
buckets: make(map[string]*bucket),
|
||
}
|
||
}
|
||
|
||
// Allow засчитывает попытку и говорит, разрешена ли она.
|
||
func (l *Limiter) Allow(key string) bool {
|
||
now := time.Now()
|
||
|
||
l.mu.Lock()
|
||
defer l.mu.Unlock()
|
||
|
||
l.cleanupLocked(now)
|
||
|
||
b := l.buckets[key]
|
||
if b == nil || now.Sub(b.start) >= l.window {
|
||
l.buckets[key] = &bucket{count: 1, start: now}
|
||
return true
|
||
}
|
||
if b.count >= l.limit {
|
||
return false
|
||
}
|
||
b.count++
|
||
return true
|
||
}
|
||
|
||
// cleanupLocked выбрасывает протухшие ключи — карта не должна расти
|
||
// бесконечно под потоком разных IP. Чистка раз в окно, а не на каждый вызов.
|
||
func (l *Limiter) cleanupLocked(now time.Time) {
|
||
if now.Sub(l.lastCleanup) < l.window {
|
||
return
|
||
}
|
||
l.lastCleanup = now
|
||
for k, b := range l.buckets {
|
||
if now.Sub(b.start) >= l.window {
|
||
delete(l.buckets, k)
|
||
}
|
||
}
|
||
}
|
||
|
||
// ByIP — ключ лимита: адрес клиента. Просто и без чтения тела запроса.
|
||
func ByIP(c *gin.Context) string {
|
||
return c.ClientIP()
|
||
}
|
||
|
||
// Middleware ограничивает количество запросов на ключ (обычно — IP)
|
||
// и отвечает 429, не давая запросу дойти до хендлера.
|
||
func Middleware(l Allower, key func(*gin.Context) string, message string) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
if !l.Allow(key(c)) {
|
||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": message})
|
||
return
|
||
}
|
||
c.Next()
|
||
}
|
||
}
|