Use caplinks' AttemptGuard and JsonRecords for the review gate

heed/review_access.py

@@ -30,19 +30,18 @@
 
 from __future__ import annotations
 
+import functools
 import hashlib
 import json
 import os
-import tempfile
 import time
-from collections.abc import Iterator, MutableMapping
 from pathlib import Path
 
 from fastapi import HTTPException, Request, Response
 from starlette.concurrency import run_in_threadpool
 
 from heed.security import RateLimiter, client_ip, make_rate_limiter, rate_limit_key
-from heed.threads import read_capped, valid_id
+from heed.threads import read_capped
 
 ACCESS_DIR = "access"
 COOKIE_PREFIX = "heed_unlock_"
@@ -62,57 +61,16 @@
 )
 
 
-class _JsonRecords(MutableMapping):
-    """``{id: dict}`` kept as ``<dir>/<id>.json`` files (what caplinks stores)."""
-
-    def __init__(self, folder: Path):
-        self.folder = folder
-
-    def _path(self, key: str) -> Path:
-        if not valid_id(key):
-            raise KeyError(key)
-        return self.folder / f"{key}.json"
-
-    def __getitem__(self, key):
-        try:
-            return json.loads(self._path(key).read_text())
-        except FileNotFoundError:
-            raise KeyError(key) from None
-
-    def __setitem__(self, key, value):
-        path = self._path(key)
-        path.parent.mkdir(parents=True, exist_ok=True)
-        fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".tmp")  # 0600: a hash
-        with os.fdopen(fd, "w") as f:
-            json.dump(value, f)
-        os.replace(tmp, path)
-
-    def __delitem__(self, key):
-        try:
-            self._path(key).unlink()
-        except FileNotFoundError:
-            raise KeyError(key) from None
-
-    def __iter__(self) -> Iterator[str]:
-        if self.folder.is_dir():
-            yield from (p.stem for p in self.folder.glob("*.json"))
-
-    def __len__(self):
-        return sum(1 for _ in self)
-
-    def __contains__(self, key):
-        try:
-            return self._path(key).is_file()
-        except KeyError:
-            return False
+def _records_dir(rootdir) -> Path:
+    return Path(rootdir).expanduser() / ACCESS_DIR
 
 
 def _links(rootdir):
     try:
-        from caplinks import CapabilityLinks
+        from caplinks import CapabilityLinks, FileTexts, JsonRecords
     except ImportError as e:
         raise ImportError(_MISSING_EXTRA) from e
-    return CapabilityLinks(_JsonRecords(Path(rootdir).expanduser() / ACCESS_DIR))
+    return CapabilityLinks(JsonRecords(FileTexts(_records_dir(rootdir))))
 
 
 def set_access(rootdir, page_id: str, *, password: str | None = None, expires=None):
@@ -146,7 +104,10 @@ def revoke_access(rootdir, page_id: str):
     """
     links = _links(rootdir)
     if links.record(page_id) is None:
-        if page_id in links.store:  # a damaged record: replace it, closed
+        path = _records_dir(rootdir) / f"{page_id}.json"
+        if path.is_symlink():  # caplinks refuses to follow it: replace it, closed
+            path.unlink()
+        elif page_id in links.store:  # a damaged record
             del links.store[page_id]
         links.create({}, expires="never", token=page_id)
     return links.revoke(page_id)
@@ -182,7 +143,7 @@ def __init__(
         lock_s: float = DFLT_ATTEMPT_WINDOW_S,
         clock=time.monotonic,
     ):
-        self.records = _JsonRecords(Path(rootdir).expanduser() / ACCESS_DIR)
+        self.rootdir = rootdir
         self.trusted_proxies = trusted_proxies
         self.rate_limiter = rate_limiter or make_rate_limiter(
             max_per_window=DFLT_MAX_ATTEMPTS,
@@ -197,16 +158,27 @@ def __init__(
         self.unlock_max_age_s = unlock_max_age_s
         self.lock_s = lock_s
         self._clock = clock
-        self._locked: dict[str, float] = {}  # page id -> when its lock ends
+
+    @functools.cached_property
+    def _guard(self):
+        """The lock-out (caplinks' ``AttemptGuard``), built on the first unlock."""
+        from caplinks import AttemptGuard
+
+        return AttemptGuard(
+            client_limiter=self.rate_limiter,
+            link_limiter=self.page_failure_limiter,
+            lock_s=self.lock_s,
+            clock=self._clock,
+        )
 
     def _open_record(self, page_id: str):
         """``(links, record)`` of a gated page that is neither revoked nor expired;
         ``None`` for an ungated one. Raises 404 (revoked), 410 (expired), 503 (no
         caplinks to judge a gated page: closed, never open)."""
-        if page_id not in self.records:
-            return None
+        if not os.path.lexists(_records_dir(self.rootdir) / f"{page_id}.json"):
+            return None  # no record: open (a bad or symlinked one gates: 404)
         try:
-            links = _links(self.records.folder.parent)
+            links = _links(self.rootdir)
         except ImportError:
             raise HTTPException(
                 503, "This page is protected and cannot be opened here"
@@ -245,16 +217,13 @@ async def unlock(self, request: Request, page_id: str) -> Response:
                 raise TypeError
         except (ValueError, KeyError, TypeError, RecursionError):
             raise HTTPException(400, "Expected a JSON object with a password") from None
-        if self._locked.get(page_id, 0) > self._clock() or not self.rate_limiter.allow(
-            f"{page_id}:{self._client(request)}"
-        ):
+        if not self._guard.allow(page_id, self._client(request)):
             raise HTTPException(
                 429, "Too many attempts. Wait a few minutes and try again."
             )
         # scrypt is slow on purpose: keep it off the event loop
         if not await run_in_threadpool(links.check_password, rec, password):
-            if not self.page_failure_limiter.allow(page_id):
-                self._locked[page_id] = self._clock() + self.lock_s
+            self._guard.failed(page_id)
             raise HTTPException(401, "That password is not right.")
         proto = request.headers.get("x-forwarded-proto", request.url.scheme)
         response = Response(status_code=204, headers={"Cache-Control": "no-store"})

pyproject.toml

@@ -47,7 +47,7 @@ github = ["ghapi>=1.0"]
 # Shared rate limiting across workers/hosts (rate_limit_backend="redis").
 redis = ["redis>=4.2"]
 # Passwords, expiry and revoke on review pages (heed.review_access).
-gate = ["caplinks>=0.0.3"]
+gate = ["caplinks>=0.0.4"]
 # Kept as an alias for existing installs: enlace is now a core dependency (for
 # enlace.ratelimit), so this extra adds nothing.
 enlace = ["enlace>=0.1.62"]
@@ -61,7 +61,7 @@ mcp = ["py2mcp>=0.1.0"]
 # only falls back to `httpx` with a deprecation warning, while starlette <1
 # knows `httpx` alone. Both are declared so the suite collects on either
 # resolution; they are distinct top-level modules and coexist fine.
-dev = ["pytest", "httpx", "httpx2", "pytest-asyncio", "Pillow>=10.0", "caplinks>=0.0.3"]
+dev = ["pytest", "httpx", "httpx2", "pytest-asyncio", "Pillow>=10.0", "caplinks>=0.0.4"]
 
 [tool.pytest.ini_options]
 # wads CI runs `pytest --doctest-modules` with NO path argument, so collection is

tests/test_review.py

@@ -379,14 +379,14 @@ def test_expiry_answers_410_and_revoke_404(root, clip):
     assert client.get(f"/threads/{page.id}").status_code == 410
     publish_review([], rootdir=root, to=page.id, expires="3d")  # extend: open again
     assert client.get(f"/pages/{page.id}").status_code == 200
-    assert main(["--dir", str(root), "revoke", page.id]) == 0
+    assert main(["--dir", str(root), "revoke", "--", page.id]) == 0
     assert client.get(f"/pages/{page.id}").status_code == 404
     assert client.get(f"/pages/{page.id}/media/{page.media[0].name}").status_code == 404
 
 
 def test_revoking_an_ungated_page_closes_it(root, clip):
     page, client = _gated(root, clip)
-    assert main(["--dir", str(root), "revoke", page.id]) == 0
+    assert main(["--dir", str(root), "revoke", "--", page.id]) == 0
     assert client.get(f"/pages/{page.id}").status_code == 404
 
 
@@ -444,7 +444,7 @@ def test_revoking_a_page_with_a_damaged_record_still_closes_it(root, clip):
     page, client = _gated(root, clip, password=PW)
     (root / "access" / f"{page.id}.json").write_text("{not json")
     assert client.get(f"/pages/{page.id}").status_code == 404  # closed, not open
-    assert main(["--dir", str(root), "revoke", page.id]) == 0
+    assert main(["--dir", str(root), "revoke", "--", page.id]) == 0
     assert client.get(f"/pages/{page.id}").status_code == 404
 
 
@@ -474,3 +474,17 @@ def test_publish_cli_reads_the_password_from_a_file(root, clip, tmp_path, capsys
         client.post(f"/pages/{page_id}/unlock", json={"password": PW}).status_code
         == 204
     )
+
+
+def test_a_dangling_or_symlinked_record_closes_the_page_and_revoke_replaces_it(
+    root, clip, tmp_path
+):
+    page, client = _gated(root, clip, password=PW)
+    record = root / "access" / f"{page.id}.json"
+    record.unlink()
+    record.symlink_to(tmp_path / "nowhere.json")  # dangling
+    assert client.get(f"/pages/{page.id}").status_code == 404
+    assert main(["--dir", str(root), "revoke", "--", page.id]) == 0
+    assert (
+        not record.is_symlink() and client.get(f"/pages/{page.id}").status_code == 404
+    )