From 9627828795ede04116f17b083e814c976139a912 Mon Sep 17 00:00:00 2001 From: Robert Sigmundsson Date: Wed, 15 Jul 2026 21:59:58 +0200 Subject: [PATCH] fix(storage): route id sanitisers through single source + brain-scope every record-id fetch, incl. concurrent batch fetches (SECURITY: cross-brain IDOR, adapted for 2.10.5 concurrency) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-source sanitisation: * migrations._sanitize_id (dash-only .replace) and an inline .replace("-","_") in tool_events were the last two id sanitisers bypassing the _ids.py choke point. Both now call _to_surreal_id (full [A-Za-z0-9_] fold), which also fixes the latent dotted-id -> non-matching-edge mismatch in migrations. Brain-scoped record fetches (cross-brain IDOR): * Every fetch-by-record-id did a bare select with no brain filter, so a caller who knows/guesses another brain's record id could read it (confirmed path: server/routes/hub.py switches brain from body.brain_id validating only the char format, app.py calls find_neurons_by_ids with client ids). All now add WHERE brain_id = $brain_id (bound param; brain_id from _get_brain_id -> _safe_brain_id, fail-closed): get_neuron, get_neuron_state, get_synapse, get_fiber, find_neurons_by_ids. They stay direct record fetches (ids pinned in FROM, not a table scan), so the dashboard-perf property is preserved. * get_neurons_batch (v2.10.5 concurrent: semaphore + gather) keeps its concurrency — the bare conn.select inside _fetch_one becomes a scoped self._query, with brain_id resolved ONCE before the semaphore/gather so an unsafe brain context fails closed before a single fetch goes out. * get_synapses_batch (NEW in v2.10.5, had the same unscoped concurrent-fetch pattern, absent from the original fc88973) gets the identical adaptation. Tests: brain-scope query shape + bound param for all single-record methods and both concurrent batch methods (asserted over conn.query.call_args_list, N calls); fail-closed brain-id for both batch methods; single-source routing for migrations (dotted-id fold) and tool_events; dashboard-perf assertion updated to expect WHERE brain_id while still forbidding a FROM-neuron table scan. (adapted from commit fc889732b8998cead61250f9ab528858c6e4cd41; extends scope to get_synapses_batch, new in v2.10.5, which had the same unscoped-fetch pattern) (cherry picked from commit 8bca0c7a6d40df2cbc70ffbebadb9a8f96dd22a8) --- .../storage/surrealdb/migrations.py | 16 +- src/surreal_memory/storage/surrealdb/store.py | 91 +++++--- .../storage/surrealdb/tool_events.py | 3 +- tests/unit/test_dashboard_perf_queries.py | 5 +- .../test_id_consolidation_and_brain_scope.py | 196 ++++++++++++++++++ 5 files changed, 272 insertions(+), 39 deletions(-) create mode 100644 tests/unit/test_id_consolidation_and_brain_scope.py diff --git a/src/surreal_memory/storage/surrealdb/migrations.py b/src/surreal_memory/storage/surrealdb/migrations.py index 82918dbd..81743a84 100644 --- a/src/surreal_memory/storage/surrealdb/migrations.py +++ b/src/surreal_memory/storage/surrealdb/migrations.py @@ -35,6 +35,7 @@ import uuid from typing import Any +from surreal_memory.storage.surrealdb._ids import _to_surreal_id from surreal_memory.storage.surrealdb.connection import ( MIN_SERVER_VERSION, parse_server_version, @@ -123,17 +124,6 @@ def _already_exists(exc: Exception) -> bool: return "already exists" in str(exc).lower() -def _sanitize_id(raw: str) -> str: - """Strip a table prefix and normalise dashes to underscores. - - Mirrors ``store._to_surreal_id`` so migrated endpoint ids match what the - query layer produces (``neuron:abc-123`` / ``abc-123`` -> ``abc_123``). - """ - if ":" in raw: - raw = raw.rsplit(":", 1)[1] - return raw.replace("-", "_") - - def _record_part(rid: Any) -> str: """Extract the identifier part of a RecordID (or a ``table:id`` string).""" part = getattr(rid, "id", None) @@ -325,8 +315,8 @@ def _to_relation_row(row: dict[str, Any]) -> dict[str, Any]: from surrealdb import RecordID # lazy sid = _record_part(row["id"]) - src = _sanitize_id(str(row.get("source_id", ""))) - tgt = _sanitize_id(str(row.get("target_id", ""))) + src = _to_surreal_id(str(row.get("source_id", ""))) + tgt = _to_surreal_id(str(row.get("target_id", ""))) relation: dict[str, Any] = { "id": RecordID("synapse", sid), "in": RecordID("neuron", src), diff --git a/src/surreal_memory/storage/surrealdb/store.py b/src/surreal_memory/storage/surrealdb/store.py index 81016f7e..9598eff8 100755 --- a/src/surreal_memory/storage/surrealdb/store.py +++ b/src/surreal_memory/storage/surrealdb/store.py @@ -636,12 +636,17 @@ async def add_neuron(self, neuron: Neuron) -> str: return neuron.id async def get_neuron(self, neuron_id: str) -> Neuron | None: - conn = self._ensure_conn() + # Scope to the current brain: a bare record select would let a caller + # read another brain's neuron by id. The record is still pinned in FROM. + brain_id = self._get_brain_id() sid = _to_surreal_id(neuron_id) try: - result = await conn.select(f"neuron:{sid}") - if result and isinstance(result, list) and len(result) > 0: - return _row_to_neuron(result[0]) + rows = await self._query( + f"SELECT * FROM neuron:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_neuron(rows[0]) except Exception: pass return None @@ -660,17 +665,29 @@ async def get_neurons_batch(self, neuron_ids: list[str]) -> dict[str, Neuron]: """ if not neuron_ids: return {} + # Scope every fetch to the current brain: these are direct record-id + # selects, so without a brain_id filter a caller could read a neuron + # owned by another brain just by knowing (or guessing) its id. The FROM + # is still a pinned record (neuron:{sid}), so this stays a direct fetch, + # not a table scan — the brain_id check only rejects a cross-brain hit. + # Resolve the brain id ONCE, before any concurrent fetch starts, so a + # hostile/unset brain context fails closed (_safe_brain_id raises) rather + # than firing 16 unscoped selects. + brain_id = self._get_brain_id() semaphore = asyncio.Semaphore(_BATCH_FETCH_CONCURRENCY) async def _fetch_one(nid: str) -> tuple[str, Neuron | None]: sid = _to_surreal_id(nid) async with semaphore: try: - result = await self._conn.select(f"neuron:{sid}") + rows = await self._query( + f"SELECT * FROM neuron:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) except Exception: return nid, None - if result and isinstance(result, list) and len(result) > 0: - return nid, _row_to_neuron(result[0]) + if rows: + return nid, _row_to_neuron(rows[0]) return nid, None pairs = await asyncio.gather(*(_fetch_one(nid) for nid in neuron_ids)) @@ -790,12 +807,19 @@ async def find_neurons_by_ids( ] if not safe: return [] + # Scope to the current brain: the FROM is a pinned record-id list (still + # a direct fetch, not a table scan), and the WHERE brain_id filter keeps + # a caller from reading another brain's neuron by id. + brain_id = self._get_brain_id() projection = "SELECT *" if include_embedding else "SELECT * OMIT embedding_vec" out: list[Neuron] = [] # Chunk to keep the FROM record-id list a sane query size. for i in range(0, len(safe), 1000): things = ", ".join(f"neuron:{s}" for s in safe[i : i + 1000]) - rows = await self._query(f"{projection} FROM {things}") + rows = await self._query( + f"{projection} FROM {things} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) out.extend(_row_to_neuron(r) for r in rows) return out @@ -939,12 +963,15 @@ async def has_neuron_by_content_hash(self, content_hash: int) -> bool: # ================================================================ async def get_neuron_state(self, neuron_id: str) -> NeuronState | None: - conn = self._ensure_conn() + brain_id = self._get_brain_id() sid = _to_surreal_id(neuron_id) try: - result = await conn.select(f"neuron_state:state_{sid}") - if result: - return _row_to_neuron_state(result[0] if isinstance(result, list) else result) + rows = await self._query( + f"SELECT * FROM neuron_state:state_{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_neuron_state(rows[0]) except Exception: pass return None @@ -1080,12 +1107,15 @@ async def add_synapses_batch( return len(synapses) async def get_synapse(self, synapse_id: str) -> Synapse | None: - conn = self._ensure_conn() + brain_id = self._get_brain_id() sid = _to_surreal_id(synapse_id) try: - result = await conn.select(f"synapse:{sid}") - if result: - return _row_to_synapse(result[0] if isinstance(result, list) else result) + rows = await self._query( + f"SELECT * FROM synapse:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_synapse(rows[0]) except Exception: pass return None @@ -1098,18 +1128,28 @@ async def get_synapses_batch(self, synapse_ids: list[str]) -> dict[str, Synapse] """ if not synapse_ids: return {} - conn = self._ensure_conn() + # Scope every fetch to the current brain: like get_neurons_batch these are + # direct record-id selects, so without a brain_id filter a caller could + # read a synapse owned by another brain just by knowing (or guessing) its + # id. The FROM stays a pinned record (synapse:{sid}) — a direct fetch, not + # a table scan — and the WHERE brain_id only rejects a cross-brain hit. + # Resolve the brain id ONCE, before any concurrent fetch starts, so a + # hostile/unset brain context fails closed (_safe_brain_id raises). + brain_id = self._get_brain_id() semaphore = asyncio.Semaphore(_BATCH_FETCH_CONCURRENCY) async def _fetch_one(syn_id: str) -> tuple[str, Synapse | None]: sid = _to_surreal_id(syn_id) async with semaphore: try: - result = await conn.select(f"synapse:{sid}") + rows = await self._query( + f"SELECT * FROM synapse:{sid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) except Exception: return syn_id, None - if result: - return syn_id, _row_to_synapse(result[0] if isinstance(result, list) else result) + if rows: + return syn_id, _row_to_synapse(rows[0]) return syn_id, None pairs = await asyncio.gather(*(_fetch_one(sid) for sid in synapse_ids)) @@ -1415,12 +1455,15 @@ async def add_fiber(self, fiber: Fiber) -> str: return fiber.id async def get_fiber(self, fiber_id: str) -> Fiber | None: - conn = self._ensure_conn() + brain_id = self._get_brain_id() fid = _to_surreal_id(fiber_id) try: - result = await conn.select(f"fiber:{fid}") - if result: - return _row_to_fiber(result[0] if isinstance(result, list) else result) + rows = await self._query( + f"SELECT * FROM fiber:{fid} WHERE brain_id = $brain_id", + brain_id=brain_id, + ) + if rows: + return _row_to_fiber(rows[0]) except Exception: pass return None diff --git a/src/surreal_memory/storage/surrealdb/tool_events.py b/src/surreal_memory/storage/surrealdb/tool_events.py index 79f66be7..933af62a 100644 --- a/src/surreal_memory/storage/surrealdb/tool_events.py +++ b/src/surreal_memory/storage/surrealdb/tool_events.py @@ -15,6 +15,7 @@ from datetime import datetime, timedelta from typing import Any +from surreal_memory.storage.surrealdb._ids import _to_surreal_id from surreal_memory.utils.timeutils import utcnow # Cap per brain to prevent unbounded growth (mirrors the SQLite mixin). @@ -66,7 +67,7 @@ async def insert_tool_events( await conn.insert( "tool_events", { - "id": event_id.replace("-", "_"), + "id": _to_surreal_id(event_id), "event_id": event_id, "brain_id": brain_id, "tool_name": ev.get("tool_name", ""), diff --git a/tests/unit/test_dashboard_perf_queries.py b/tests/unit/test_dashboard_perf_queries.py index 2f3b1832..11811409 100644 --- a/tests/unit/test_dashboard_perf_queries.py +++ b/tests/unit/test_dashboard_perf_queries.py @@ -103,7 +103,10 @@ async def test_fetches_records_directly_with_omit(self): await st.find_neurons_by_ids(["abc-123", "def-456"]) (sql,) = _queries(conn) assert sql.startswith("SELECT * OMIT embedding_vec FROM neuron:abc_123, neuron:def_456") - assert "WHERE" not in sql # record fetch, not a table scan + # Still a pinned-record fetch, not a `FROM neuron` table scan… + assert "FROM neuron WHERE" not in sql + # …but scoped to the current brain (no cross-brain id read). + assert sql.endswith("WHERE brain_id = $brain_id") async def test_empty_and_unsafe_ids_are_dropped(self): st, conn = _store_with_mock_conn() diff --git a/tests/unit/test_id_consolidation_and_brain_scope.py b/tests/unit/test_id_consolidation_and_brain_scope.py new file mode 100644 index 00000000..e7df9022 --- /dev/null +++ b/tests/unit/test_id_consolidation_and_brain_scope.py @@ -0,0 +1,196 @@ +"""Single-source id sanitisation + brain-scoped record fetch. + +Two follow-ups on the id-sanitiser consolidation: + + * two weak sanitisers still bypassed the ``_ids.py`` single choke point — + ``migrations._sanitize_id`` (dash-only ``.replace``) and an inline + ``.replace("-", "_")`` in ``tool_events`` — so both now route through + ``_to_surreal_id``; + * every fetch-by-record-id (``get_neuron``, ``get_neuron_state``, + ``get_synapse``, ``get_fiber``, ``find_neurons_by_ids`` and the two + concurrent batch fetches ``get_neurons_batch`` / ``get_synapses_batch``) + issued a bare record select with no ``WHERE brain_id`` — a caller could read + another brain's record just by knowing its id — so all now scope to the + current brain (cross-brain IDOR). + +The concurrent batch fetches (v2.10.5, bounded ``asyncio.Semaphore`` + +``gather``) fire ONE ``self._query`` per id, so their brain-scope is asserted +over ``conn.query.call_args_list`` (N calls), not a single ``call_args``. The +brain id is resolved once, before any fetch starts, so an unsafe brain context +fails closed (``_safe_brain_id`` raises) before a single query goes out. + +The surrealdb SDK is stubbed (repo convention); the connection is an AsyncMock. +""" + +from __future__ import annotations + +import sys +from unittest.mock import AsyncMock, MagicMock + +if "surrealdb" not in sys.modules: + sys.modules["surrealdb"] = MagicMock() + sys.modules["surrealdb.errors"] = MagicMock() + +import surreal_memory.storage.surrealdb.tool_events as tool_events_mod +from surreal_memory.storage.surrealdb import migrations +from surreal_memory.storage.surrealdb._ids import _to_surreal_id +from surreal_memory.storage.surrealdb.store import SurrealDBStorage + + +def _store_with_mock_conn() -> tuple[SurrealDBStorage, AsyncMock]: + st = SurrealDBStorage(url="http://localhost:8001") + conn = AsyncMock() + conn.query = AsyncMock(return_value=[]) + conn.insert = AsyncMock(return_value=None) + st._conn = conn + st._current_brain_id = "b1" + return st, conn + + +# --------------------------------------------------------------------------- # +# brain-scoped batch fetch +# --------------------------------------------------------------------------- # +class TestBrainScopedBatchFetch: + async def test_find_neurons_by_ids_scopes_to_current_brain(self): + st, conn = _store_with_mock_conn() + await st.find_neurons_by_ids(["abc-123", "def-456"]) + sql, params = conn.query.call_args.args + # Still a pinned-record fetch (ids in FROM), never a `FROM neuron` scan… + assert sql.startswith("SELECT * OMIT embedding_vec FROM neuron:abc_123, neuron:def_456") + # …but now scoped, so a foreign-brain id cannot be read back. + assert sql.endswith("WHERE brain_id = $brain_id") + assert params == {"brain_id": "b1"} + + async def test_get_neurons_batch_scopes_to_current_brain(self): + # v2.10.5 made this concurrent (semaphore + gather): one self._query per + # id, so assert brain-scope over EVERY recorded query, not just the last. + st, conn = _store_with_mock_conn() + await st.get_neurons_batch(["abc-123", "def-456"]) + calls = conn.query.call_args_list + assert len(calls) == 2 + seen_ids = set() + for call in calls: + sql, params = call.args + assert sql.startswith("SELECT * FROM neuron:") + assert sql.endswith("WHERE brain_id = $brain_id") + assert params == {"brain_id": "b1"} + seen_ids.add(sql) + assert "SELECT * FROM neuron:abc_123 WHERE brain_id = $brain_id" in seen_ids + assert "SELECT * FROM neuron:def_456 WHERE brain_id = $brain_id" in seen_ids + + async def test_get_neurons_batch_uses_safe_brain_id(self): + # An unsafe brain context must fail closed (via _safe_brain_id) rather + # than inline a hostile id or fire a single unscoped select — the brain + # id is resolved once, before any concurrent fetch starts. + st, conn = _store_with_mock_conn() + st._current_brain_id = 'bad"; DELETE neuron' + import pytest + + with pytest.raises(ValueError): + await st.get_neurons_batch(["abc-123"]) + conn.query.assert_not_called() + + async def test_get_synapses_batch_scopes_to_current_brain(self): + # New in v2.10.5, same concurrent unscoped-fetch pattern as + # get_neurons_batch — assert brain-scope over every recorded query. + st, conn = _store_with_mock_conn() + await st.get_synapses_batch(["s-1", "s-2"]) + calls = conn.query.call_args_list + assert len(calls) == 2 + seen = set() + for call in calls: + sql, params = call.args + assert sql.startswith("SELECT * FROM synapse:") + assert sql.endswith("WHERE brain_id = $brain_id") + assert params == {"brain_id": "b1"} + seen.add(sql) + assert "SELECT * FROM synapse:s_1 WHERE brain_id = $brain_id" in seen + assert "SELECT * FROM synapse:s_2 WHERE brain_id = $brain_id" in seen + + async def test_get_synapses_batch_uses_safe_brain_id(self): + st, conn = _store_with_mock_conn() + st._current_brain_id = 'bad"; DELETE synapse' + import pytest + + with pytest.raises(ValueError): + await st.get_synapses_batch(["s-1"]) + conn.query.assert_not_called() + + +# --------------------------------------------------------------------------- # +# single-source sanitisation +# --------------------------------------------------------------------------- # +class TestSanitizerConsolidation: + def test_migrations_has_no_local_sanitiser(self): + # The dash-only duplicate is gone; the single source is _ids._to_surreal_id. + assert not hasattr(migrations, "_sanitize_id") + + def test_relation_row_folds_via_single_source(self, monkeypatch): + # A dotted endpoint id must fold '.' -> '_' (full [A-Za-z0-9_] fold), + # which the removed dash-only sanitiser did NOT do. This both proves the + # single source is used and fixes a latent bug: neuron records are stored + # folded, so a dotted legacy id used to migrate to a non-matching edge. + rid_calls: list[tuple[str, str]] = [] + + def _fake_rid(table: str, ident: object) -> str: + rid_calls.append((table, str(ident))) + return f"{table}:{ident}" + + monkeypatch.setattr(sys.modules["surrealdb"], "RecordID", _fake_rid) + migrations._to_relation_row( + { + "id": "synapse:s1", + "source_id": "a.b-c", + "target_id": "x-y", + "brain_id": "default", + "type": "assoc", + "weight": 1.0, + } + ) + assert ("neuron", _to_surreal_id("a.b-c")) in rid_calls # -> a_b_c + assert ("neuron", "a_b_c") in rid_calls + assert ("neuron", "x_y") in rid_calls + + async def test_tool_events_id_routes_through_single_source(self, monkeypatch): + st, conn = _store_with_mock_conn() + monkeypatch.setattr(tool_events_mod, "_to_surreal_id", lambda s: "SENTINEL") + await st.insert_tool_events( + "b1", [{"tool_name": "t", "created_at": "2026-01-01T00:00:00+00:00"}] + ) + table, doc = conn.insert.call_args.args + assert table == "tool_events" + # The record id came from the single-source folder, not a raw .replace. + assert doc["id"] == "SENTINEL" + + +# --------------------------------------------------------------------------- # +# single-record fetches close the same read-by-id class +# --------------------------------------------------------------------------- # +class TestSingleRecordBrainScope: + async def test_get_neuron_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_neuron("abc-123") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM neuron:abc_123 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_neuron_state_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_neuron_state("abc-123") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM neuron_state:state_abc_123 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_synapse_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_synapse("s-1") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM synapse:s_1 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"} + + async def test_get_fiber_scopes_to_brain(self): + st, conn = _store_with_mock_conn() + await st.get_fiber("f-1") + sql, params = conn.query.call_args.args + assert sql == "SELECT * FROM fiber:f_1 WHERE brain_id = $brain_id" + assert params == {"brain_id": "b1"}