Harden FS source against symlink escape; surface pool/poller liveness in /health, additional auth tests
This commit is contained in:
parent
4aee18dcbe
commit
ef9cacf981
8 changed files with 206 additions and 15 deletions
|
|
@ -9,15 +9,27 @@ router = APIRouter()
|
||||||
@router.get("/health", response_model=HealthResponse)
|
@router.get("/health", response_model=HealthResponse)
|
||||||
async def health(state: APIState = Depends(get_state)) -> HealthResponse:
|
async def health(state: APIState = Depends(get_state)) -> HealthResponse:
|
||||||
"""Liveness signal + queue/worker overview. Unauthenticated so load
|
"""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()
|
counts = await state.job_repo.counts_by_status()
|
||||||
worker_count = (
|
worker_count = (
|
||||||
state.config.ingester.workers.worker_count if state.pool is not None else 0
|
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
|
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(
|
return HealthResponse(
|
||||||
status="ok",
|
status="degraded" if degraded else "ok",
|
||||||
queue_counts=counts,
|
queue_counts=counts,
|
||||||
worker_count=worker_count,
|
worker_count=worker_count,
|
||||||
poller_count=poller_count,
|
poller_count=poller_count,
|
||||||
|
workers_alive=workers_alive,
|
||||||
|
pollers_alive=pollers_alive,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -8,6 +8,11 @@ class HealthResponse(BaseModel):
|
||||||
queue_counts: dict[str, int]
|
queue_counts: dict[str, int]
|
||||||
worker_count: int
|
worker_count: int
|
||||||
poller_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):
|
class SourceSummary(BaseModel):
|
||||||
|
|
|
||||||
|
|
@ -87,3 +87,9 @@ class PollerManager:
|
||||||
@property
|
@property
|
||||||
def pollers(self) -> list[BasePoller]:
|
def pollers(self) -> list[BasePoller]:
|
||||||
return list(self._pollers)
|
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())
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,12 @@
|
||||||
import hashlib
|
import hashlib
|
||||||
import mimetypes
|
import mimetypes
|
||||||
|
import os
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncIterator
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from urllib.parse import unquote, urlparse
|
from urllib.parse import unquote, urlparse
|
||||||
|
|
||||||
|
from haiku.rag.client.exceptions import UnsupportedSourceError
|
||||||
from haiku.rag.ingester.sources.base import (
|
from haiku.rag.ingester.sources.base import (
|
||||||
FetchResult,
|
FetchResult,
|
||||||
RevisionSnapshot,
|
RevisionSnapshot,
|
||||||
|
|
@ -51,27 +53,38 @@ class FSSource:
|
||||||
supported_extensions=self.supported_extensions,
|
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:
|
def supports(self, uri: str) -> bool:
|
||||||
scheme = urlparse(uri).scheme
|
scheme = urlparse(uri).scheme
|
||||||
if scheme not in ("", "file"):
|
if scheme not in ("", "file"):
|
||||||
return False
|
return False
|
||||||
try:
|
return self._resolve_within_root(uri) is not None
|
||||||
_uri_to_path(uri)
|
|
||||||
except ValueError:
|
|
||||||
return False
|
|
||||||
return True
|
|
||||||
|
|
||||||
async def head(self, uri: str) -> str | None:
|
async def head(self, uri: str) -> str | None:
|
||||||
path = _uri_to_path(uri).absolute()
|
path = self._resolve_within_root(uri)
|
||||||
if not path.exists():
|
if path is None or not path.exists():
|
||||||
return None
|
return None
|
||||||
return str(path.stat().st_mtime_ns)
|
return str(path.stat().st_mtime_ns)
|
||||||
|
|
||||||
async def fetch(self, uri: str) -> FetchResult:
|
async def fetch(self, uri: str) -> FetchResult:
|
||||||
# Absolute path is needed for as_uri() and matches the old
|
path = self._resolve_within_root(uri)
|
||||||
# _create_document_from_file behavior (which keyed docs on the
|
if path is None:
|
||||||
# absolute file:// URI).
|
raise UnsupportedSourceError(f"Path escapes FS root ({self.root}): {uri}")
|
||||||
path = _uri_to_path(uri).absolute()
|
|
||||||
body = path.read_bytes()
|
body = path.read_bytes()
|
||||||
content_type, _ = mimetypes.guess_type(path.name)
|
content_type, _ = mimetypes.guess_type(path.name)
|
||||||
if content_type is None:
|
if content_type is None:
|
||||||
|
|
@ -95,7 +108,21 @@ class FSSource:
|
||||||
now = datetime.now(UTC)
|
now = datetime.now(UTC)
|
||||||
seen: set[str] = set()
|
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():
|
if not path.is_file():
|
||||||
continue
|
continue
|
||||||
if not self.filter.include_file(str(path)):
|
if not self.filter.include_file(str(path)):
|
||||||
|
|
|
||||||
|
|
@ -50,6 +50,12 @@ class WorkerPool:
|
||||||
self._workers: list[asyncio.Task] = []
|
self._workers: list[asyncio.Task] = []
|
||||||
self._reaper: asyncio.Task | None = None
|
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:
|
async def start(self) -> None:
|
||||||
if self._workers:
|
if self._workers:
|
||||||
raise RuntimeError("WorkerPool already started")
|
raise RuntimeError("WorkerPool already started")
|
||||||
|
|
|
||||||
|
|
@ -63,8 +63,13 @@ async def _apply_canonical_metadata_keys(store: Store) -> None:
|
||||||
try:
|
try:
|
||||||
meta = json.loads(raw)
|
meta = json.loads(raw)
|
||||||
except Exception:
|
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(
|
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
|
skipped += 1
|
||||||
continue
|
continue
|
||||||
|
|
|
||||||
|
|
@ -71,6 +71,31 @@ async def test_health_ok_with_counts(state, jobs):
|
||||||
assert body["queue_counts"] == {"queued": 1, "dead": 1}
|
assert body["queue_counts"] == {"queued": 1, "dead": 1}
|
||||||
assert body["worker_count"] == 0 # pool not attached in the test state
|
assert body["worker_count"] == 0 # pool not attached in the test state
|
||||||
assert body["poller_count"] == 0
|
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
|
@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
|
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 ---
|
# --- /jobs ---
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.client.exceptions import UnsupportedSourceError
|
||||||
from haiku.rag.ingester.sources.base import SourceEventKind
|
from haiku.rag.ingester.sources.base import SourceEventKind
|
||||||
from haiku.rag.ingester.sources.fs import FSSource
|
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)}
|
uris = {e.uri async for e in src.discover(since=None)}
|
||||||
assert (fs_root / "b.txt").as_uri() not in uris
|
assert (fs_root / "b.txt").as_uri() not in uris
|
||||||
assert (fs_root / "a.md").as_uri() 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)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue