diff --git a/headroom/memory/storage_router.py b/headroom/memory/storage_router.py index 5337e32f7..45f980584 100644 --- a/headroom/memory/storage_router.py +++ b/headroom/memory/storage_router.py @@ -160,7 +160,15 @@ class ProjectResolver: if explicit: safe = self._sanitize_basename(explicit) if safe: - return safe, explicit + # Append a digest of the raw id like `_identity_from_cwd` does: + # `_sanitize_basename` maps every disallowed character to a dash + # and truncates to 64 chars, so distinct ids such as "acme/api" + # and "acme api" both collapse to "acme-api" and would otherwise + # share one project store (cross-project memory leak). The digest + # keeps distinct ids on distinct keys while the sanitized prefix + # stays human-readable on disk. + digest = hashlib.sha256(explicit.encode("utf-8")).hexdigest()[:16] + return f"{safe}-{digest}", explicit # Tier 2: client-provided explicit cwd (any client). explicit_cwd = self._first_nonempty_header(ctx.headers, "x-headroom-cwd") @@ -294,12 +302,20 @@ class BackendRouter: if mode is MemoryStorageMode.USER: user_safe = ProjectResolver._sanitize_basename(ctx.base_user_id) or "default" - db_path = self._config.root_dir / "users" / user_safe / "memory.db" + # Append a digest of the raw user id for the same reason as the + # project keys above: `_sanitize_basename` collapses distinct ids + # ("alice/qa", "alice qa", "alice@qa") to the same "alice-qa", which + # in USER mode would pool two different users into one memory.db — + # a cross-user data-isolation leak, the one thing USER mode exists to + # prevent. The digest keeps distinct users on distinct stores. + digest = hashlib.sha256(ctx.base_user_id.encode("utf-8")).hexdigest()[:16] + user_key = f"{user_safe}-{digest}" + db_path = self._config.root_dir / "users" / user_key / "memory.db" return ResolvedScope( mode=MemoryStorageMode.USER, db_path=db_path, display_name=ctx.base_user_id, - project_key=user_safe, + project_key=user_key, ) # PROJECT mode. diff --git a/tests/test_memory_storage_router.py b/tests/test_memory_storage_router.py index 9a1c7857c..90b36c946 100644 --- a/tests/test_memory_storage_router.py +++ b/tests/test_memory_storage_router.py @@ -48,10 +48,25 @@ def test_resolver_tier1_explicit_project_id_wins() -> None: ) assert out is not None key, display = out - assert key == "billing-svc" + # The sanitized id stays as a human-readable prefix; a sha256 digest is + # appended so distinct ids that sanitize alike cannot collide. + assert key.startswith("billing-svc-") + assert len(key.split("-")[-1]) == 16 assert display == "billing-svc" +def test_resolver_tier1_distinct_ids_that_sanitize_alike_dont_collide() -> None: + r = ProjectResolver() + # "acme/api" and "acme api" both sanitize to "acme-api"; without the digest + # they would share one project store (cross-project memory leak). + k1, _ = r.resolve(_ctx(headers={"x-headroom-project-id": "acme/api"})) # type: ignore[misc] + k2, _ = r.resolve(_ctx(headers={"x-headroom-project-id": "acme api"})) # type: ignore[misc] + assert k1 != k2 + # Same id resolves to a stable key across calls. + k1b, _ = r.resolve(_ctx(headers={"x-headroom-project-id": "acme/api"})) # type: ignore[misc] + assert k1 == k1b + + def test_resolver_tier2_explicit_cwd_header() -> None: r = ProjectResolver() out = r.resolve(_ctx(headers={"x-headroom-cwd": "/Users/foo/code/project-b"})) @@ -353,6 +368,20 @@ def test_router_user_mode_partitions_by_user( assert scope_b.display_name == "bob" +def test_router_user_mode_distinct_ids_that_sanitize_alike_dont_collide( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # "alice/qa" and "alice qa" both sanitize to "alice-qa"; without the digest + # they would share one users/alice-qa/memory.db — a cross-user leak, the one + # thing USER mode exists to prevent. + router = _make_router(tmp_path, MemoryStorageMode.USER, monkeypatch) + + _, scope_a = router.backend_for(_ctx(base_user_id="alice/qa")) + _, scope_b = router.backend_for(_ctx(base_user_id="alice qa")) + + assert scope_a.db_path != scope_b.db_path + + def test_router_global_mode_reuses_legacy_path( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index 7845396b0..7fb8f442e 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -896,7 +896,8 @@ def test_resolve_ccr_workspace_explicit_project_id_wins() -> None: request = _fake_request({"x-headroom-project-id": "my-cool-project"}) body = {} key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body) - assert key == "my-cool-project" + assert key.startswith("my-cool-project-") + assert len(key.split("-")[-1]) == 16 assert label == "my-cool-project"