From 4c74bc0005f770b2a0d46072c5eefe6ed900070b Mon Sep 17 00:00:00 2001 From: Sarthak Agrawal Date: Fri, 12 Jun 2026 03:25:51 +0530 Subject: [PATCH 1/3] fix raw object key idempotency --- src/kb/storage/objects.py | 4 ++-- tests/test_objects.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) create mode 100644 tests/test_objects.py diff --git a/src/kb/storage/objects.py b/src/kb/storage/objects.py index a11d1e5..aae93be 100644 --- a/src/kb/storage/objects.py +++ b/src/kb/storage/objects.py @@ -2,7 +2,7 @@ Both raw files and cached parse artifacts (Unstructured element JSON) live here. Keying convention: - raw/// + raw// parse//elements.json """ @@ -125,7 +125,7 @@ def _get_backend() -> _Backend: async def put_raw_file(*, domain: str, filename: str, blob: bytes) -> tuple[str, str]: """Store a raw upload; returns (object_key, content_hash). Idempotent by content_hash.""" h = _content_hash(blob) - key = f"raw/{domain}/{h}/{filename}" + key = f"raw/{domain}/{h}" backend = _get_backend() if not await backend.exists(key): await backend.put(key, blob) diff --git a/tests/test_objects.py b/tests/test_objects.py new file mode 100644 index 0000000..d7837f0 --- /dev/null +++ b/tests/test_objects.py @@ -0,0 +1,33 @@ +"""Raw object store keying and idempotency.""" + +from __future__ import annotations + +import asyncio +import hashlib + +from kb.storage import objects + + +class FakeBackend: + def __init__(self) -> None: + self.keys: list[str] = [] + self._exists: set[str] = set() + + async def exists(self, key: str) -> bool: + return key in self._exists + + async def put(self, key: str, blob: bytes, mime: str | None = None) -> None: + self.keys.append(key) + self._exists.add(key) + + +def test_put_raw_file_is_idempotent_by_content_hash(monkeypatch) -> None: + backend = FakeBackend() + monkeypatch.setattr(objects, "_backend", backend) + + first = asyncio.run(objects.put_raw_file(domain="sec", filename="a.pdf", blob=b"same")) + second = asyncio.run(objects.put_raw_file(domain="sec", filename="b.pdf", blob=b"same")) + + expected_key = f"raw/sec/{hashlib.sha256(b'same').hexdigest()}" + assert first == second + assert backend.keys == [expected_key] From 8f56766129ecce79020dc83bbc3746efbc9b72fb Mon Sep 17 00:00:00 2001 From: Sarthak Agrawal Date: Fri, 12 Jun 2026 03:44:45 +0530 Subject: [PATCH 2/3] fix local object store key safety --- src/kb/storage/objects.py | 9 ++++++++- tests/test_objects.py | 7 +++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/kb/storage/objects.py b/src/kb/storage/objects.py index aae93be..60a4c4e 100644 --- a/src/kb/storage/objects.py +++ b/src/kb/storage/objects.py @@ -24,6 +24,13 @@ def _content_hash(blob: bytes) -> str: return hashlib.sha256(blob).hexdigest() +def _safe_key_segment(segment: str) -> str: + """Reject path-traversal-shaped segments before building object keys.""" + if not segment or segment in {".", ".."} or "/" in segment or "\\" in segment: + raise ValueError(f"unsafe object-store key segment: {segment!r}") + return segment + + # ─── Backend interface ──────────────────────────────────────────────────── class _Backend: async def put(self, key: str, blob: bytes, mime: str | None = None) -> None: ... @@ -125,7 +132,7 @@ def _get_backend() -> _Backend: async def put_raw_file(*, domain: str, filename: str, blob: bytes) -> tuple[str, str]: """Store a raw upload; returns (object_key, content_hash). Idempotent by content_hash.""" h = _content_hash(blob) - key = f"raw/{domain}/{h}" + key = f"raw/{_safe_key_segment(domain)}/{h}" backend = _get_backend() if not await backend.exists(key): await backend.put(key, blob) diff --git a/tests/test_objects.py b/tests/test_objects.py index d7837f0..00d6e2f 100644 --- a/tests/test_objects.py +++ b/tests/test_objects.py @@ -5,6 +5,8 @@ import asyncio import hashlib +import pytest + from kb.storage import objects @@ -31,3 +33,8 @@ def test_put_raw_file_is_idempotent_by_content_hash(monkeypatch) -> None: expected_key = f"raw/sec/{hashlib.sha256(b'same').hexdigest()}" assert first == second assert backend.keys == [expected_key] + + +def test_put_raw_file_rejects_unsafe_domain_segment() -> None: + with pytest.raises(ValueError): + asyncio.run(objects.put_raw_file(domain="../evil", filename="a.pdf", blob=b"same")) From ee2383077d63c36a94230e48b9f1d55ed64a8ac5 Mon Sep 17 00:00:00 2001 From: Sarthak Agrawal Date: Fri, 12 Jun 2026 13:29:54 +0530 Subject: [PATCH 3/3] Fix pgvector parent_id scope filtering --- src/kb/vector/pgvector_store.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/kb/vector/pgvector_store.py b/src/kb/vector/pgvector_store.py index 232d33f..1c081f3 100644 --- a/src/kb/vector/pgvector_store.py +++ b/src/kb/vector/pgvector_store.py @@ -12,7 +12,11 @@ from kb.vector.embed import embed_dense # Whitelist filter columns — anything not in here is silently dropped to keep SQL safe. -_ALLOWED_FILTER_COLS = {"project", "file_id", "entity_id", "parent_chunk", "domain"} +# +# `parent_id` is the public query-scoping key used by callers; `parent_chunk` +# is the backing column in `chunks`. Support both so scope filters behave the +# same across vector backends. +_ALLOWED_FILTER_COLS = {"project", "file_id", "entity_id", "parent_id", "parent_chunk", "domain"} def _vec_literal(v: list[float]) -> str: @@ -140,7 +144,8 @@ async def hybrid_search( for k, v in filters.items(): if v is None or k not in _ALLOWED_FILTER_COLS: continue - where.append(f"{k} = :flt_{k}") + col = "parent_chunk" if k == "parent_id" else k + where.append(f"{col} = :flt_{k}") params[f"flt_{k}"] = v wsql = " AND ".join(where)