89 lines
4.4 KiB
Go
89 lines
4.4 KiB
Go
package ratelimit
|
||
|
||
import (
|
||
"context"
|
||
"log/slog"
|
||
"time"
|
||
|
||
"github.com/jackc/pgx/v5"
|
||
"github.com/jackc/pgx/v5/pgconn"
|
||
)
|
||
|
||
// Querier — всё, что лимитеру нужно от pgx: пул (pgxpool.Pool) или
|
||
// транзакция (pgx.Tx). Интерфейс свой, а не из repository: пакет
|
||
// лимитера не должен тянуть за собой весь слой репозиториев.
|
||
type Querier interface {
|
||
Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
|
||
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
|
||
}
|
||
|
||
// DBLimiter — лимитер на Postgres (таблица rate_limits, миграция 0009).
|
||
// В отличие от in-memory Limiter, счётчик переживает рестарт процесса
|
||
// и общий для нескольких инстансов API. Ключ хранится как «имя|IP»:
|
||
// у каждого лимита своё имя, поэтому потолки разных эндпоинтов
|
||
// не складываются в одну кучу.
|
||
type DBLimiter struct {
|
||
q Querier
|
||
name string
|
||
limit int
|
||
window time.Duration
|
||
}
|
||
|
||
// NewDB создаёт лимитер name на limit запросов за window по одному ключу.
|
||
// Имена — «login», «refresh», «register», «reset»: видно в таблице,
|
||
// какой эндпоинт упёрся в потолок.
|
||
func NewDB(q Querier, name string, limit int, window time.Duration) *DBLimiter {
|
||
return &DBLimiter{q: q, name: name, limit: limit, window: window}
|
||
}
|
||
|
||
// upsertSQL засчитывает попытку одним round-trip'ом: окно кончилось —
|
||
// счётчик и окно начинаются заново, нет — инкремент. Возвращается
|
||
// счётчик ПОСЛЕ засчёта; гонки двух параллельных запросов не боятся:
|
||
// ON CONFLICT сериализует обновление одной строки.
|
||
// clock_timestamp(), а не now(): now() в Postgres — время начала
|
||
// транзакции, внутри длинной транзакции окно не сдвинулось бы.
|
||
const upsertSQL = `
|
||
INSERT INTO rate_limits (key, window_start, count)
|
||
VALUES ($1, clock_timestamp(), 1)
|
||
ON CONFLICT (key) DO UPDATE SET
|
||
count = CASE WHEN rate_limits.window_start + make_interval(secs => $2) > clock_timestamp()
|
||
THEN rate_limits.count + 1 ELSE 1 END,
|
||
window_start = CASE WHEN rate_limits.window_start + make_interval(secs => $2) > clock_timestamp()
|
||
THEN rate_limits.window_start ELSE clock_timestamp() END
|
||
RETURNING count`
|
||
|
||
// Allow засчитывает попытку и говорит, разрешена ли она.
|
||
// Ошибка БД не блокирует запрос (fail-open): лимитер — не то место,
|
||
// из-за которого падает вход; доступ к базе и так нужен самому логину,
|
||
// а ошибку видно в логе.
|
||
func (l *DBLimiter) Allow(key string) bool {
|
||
var count int
|
||
err := l.q.QueryRow(context.Background(), upsertSQL, l.name+"|"+key, float64(l.window.Seconds())).Scan(&count)
|
||
if err != nil {
|
||
slog.Warn("rate limit: db check failed, request allowed", "limiter", l.name, "err", err)
|
||
return true
|
||
}
|
||
return count <= l.limit
|
||
}
|
||
|
||
// RunCleanup периодически выбрасывает протухшие строки: самое длинное
|
||
// окно — «3 письма в сутки» на forgotten-пароли (24 часа), поэтому
|
||
// строки старше 25 часов (час запаса) не понадобятся никогда. Короткие
|
||
// минутные окна живут дольше нужного — для шкалы rate_limits это
|
||
// копейки (строка на IP/ящик), а чистка одна на все лимиты.
|
||
// Останавливается по отмене ctx (graceful shutdown в cmd/api).
|
||
func RunCleanup(ctx context.Context, q Querier, every time.Duration) {
|
||
t := time.NewTicker(every)
|
||
defer t.Stop()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case <-t.C:
|
||
if _, err := q.Exec(ctx, `DELETE FROM rate_limits WHERE window_start < now() - interval '25 hours'`); err != nil {
|
||
slog.Warn("rate limit: cleanup failed", "err", err)
|
||
}
|
||
}
|
||
}
|
||
}
|