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()
|
|||
|
|
}
|
|||
|
|
}
|