mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(memory): make explicit-project and user store keys collision-resistant (#2231)
## Description
Two of the memory storage router's key-derivation paths can pool
distinct identities into one store.
`ProjectResolver._identity_from_cwd` builds a collision-resistant key by
appending a `sha256` digest to the sanitized basename:
```python
safe_basename = cls._sanitize_basename(basename) or "project"
digest = hashlib.sha256(normalised.encode("utf-8")).hexdigest()[:16]
key = f"{safe_basename}-{digest}"
```
But the two non-cwd paths use the bare sanitized basename as the key:
```python
# Tier 1 — explicit x-headroom-project-id
safe = self._sanitize_basename(explicit)
if safe:
return safe, explicit # <-- no digest
# USER mode
user_safe = ProjectResolver._sanitize_basename(ctx.base_user_id) or "default"
db_path = self._config.root_dir / "users" / user_safe / "memory.db" # <-- no digest
```
`_sanitize_basename` maps every disallowed character to a single dash,
so distinct inputs collapse to the same basename:
- `acme/api` and `acme api` (and `acme@api`) all → `acme-api`
- user ids `alice/qa` and `alice qa` → `alice-qa`
Both the project key (`root/projects/<key>/memory.db`) and the USER key
(`root/users/<key>/memory.db`) are derived directly from that basename,
so two distinct project ids — or, in USER mode, two distinct **users** —
resolve to the same `memory.db` and share each other's memories. USER
mode exists specifically to isolate users, so this is a cross-user
data-isolation leak; the explicit-project-id path is the same leak
across projects. Both are client-controlled (`x-headroom-project-id` /
`x-headroom-user-id` headers), so the collision is easy to hit and could
even be provoked deliberately.
## Fix
Append the same digest of the raw id to both keys, exactly as
`_identity_from_cwd` does, keeping the sanitized basename as a
human-readable prefix:
```python
digest = hashlib.sha256(explicit.encode("utf-8")).hexdigest()[:16]
return f"{safe}-{digest}", explicit
```
```python
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"
```
Distinct ids now always land on distinct stores; the same id remains
stable across calls.
**Migration note:** this changes the on-disk key format for the
explicit-project and USER stores (`<basename>` → `<basename>-<digest>`).
Memories written under the old bare-basename paths are not migrated; the
router will start a fresh store at the new path. GLOBAL and cwd-derived
PROJECT stores (which already carried the digest) are unaffected.
Flagging this explicitly so you can decide whether a migration shim is
wanted before merge.
Closes #
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/memory/storage_router.py`: append a `sha256` digest to the
explicit-project-id key (Tier 1) and the USER-mode key, matching
`_identity_from_cwd`.
- `tests/test_memory_storage_router.py`: update the Tier-1 key assertion
to the prefix+digest form; add collision regression tests for the
explicit-project and USER paths.
- `CHANGELOG.md`: Bug Fixes entry (including the migration note).
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/memory/storage_router.py tests/test_memory_storage_router.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/memory/storage_router.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the key derivation with a dependency-free script mirroring
`_sanitize_basename` + the digest, and left the full pytest to CI.
- Exact command / steps: derived keys for `alice/qa` and `alice qa`
under the OLD bare-basename scheme and the NEW digest scheme.
- Observed result: OLD → both `alice-qa` (identical → shared store); NEW
→ `alice-qa-7e02fc2dfbc447b4` vs `alice-qa-4c9241514a374ba3` (distinct),
stable per input, with the `alice-qa-` prefix retained.
- Not tested: a live proxy with two colliding tenants; full local
`pytest` deferred to CI (OOM).
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" box is unchecked because the full suite
imports the ML stack, which I can't run here. The changed/added tests
use the existing `tests/test_memory_storage_router.py` harness so they
run under the normal CI pytest job; behaviour is additionally verified
by the standalone proof above. I updated
`test_resolver_tier1_explicit_project_id_wins` to assert the new
prefix+digest key. Happy to add a migration shim (read the old path if
the new one is empty) if you'd prefer that over the fresh-store
behavior.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
29d8a5e563
commit
f840d5f2fe
3 changed files with 51 additions and 5 deletions
|
|
@ -160,7 +160,15 @@ class ProjectResolver:
|
||||||
if explicit:
|
if explicit:
|
||||||
safe = self._sanitize_basename(explicit)
|
safe = self._sanitize_basename(explicit)
|
||||||
if safe:
|
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).
|
# Tier 2: client-provided explicit cwd (any client).
|
||||||
explicit_cwd = self._first_nonempty_header(ctx.headers, "x-headroom-cwd")
|
explicit_cwd = self._first_nonempty_header(ctx.headers, "x-headroom-cwd")
|
||||||
|
|
@ -294,12 +302,20 @@ class BackendRouter:
|
||||||
|
|
||||||
if mode is MemoryStorageMode.USER:
|
if mode is MemoryStorageMode.USER:
|
||||||
user_safe = ProjectResolver._sanitize_basename(ctx.base_user_id) or "default"
|
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(
|
return ResolvedScope(
|
||||||
mode=MemoryStorageMode.USER,
|
mode=MemoryStorageMode.USER,
|
||||||
db_path=db_path,
|
db_path=db_path,
|
||||||
display_name=ctx.base_user_id,
|
display_name=ctx.base_user_id,
|
||||||
project_key=user_safe,
|
project_key=user_key,
|
||||||
)
|
)
|
||||||
|
|
||||||
# PROJECT mode.
|
# PROJECT mode.
|
||||||
|
|
|
||||||
|
|
@ -48,10 +48,25 @@ def test_resolver_tier1_explicit_project_id_wins() -> None:
|
||||||
)
|
)
|
||||||
assert out is not None
|
assert out is not None
|
||||||
key, display = out
|
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"
|
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:
|
def test_resolver_tier2_explicit_cwd_header() -> None:
|
||||||
r = ProjectResolver()
|
r = ProjectResolver()
|
||||||
out = r.resolve(_ctx(headers={"x-headroom-cwd": "/Users/foo/code/project-b"}))
|
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"
|
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(
|
def test_router_global_mode_reuses_legacy_path(
|
||||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||||
) -> None:
|
) -> None:
|
||||||
|
|
|
||||||
|
|
@ -896,7 +896,8 @@ def test_resolve_ccr_workspace_explicit_project_id_wins() -> None:
|
||||||
request = _fake_request({"x-headroom-project-id": "my-cool-project"})
|
request = _fake_request({"x-headroom-project-id": "my-cool-project"})
|
||||||
body = {}
|
body = {}
|
||||||
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, 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"
|
assert label == "my-cool-project"
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue