diff --git a/tests/test_security_base.py b/tests/test_security_base.py index e8a9669..faddb4a 100644 --- a/tests/test_security_base.py +++ b/tests/test_security_base.py @@ -171,51 +171,6 @@ async def scenario(): _run(scenario()) -# --------------------------------------------------------------------------- -# MemoryBackend sync methods -# --------------------------------------------------------------------------- - - -def test_memory_backend_incr_sync(): - backend = MemoryBackend() - assert backend.incr_sync("key1", window=60.0) == 1 - assert backend.incr_sync("key1", window=60.0) == 2 - assert backend.incr_sync("key1", window=60.0) == 3 - - -def test_memory_backend_get_sync(): - backend = MemoryBackend() - assert backend.get_sync("missing") == 0 - backend.incr_sync("key1", window=60.0) - backend.incr_sync("key1", window=60.0) - assert backend.get_sync("key1") == 2 - - -def test_memory_backend_reset_sync(): - backend = MemoryBackend() - backend.incr_sync("key1", window=60.0) - backend.incr_sync("key1", window=60.0) - assert backend.get_sync("key1") == 2 - backend.reset_sync("key1") - assert backend.get_sync("key1") == 0 - - -def test_memory_backend_sync_matches_async(): - """Sync and async methods should produce identical results.""" - async def scenario(): - async_backend = MemoryBackend() - sync_backend = MemoryBackend() - - for _i in range(5): - async_result = await async_backend.incr("key1", window=60.0) - sync_result = sync_backend.incr_sync("key1", window=60.0) - assert async_result == sync_result - - assert await async_backend.get("key1") == sync_backend.get_sync("key1") - - _run(scenario()) - - # --------------------------------------------------------------------------- # HookManager # --------------------------------------------------------------------------- diff --git a/tests/test_security_brute_force.py b/tests/test_security_brute_force.py index bca6b16..847f777 100644 --- a/tests/test_security_brute_force.py +++ b/tests/test_security_brute_force.py @@ -34,19 +34,25 @@ def _make_bf(**kwargs): def test_record_failure_increments(): - bf = _make_bf(max_attempts=5, window_seconds=60) - assert bf.record_failure("user:1.2.3.4") == 1 - assert bf.record_failure("user:1.2.3.4") == 2 - assert bf.record_failure("user:1.2.3.4") == 3 + async def scenario(): + bf = _make_bf(max_attempts=5, window_seconds=60) + assert await bf.record_failure("user:1.2.3.4") == 1 + assert await bf.record_failure("user:1.2.3.4") == 2 + assert await bf.record_failure("user:1.2.3.4") == 3 + + _run(scenario()) def test_is_locked_after_threshold(): - bf = _make_bf(max_attempts=3, window_seconds=60, lockout_seconds=120) - bf.record_failure("user:1") - bf.record_failure("user:1") - assert bf.is_locked("user:1") is False # 2 < 3 - bf.record_failure("user:1") # 3 >= 3, locked - assert bf.is_locked("user:1") is True + async def scenario(): + bf = _make_bf(max_attempts=3, window_seconds=60, lockout_seconds=120) + await bf.record_failure("user:1") + await bf.record_failure("user:1") + assert bf.is_locked("user:1") is False # 2 < 3 + await bf.record_failure("user:1") # 3 >= 3, locked + assert bf.is_locked("user:1") is True + + _run(scenario()) def test_is_locked_returns_false_for_unknown_key(): @@ -55,42 +61,54 @@ def test_is_locked_returns_false_for_unknown_key(): def test_mark_success_resets(): - bf = _make_bf(max_attempts=3, window_seconds=60, lockout_seconds=120) - bf.record_failure("user:1") - bf.record_failure("user:1") - bf.record_failure("user:1") - assert bf.is_locked("user:1") is True - bf.mark_success("user:1") - assert bf.is_locked("user:1") is False + async def scenario(): + bf = _make_bf(max_attempts=3, window_seconds=60, lockout_seconds=120) + await bf.record_failure("user:1") + await bf.record_failure("user:1") + await bf.record_failure("user:1") + assert bf.is_locked("user:1") is True + await bf.mark_success("user:1") + assert bf.is_locked("user:1") is False + + _run(scenario()) def test_mark_success_clears_counter(): - bf = _make_bf(max_attempts=3, window_seconds=60, lockout_seconds=120) - bf.record_failure("user:1") - bf.record_failure("user:1") - bf.mark_success("user:1") - # Counter reset — should need 3 new failures - bf.record_failure("user:1") - bf.record_failure("user:1") - assert bf.is_locked("user:1") is False - bf.record_failure("user:1") - assert bf.is_locked("user:1") is True + async def scenario(): + bf = _make_bf(max_attempts=3, window_seconds=60, lockout_seconds=120) + await bf.record_failure("user:1") + await bf.record_failure("user:1") + await bf.mark_success("user:1") + # Counter reset — should need 3 new failures + await bf.record_failure("user:1") + await bf.record_failure("user:1") + assert bf.is_locked("user:1") is False + await bf.record_failure("user:1") + assert bf.is_locked("user:1") is True + + _run(scenario()) def test_get_retry_after(): - bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=300) - assert bf.get_retry_after("user:1") == 0 # not locked - bf.record_failure("user:1") - bf.record_failure("user:1") - assert bf.get_retry_after("user:1") > 0 + async def scenario(): + bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=300) + assert bf.get_retry_after("user:1") == 0 # not locked + await bf.record_failure("user:1") + await bf.record_failure("user:1") + assert bf.get_retry_after("user:1") > 0 + + _run(scenario()) def test_separate_keys_independent(): - bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=120) - bf.record_failure("user:A") - bf.record_failure("user:A") - assert bf.is_locked("user:A") is True - assert bf.is_locked("user:B") is False + async def scenario(): + bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=120) + await bf.record_failure("user:A") + await bf.record_failure("user:A") + assert bf.is_locked("user:A") is True + assert bf.is_locked("user:B") is False + + _run(scenario()) # --------------------------------------------------------------------------- @@ -99,13 +117,52 @@ def test_separate_keys_independent(): def test_custom_backend(): - backend = MemoryBackend() - bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=60, backend=backend) - bf.record_failure("test") - bf.record_failure("test") - assert bf.is_locked("test") is True - # Verify it used our backend - assert backend.get_sync("bf:test") == 2 + async def scenario(): + backend = MemoryBackend() + bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=60, backend=backend) + await bf.record_failure("test") + await bf.record_failure("test") + assert bf.is_locked("test") is True + # Verify it used our backend + assert await backend.get("bf:test") == 2 + + _run(scenario()) + + +class _FakeAsyncOnlyBackend: + """A minimal StorageBackend implementer with only the async Protocol + methods — no incr_sync/reset_sync. Stands in for RedisBackend (or any + other real backend) to prove record_failure/mark_success don't require + a MemoryBackend-only sync shim.""" + + def __init__(self) -> None: + self._counts: dict[str, int] = {} + + async def incr(self, key: str, window: float) -> int: + self._counts[key] = self._counts.get(key, 0) + 1 + return self._counts[key] + + async def get(self, key: str) -> int: + return self._counts.get(key, 0) + + async def reset(self, key: str) -> None: + self._counts.pop(key, None) + + +def test_works_with_any_storage_backend_protocol_implementer(): + """Regression test for issue #12: a backend without incr_sync/reset_sync + (e.g. RedisBackend) used to AttributeError on the first failed attempt.""" + + async def scenario(): + backend = _FakeAsyncOnlyBackend() + bf = _make_bf(max_attempts=2, window_seconds=60, lockout_seconds=60, backend=backend) + assert await bf.record_failure("user:1") == 1 + assert await bf.record_failure("user:1") == 2 + assert bf.is_locked("user:1") is True + await bf.mark_success("user:1") + assert bf.is_locked("user:1") is False + + _run(scenario()) # --------------------------------------------------------------------------- @@ -160,8 +217,8 @@ async def scenario(): # Lock the key directly (this is the same IP TestClient/the request # scope above resolves to via the default IP-based key_func) - middleware.record_failure("testclient") - middleware.record_failure("testclient") + await middleware.record_failure("testclient") + await middleware.record_failure("testclient") resp = await middleware(request) assert resp.status_code == 429 diff --git a/velocix/security/brute_force.py b/velocix/security/brute_force.py index a46fdef..5150b65 100644 --- a/velocix/security/brute_force.py +++ b/velocix/security/brute_force.py @@ -43,13 +43,13 @@ bf._backend = MemoryBackend() bf._lockouts: dict[str, float] = {} - # In your login handler: + # In your (async) login handler: key = f"login:{username}:{ip}" if bf.is_locked(key): return JSONResponse({"error": "Account locked"}, status_code=429) # ... verify credentials ... - bf.record_failure(key) # on bad password - bf.mark_success(key) # on good password + await bf.record_failure(key) # on bad password + await bf.mark_success(key) # on good password """ import time @@ -142,24 +142,19 @@ def _lock(self, key: str) -> None: """Lock a key for ``_lockout_seconds`` from now.""" self._lockouts[key] = time.time() + self._lockout_seconds - def record_failure(self, key: str) -> int: + async def record_failure(self, key: str) -> int: """Record a failed attempt. Returns the current count within the window. If the count exceeds ``max_attempts``, the key is locked out. """ - # incr_sync/reset_sync aren't part of the async StorageBackend Protocol — - # _patch_backend() monkey-patches them onto MemoryBackend only (its ops - # are plain dict access under the hood, so a sync shim is safe there). - # A backend that doesn't get patched (e.g. RedisBackend) would AttributeError - # here; see issue tracking async-backend support for record_failure/mark_success. - count: int = self._backend.incr_sync(f"bf:{key}", self._window_seconds) # type: ignore[attr-defined] + count = await self._backend.incr(f"bf:{key}", self._window_seconds) if count >= self._max_attempts: self._lock(key) return count - def mark_success(self, key: str) -> None: + async def mark_success(self, key: str) -> None: """Reset the failure counter and lockout for a key after successful auth.""" - self._backend.reset_sync(f"bf:{key}") # type: ignore[attr-defined] + await self._backend.reset(f"bf:{key}") self._lockouts.pop(key, None) def get_retry_after(self, key: str) -> int: @@ -191,47 +186,3 @@ async def _on_request(self, request: Request) -> Response: ) return await self.app(request) - - -# --------------------------------------------------------------------------- -# StorageBackend sync wrappers for brute force counting -# The async StorageBackend protocol is designed for middleware, but brute -# force often needs sync access from non-async login handlers. These thin -# wrappers run the async method in a new event loop if needed. -# --------------------------------------------------------------------------- - -def _patch_backend() -> None: - """Add sync methods to MemoryBackend and StorageBackend implementations. - - MemoryBackend operations are actually synchronous dict ops, so we can - call them directly. For RedisBackend, we'd need to handle async properly. - This patch adds ``incr_sync`` and ``reset_sync`` methods. - """ - def _memory_incr_sync(self: MemoryBackend, key: str, window: float) -> int: - now = time.time() - entry = self._store.get(key) - if entry is None or entry[1] <= now: - self._store[key] = (1, now + window) - return 1 - count = entry[0] + 1 - self._store[key] = (count, entry[1]) - return count - - def _memory_reset_sync(self: MemoryBackend, key: str) -> None: - self._store.pop(key, None) - - def _memory_get_sync(self: MemoryBackend, key: str) -> int: - entry = self._store.get(key) - if entry is None or entry[1] <= time.time(): - return 0 - return entry[0] - - if not hasattr(MemoryBackend, "incr_sync"): - MemoryBackend.incr_sync = _memory_incr_sync # type: ignore[attr-defined] - if not hasattr(MemoryBackend, "reset_sync"): - MemoryBackend.reset_sync = _memory_reset_sync # type: ignore[attr-defined] - if not hasattr(MemoryBackend, "get_sync"): - MemoryBackend.get_sync = _memory_get_sync # type: ignore[attr-defined] - - -_patch_backend()