From ef9cacf981e7bd2dfafb7d745986b5f6584e5367 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 25 May 2026 17:36:31 +0300 Subject: [PATCH] Harden FS source against symlink escape; surface pool/poller liveness in /health, additional auth tests --- .../haiku/rag/ingester/api/routes/health.py | 16 +++- .../haiku/rag/ingester/api/schemas.py | 5 ++ .../haiku/rag/ingester/pollers/manager.py | 6 ++ .../haiku/rag/ingester/sources/fs.py | 51 ++++++++++--- .../haiku/rag/ingester/workers/pool.py | 6 ++ .../haiku/rag/store/upgrades/v0_50_0.py | 7 +- tests/ingester/test_api.py | 56 ++++++++++++++ tests/ingester/test_fs_source.py | 74 +++++++++++++++++++ 8 files changed, 206 insertions(+), 15 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py b/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py index 573dfb20..e6f1fd79 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/routes/health.py @@ -9,15 +9,27 @@ router = APIRouter() @router.get("/health", response_model=HealthResponse) async def health(state: APIState = Depends(get_state)) -> HealthResponse: """Liveness signal + queue/worker overview. Unauthenticated so load - balancers and uptime monitors can hit it without a token.""" + balancers and uptime monitors can hit it without a token. + + `status="degraded"` when at least one configured worker or poller task + has died (workers_alive < worker_count or pollers_alive < poller_count). + The DB query alone can't catch that — workers might be all dead and the + queue would still report counts cheerfully — so this branch is what + actually distinguishes "alive" from "alive but doing nothing useful". + """ counts = await state.job_repo.counts_by_status() worker_count = ( state.config.ingester.workers.worker_count if state.pool is not None else 0 ) + workers_alive = state.pool.live_workers if state.pool is not None else 0 poller_count = len(state.pollers.pollers) if state.pollers is not None else 0 + pollers_alive = state.pollers.live_pollers if state.pollers is not None else 0 + degraded = workers_alive < worker_count or pollers_alive < poller_count return HealthResponse( - status="ok", + status="degraded" if degraded else "ok", queue_counts=counts, worker_count=worker_count, poller_count=poller_count, + workers_alive=workers_alive, + pollers_alive=pollers_alive, ) diff --git a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py index 6a94ea24..4e1f5caa 100644 --- a/haiku_rag_slim/haiku/rag/ingester/api/schemas.py +++ b/haiku_rag_slim/haiku/rag/ingester/api/schemas.py @@ -8,6 +8,11 @@ class HealthResponse(BaseModel): queue_counts: dict[str, int] worker_count: int poller_count: int + # Live counters — non-zero shortfalls vs the configured count signal a + # crashed task. status="degraded" when either shortfall is non-zero so + # uptime monitors can alert without needing to do the math themselves. + workers_alive: int + pollers_alive: int class SourceSummary(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py index e9593f91..d058eb71 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/manager.py @@ -87,3 +87,9 @@ class PollerManager: @property def pollers(self) -> list[BasePoller]: return list(self._pollers) + + @property + def live_pollers(self) -> int: + """Poller tasks that are still running. Equal to len(pollers) under + normal operation; less when a poller has crashed.""" + return sum(1 for t in self._tasks if not t.done()) diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 91878a42..7159b950 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -1,10 +1,12 @@ import hashlib import mimetypes +import os from collections.abc import AsyncIterator from datetime import UTC, datetime from pathlib import Path from urllib.parse import unquote, urlparse +from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.ingester.sources.base import ( FetchResult, RevisionSnapshot, @@ -51,27 +53,38 @@ class FSSource: supported_extensions=self.supported_extensions, ) + def _resolve_within_root(self, uri: str) -> Path | None: + """Resolve a URI to a real path guaranteed to live under ``self.root``. + + Returns ``None`` if the URI parses but resolves outside the root + (path-traversal via ``..``, symlinks pointing elsewhere). Callers + treat this as "not ours" — `supports()` returns False, `head()` + returns None, `fetch()` raises ``UnsupportedSourceError``. + """ + try: + path = _uri_to_path(uri).resolve(strict=False) + except (ValueError, OSError): + return None + if not path.is_relative_to(self.root): + return None + return path + def supports(self, uri: str) -> bool: scheme = urlparse(uri).scheme if scheme not in ("", "file"): return False - try: - _uri_to_path(uri) - except ValueError: - return False - return True + return self._resolve_within_root(uri) is not None async def head(self, uri: str) -> str | None: - path = _uri_to_path(uri).absolute() - if not path.exists(): + path = self._resolve_within_root(uri) + if path is None or not path.exists(): return None return str(path.stat().st_mtime_ns) async def fetch(self, uri: str) -> FetchResult: - # Absolute path is needed for as_uri() and matches the old - # _create_document_from_file behavior (which keyed docs on the - # absolute file:// URI). - path = _uri_to_path(uri).absolute() + path = self._resolve_within_root(uri) + if path is None: + raise UnsupportedSourceError(f"Path escapes FS root ({self.root}): {uri}") body = path.read_bytes() content_type, _ = mimetypes.guess_type(path.name) if content_type is None: @@ -95,7 +108,21 @@ class FSSource: now = datetime.now(UTC) seen: set[str] = set() - for path in sorted(self.root.rglob("*")): + # os.walk with followlinks=False so symlinked directories aren't + # traversed. Then per-file: skip individual file-symlinks too, since + # they could point outside root and reading them would leak data. + # Operators wanting to ingest content from outside root should + # bind-mount it in or configure a second source. + candidates: list[Path] = [] + for dirpath, _dirnames, filenames in os.walk(self.root, followlinks=False): + for filename in filenames: + path = Path(dirpath) / filename + if path.is_symlink(): + continue + candidates.append(path) + candidates.sort() + + for path in candidates: if not path.is_file(): continue if not self.filter.include_file(str(path)): diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index edbd1aef..a8dbf475 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -50,6 +50,12 @@ class WorkerPool: self._workers: list[asyncio.Task] = [] self._reaper: asyncio.Task | None = None + @property + def live_workers(self) -> int: + """Worker tasks that are still running. Equal to worker_count under + normal operation; less when a worker has crashed.""" + return sum(1 for t in self._workers if not t.done()) + async def start(self) -> None: if self._workers: raise RuntimeError("WorkerPool already started") diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_50_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_50_0.py index 9c574c44..36096e3f 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_50_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_50_0.py @@ -63,8 +63,13 @@ async def _apply_canonical_metadata_keys(store: Store) -> None: try: meta = json.loads(raw) except Exception: + # exc_info=True so the actual JSONDecodeError reaches the logs; + # otherwise a permanently malformed row stays malformed forever + # and the operator has nothing to grep for. logger.warning( - "Could not parse metadata JSON for document %s; skipping", doc_id + "Could not parse metadata JSON for document %s; skipping", + doc_id, + exc_info=True, ) skipped += 1 continue diff --git a/tests/ingester/test_api.py b/tests/ingester/test_api.py index def21535..974fa88f 100644 --- a/tests/ingester/test_api.py +++ b/tests/ingester/test_api.py @@ -71,6 +71,31 @@ async def test_health_ok_with_counts(state, jobs): assert body["queue_counts"] == {"queued": 1, "dead": 1} assert body["worker_count"] == 0 # pool not attached in the test state assert body["poller_count"] == 0 + assert body["workers_alive"] == 0 + assert body["pollers_alive"] == 0 + + +@pytest.mark.asyncio +async def test_health_degraded_when_worker_died(jobs, sync): + """If a worker task crashed (live_workers < worker_count), /health must + flip to status='degraded' so uptime monitors notice.""" + from unittest.mock import MagicMock + + from haiku.rag.config import AppConfig + + config = AppConfig() + config.ingester.workers.worker_count = 4 + + pool = MagicMock() + pool.live_workers = 3 # one dead + + state = APIState(config=config, job_repo=jobs, sync_repo=sync, pool=pool) + async with _client(state) as client: + resp = await client.get("/health") + body = resp.json() + assert body["status"] == "degraded" + assert body["worker_count"] == 4 + assert body["workers_alive"] == 3 @pytest.mark.asyncio @@ -111,6 +136,37 @@ async def test_no_auth_token_allows_everything(state, jobs): assert (await client.get("/health")).status_code == 200 +@pytest.mark.asyncio +async def test_mutation_endpoints_require_auth(state, jobs): + """Existing tests prove auth gates GETs; this pins that the *mutation* + endpoints (retry, cancel, DLQ requeue, source refresh) also require the + bearer. A missing-auth regression on these would silently let anyone + cancel jobs or reset the DLQ.""" + j = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert j is not None + await jobs.mark_dead(j.id, "boom") + + async with _client(state, auth_token="secret") as client: + # Cancel: blocked without token + resp = await client.delete(f"/jobs/{j.id}") + assert resp.status_code == 401 + # Retry: blocked without token + resp = await client.post(f"/jobs/{j.id}/retry") + assert resp.status_code == 401 + # DLQ retry: blocked without token + resp = await client.post(f"/dlq/{j.id}/retry") + assert resp.status_code == 401 + # Source refresh: blocked without token + resp = await client.post("/sources/anything/refresh") + assert resp.status_code == 401 + + # With correct token: 200 for retry (job is dead, gets resurrected). + ok = await client.post( + f"/jobs/{j.id}/retry", headers={"Authorization": "Bearer secret"} + ) + assert ok.status_code == 200 + + # --- /jobs --- diff --git a/tests/ingester/test_fs_source.py b/tests/ingester/test_fs_source.py index 1ef7cfc0..816b2fbb 100644 --- a/tests/ingester/test_fs_source.py +++ b/tests/ingester/test_fs_source.py @@ -3,6 +3,7 @@ from pathlib import Path import pytest +from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.ingester.sources.base import SourceEventKind from haiku.rag.ingester.sources.fs import FSSource @@ -152,3 +153,76 @@ async def test_fs_source_discover_respects_include_patterns(fs_root: Path): uris = {e.uri async for e in src.discover(since=None)} assert (fs_root / "b.txt").as_uri() not in uris assert (fs_root / "a.md").as_uri() in uris + + +# --- symlink escape defenses --- + + +def test_fs_source_supports_rejects_paths_outside_root(fs_root: Path, tmp_path: Path): + """A URI for a file outside the configured root must not match — even + if it's a valid file:// URI. Otherwise the resolve_fetcher chain could + end up handing /etc/passwd to FSSource.fetch().""" + outside = tmp_path.parent / "outside.md" + outside.write_text("not yours") + src = FSSource(root=fs_root) + assert src.supports(outside.as_uri()) is False + + +@pytest.mark.asyncio +async def test_fs_source_fetch_rejects_paths_outside_root( + fs_root: Path, tmp_path: Path +): + outside = tmp_path.parent / "outside.md" + outside.write_text("not yours") + src = FSSource(root=fs_root) + with pytest.raises(UnsupportedSourceError, match="escapes FS root"): + await src.fetch(outside.as_uri()) + + +@pytest.mark.asyncio +async def test_fs_source_fetch_rejects_symlink_to_outside_file( + fs_root: Path, tmp_path: Path +): + """Symlink under root that points outside — the classic FS escape. + resolve() chases the link to the real path, which fails the root check.""" + secret = tmp_path.parent / "secret.md" + secret.write_text("sensitive") + link = fs_root / "looks_local.md" + link.symlink_to(secret) + src = FSSource(root=fs_root) + with pytest.raises(UnsupportedSourceError, match="escapes FS root"): + await src.fetch(link.as_uri()) + + +@pytest.mark.asyncio +async def test_fs_source_discover_skips_symlinks(fs_root: Path, tmp_path: Path): + """rglob (and os.walk by default) follows symlinks. We use + followlinks=False AND an explicit per-file is_symlink() filter so a + malicious symlink under root can't be discovered, can't be queued, and + can't be fetched even if its URI ends up enqueued some other way.""" + secret = tmp_path.parent / "secret.md" + secret.write_text("sensitive") + link = fs_root / "evil.md" + link.symlink_to(secret) + src = FSSource(root=fs_root, supported_extensions=[".md"]) + uris = {e.uri async for e in src.discover(since=None)} + assert link.as_uri() not in uris + # And the legitimate files in fs_root still come through. + assert (fs_root / "a.md").as_uri() in uris + + +@pytest.mark.asyncio +async def test_fs_source_discover_skips_symlinked_directories( + fs_root: Path, tmp_path: Path +): + """os.walk(followlinks=False) must NOT descend into directory symlinks — + otherwise a `ln -s /etc /docs/escape` would walk into /etc and try to + yield its contents.""" + outside_dir = tmp_path.parent / "outside_dir" + outside_dir.mkdir() + (outside_dir / "stolen.md").write_text("not yours") + (fs_root / "escape").symlink_to(outside_dir) + src = FSSource(root=fs_root, supported_extensions=[".md"]) + uris = {e.uri async for e in src.discover(since=None)} + # No URI under /escape/* should appear. + assert not any("escape" in u for u in uris)