93 lines
4.0 KiB
Python
93 lines
4.0 KiB
Python
 
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from datetime import UTC, datetime, timedelta
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from pshelpers import needs_pg
|
||
|
|
|
||
|
|
from packetstation import db
|
||
|
|
from packetstation.auth import ratelimit
|
||
|
|
from packetstation.settings import Settings
|
||
|
|
|
||
|
|
T0 = datetime(2026, 10, 16, 8, 0, tzinfo=UTC)
|
||
|
|
W = timedelta(minutes=5)
|
||
|
|
C = timedelta(minutes=10)
|
||
|
|
|
||
|
|
|
||
|
|
def m(n: float) -> datetime:
|
||
|
|
return T0 + timedelta(minutes=n)
|
||
|
|
|
||
|
|
|
||
|
|
def test_account_under_limit_is_free() -> None:
|
||
|
|
assert ratelimit.account_block_until([m(0), m(1), m(2), m(3)], limit=5, window=W, cooldown=C) is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_account_limit_hit_blocks_for_cooldown_from_the_fifth() -> None:
|
||
|
|
fails = [m(0), m(1), m(2), m(3), m(4)]
|
||
|
|
until = ratelimit.account_block_until(fails, limit=5, window=W, cooldown=C)
|
||
|
|
assert until == m(14)
|
||
|
|
# Still refused after the failures leave the window: the cooldown is longer than the window
|
||
|
|
assert ratelimit.retry_seconds(until, m(9)) == 300
|
||
|
|
assert ratelimit.retry_seconds(until, m(14)) is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_account_failures_spread_wider_than_window_do_not_block() -> None:
|
||
|
|
fails = [m(0), m(2), m(4), m(6), m(8), m(10)]
|
||
|
|
assert ratelimit.account_block_until(fails, limit=5, window=W, cooldown=C) is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_ip_rule_counts_only_the_last_window() -> None:
|
||
|
|
now = m(10)
|
||
|
|
fails = [m(4.5)] + [m(6 + i * 0.1) for i in range(29)]
|
||
|
|
assert ratelimit.ip_block_until(fails, now=now, limit=30, window=W) is None # m(4.5) is out of the window
|
||
|
|
fails.append(m(9))
|
||
|
|
until = ratelimit.ip_block_until(fails, now=now, limit=30, window=W)
|
||
|
|
assert until == m(11) # the oldest of the 30 (m(6)) leaves the window at m(11)
|
||
|
|
assert ratelimit.retry_seconds(until, now) == 60
|
||
|
|
|
||
|
|
|
||
|
|
def test_retry_seconds_rounds_up() -> None:
|
||
|
|
assert ratelimit.retry_seconds(T0 + timedelta(seconds=0.2), T0) == 1
|
||
|
|
assert ratelimit.retry_seconds(None, T0) is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_record_refuses_unknown_kind() -> None:
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
ratelimit.record(None, account="x", ip=None, kind="sneaky", ok=False)
|
||
|
|
|
||
|
|
|
||
|
|
@needs_pg
|
||
|
|
def test_pg_check_with_explicit_times(pg_settings: Settings) -> None:
|
||
|
|
with db.connect_as("app", pg_settings, autocommit=True) as conn:
|
||
|
|
now = ratelimit.db_now(conn)
|
||
|
|
acct = "weld.andy"
|
||
|
|
for i in range(4):
|
||
|
|
ratelimit.record(
|
||
|
|
conn, account=acct, ip="10.0.0.9", kind="password", ok=False, reason="bad_password", at=now - W / 5 * i
|
||
|
|
)
|
||
|
|
# A success and rate-limited refusals never count
|
||
|
|
ratelimit.record(conn, account=acct, ip="10.0.0.9", kind="password", ok=True, at=now)
|
||
|
|
ratelimit.record(conn, account=acct, ip="10.0.0.9", kind="password", ok=False, reason="rate_limited", at=now)
|
||
|
|
assert ratelimit.check(conn, pg_settings, accounts=[acct], ip="10.0.0.9", now=now) is None
|
||
|
|
ratelimit.record(conn, account=acct, ip="10.0.0.9", kind="password", ok=False, reason="bad_password", at=now)
|
||
|
|
wait = ratelimit.check(conn, pg_settings, accounts=[acct], ip="10.0.0.9", now=now)
|
||
|
|
assert wait == 600
|
||
|
|
# Another account from the same address is still allowed
|
||
|
|
assert ratelimit.check(conn, pg_settings, accounts=["clerk.clara"], ip="10.0.0.9", now=now) is None
|
||
|
|
# Aliases count together
|
||
|
|
assert ratelimit.check(conn, pg_settings, accounts=["andy@x", acct], ip=None, now=now) == 600
|
||
|
|
# After the cooldown the account may try again
|
||
|
|
assert ratelimit.check(conn, pg_settings, accounts=[acct], ip=None, now=now + C + timedelta(seconds=1)) is None
|
||
|
|
|
||
|
|
|
||
|
|
@needs_pg
|
||
|
|
def test_pg_ip_limit(pg_settings: Settings) -> None:
|
||
|
|
with db.connect_as("app", pg_settings, autocommit=True) as conn:
|
||
|
|
now = ratelimit.db_now(conn)
|
||
|
|
for i in range(30):
|
||
|
|
ratelimit.record(
|
||
|
|
conn, account=f"user{i}", ip="10.9.9.9", kind="password", ok=False, reason="no_such_user", at=now
|
||
|
|
)
|
||
|
|
assert ratelimit.check(conn, pg_settings, accounts=["fresh"], ip="10.9.9.9", now=now) == 300
|
||
|
|
assert ratelimit.check(conn, pg_settings, accounts=["fresh"], ip="10.9.9.8", now=now) is None
|