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