feat(audit): opt-in loader size/path safety (#112 #113)

taosmd/loaders/__init__.py

@@ -25,6 +25,7 @@
 See ``reference_memory_systems_survey.md`` for the design rationale.
 """
 
+from ._safety import DEFAULT_MAX_BYTES, check_size, resolve_within
 from .blob import (
     Blob,
     BlobType,
@@ -59,4 +60,7 @@
     "REGISTRY",
     "pick_loader",
     "register_loader",
+    "DEFAULT_MAX_BYTES",
+    "check_size",
+    "resolve_within",
 ]

taosmd/loaders/_safety.py

@@ -0,0 +1,81 @@
+"""Opt-in safety guards shared by every loader.
+
+Two small, stdlib-only helpers that loaders call from their ``load()``
+entry point before touching a file:
+
+  * ``check_size`` — refuse files larger than ``max_bytes`` so a stray
+    multi-gigabyte file can't blow up memory on a ``f.read()``. The cap
+    is generous by default (100 MB) and configurable per call.
+
+  * ``resolve_within`` — when a caller pins a ``base_dir``, refuse paths
+    that resolve outside it (``../../etc/passwd``, an absolute escape, a
+    symlink pointing out of the tree). When ``base_dir`` is ``None`` —
+    the default — nothing is restricted and the resolved path is
+    returned unchanged, so standalone use against an arbitrary path is
+    never broken.
+
+Both are opt-in: a loader called the old way (no ``base_dir``, default
+cap) behaves exactly as before for any real file.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+# Generous default — large enough that no normal chat / transcript /
+# email / doc file trips it, small enough to stop a runaway read.
+DEFAULT_MAX_BYTES = 100 * 1024 * 1024  # 100 MB
+
+
+def check_size(path: str | Path, max_bytes: int | None = DEFAULT_MAX_BYTES) -> None:
+    """Raise ``ValueError`` if ``path`` is larger than ``max_bytes``.
+
+    ``max_bytes=None`` disables the check entirely (explicit opt-out).
+    Anything else is compared against ``os.path.getsize``. The error
+    names the path and both sizes so the caller can see the overage.
+    """
+    if max_bytes is None:
+        return
+    size = os.path.getsize(path)
+    if size > max_bytes:
+        raise ValueError(
+            f"{path} is {size} bytes, which exceeds the loader size "
+            f"limit of {max_bytes} bytes. Pass a larger max_bytes "
+            f"(or max_bytes=None) to override."
+        )
+
+
+def resolve_within(path: str | Path, base_dir: str | Path | None = None) -> Path:
+    """Resolve ``path`` and, when ``base_dir`` is set, confine it there.
+
+    ``base_dir`` is expected to be a directory that contains (directly or
+    transitively) the file at ``path``; passing a file as ``base_dir``
+    only ever matches that exact file.
+
+    With ``base_dir=None`` (the default) this is just ``Path(path)``
+    resolved — no restriction, so direct standalone use of any path is
+    unaffected.
+
+    With ``base_dir`` given, the resolved path must sit inside the
+    resolved ``base_dir`` or a ``ValueError`` is raised. Resolving first
+    means traversal (``../``), absolute escapes, and symlinks that point
+    out of the tree are all caught.
+
+    This is a containment check, not an atomic open: there is an inherent
+    TOCTOU window between resolving the path here and the caller opening
+    it, so a symlink swapped in after this returns is not caught. For the
+    local-first, single-user ingest path this guards against accidental
+    escapes, not a concurrent adversary on the same machine.
+    """
+    resolved = Path(path).resolve()
+    if base_dir is None:
+        return resolved
+
+    base = Path(base_dir).resolve()
+    if resolved != base and base not in resolved.parents:
+        raise ValueError(
+            f"{path} resolves to {resolved}, which is outside the "
+            f"allowed base path {base}."
+        )
+    return resolved

taosmd/loaders/chat_loader.py

@@ -25,6 +25,7 @@
 from pathlib import Path
 from typing import ClassVar
 
+from ._safety import DEFAULT_MAX_BYTES, check_size, resolve_within
 from .blob import ChatBlob, ChatMessage
 from .interface import LoaderInterface
 
@@ -44,9 +45,18 @@ def can_handle(cls, extension: str = "", mime_type: str = "") -> bool:
         ext = extension.lower().lstrip(".") if extension else ""
         return ext.endswith("chat.json") or ext.endswith("messages.json")
 
-    async def load(self, file_path: str | Path, **kwargs) -> ChatBlob:
+    async def load(
+        self,
+        file_path: str | Path,
+        *,
+        max_bytes: int | None = DEFAULT_MAX_BYTES,
+        base_dir: str | Path | None = None,
+        **kwargs,
+    ) -> ChatBlob:
         path = Path(file_path)
-        with open(path) as f:
+        safe_path = resolve_within(file_path, base_dir)
+        check_size(safe_path, max_bytes)
+        with open(safe_path) as f:
             data = json.load(f)
         if isinstance(data, dict) and "messages" in data:
             raw_messages = data["messages"]

taosmd/loaders/doc_loader.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 from typing import ClassVar
 
+from ._safety import DEFAULT_MAX_BYTES, check_size, resolve_within
 from .blob import DocBlob
 from .interface import LoaderInterface
 
@@ -23,9 +24,18 @@ class DocLoader(LoaderInterface):
         "text/plain", "text/markdown",
     )
 
-    async def load(self, file_path: str | Path, **kwargs) -> DocBlob:
+    async def load(
+        self,
+        file_path: str | Path,
+        *,
+        max_bytes: int | None = DEFAULT_MAX_BYTES,
+        base_dir: str | Path | None = None,
+        **kwargs,
+    ) -> DocBlob:
         path = Path(file_path)
-        with open(path, encoding="utf-8", errors="replace") as f:
+        safe_path = resolve_within(file_path, base_dir)
+        check_size(safe_path, max_bytes)
+        with open(safe_path, encoding="utf-8", errors="replace") as f:
             content = f.read()
 
         title = ""

taosmd/loaders/email_loader.py

@@ -18,6 +18,7 @@
 from pathlib import Path
 from typing import ClassVar
 
+from ._safety import DEFAULT_MAX_BYTES, check_size, resolve_within
 from .blob import EmailBlob
 from .interface import LoaderInterface
 
@@ -29,9 +30,18 @@ class EmailLoader(LoaderInterface):
         "message/rfc822", "application/mbox",
     )
 
-    async def load(self, file_path: str | Path, **kwargs) -> EmailBlob:
+    async def load(
+        self,
+        file_path: str | Path,
+        *,
+        max_bytes: int | None = DEFAULT_MAX_BYTES,
+        base_dir: str | Path | None = None,
+        **kwargs,
+    ) -> EmailBlob:
         path = Path(file_path)
-        with open(path, "rb") as f:
+        safe_path = resolve_within(file_path, base_dir)
+        check_size(safe_path, max_bytes)
+        with open(safe_path, "rb") as f:
             msg = email.message_from_binary_file(f, policy=policy.default)
 
         # body — prefer text/plain part, fall back to flat string.

taosmd/loaders/interface.py

@@ -17,6 +17,7 @@
 from pathlib import Path
 from typing import ClassVar
 
+from ._safety import DEFAULT_MAX_BYTES
 from .blob import Blob
 
 
@@ -47,13 +48,34 @@ def can_handle(cls, extension: str = "", mime_type: str = "") -> bool:
         return False
 
     @abstractmethod
-    async def load(self, file_path: str | Path, **kwargs) -> Blob:
+    async def load(
+        self,
+        file_path: str | Path,
+        *,
+        max_bytes: int | None = DEFAULT_MAX_BYTES,
+        base_dir: str | Path | None = None,
+        **kwargs,
+    ) -> Blob:
         """Read the file at ``file_path`` and return a typed ``Blob``.
 
         Implementations should set ``Blob.source_path`` to a string of
         ``file_path`` and populate ``raw_text`` opportunistically — when
         a string view is cheap to derive, it lets legacy ingest paths
         keep working. When it isn't cheap, leave it empty and rely on
         the typed fields.
+
+        Two opt-in safety guards, both with standalone-safe defaults:
+
+          ``max_bytes`` — the file is rejected (``ValueError``) if it is
+          larger than this. Defaults to a generous 100 MB; pass ``None``
+          to disable the check.
+
+          ``base_dir`` — when given, the resolved ``file_path`` must sit
+          inside it or a ``ValueError`` is raised (path-traversal /
+          symlink containment). Defaults to ``None`` (no restriction),
+          so direct use against any path keeps working.
+
+        Implementations enforce both via ``taosmd.loaders._safety``
+        before reading the file.
         """
         raise NotImplementedError

taosmd/loaders/transcript_loader.py

@@ -22,6 +22,7 @@
 from pathlib import Path
 from typing import Any, ClassVar
 
+from ._safety import DEFAULT_MAX_BYTES, check_size, resolve_within
 from .blob import TranscriptBlob, TranscriptStamp
 from .interface import LoaderInterface
 
@@ -55,9 +56,18 @@ def can_handle(cls, extension: str = "", mime_type: str = "") -> bool:
         ext = extension.lower().lstrip(".") if extension else ""
         return ext.endswith("transcript.json") or ext.endswith("whisper.json")
 
-    async def load(self, file_path: str | Path, **kwargs) -> TranscriptBlob:
+    async def load(
+        self,
+        file_path: str | Path,
+        *,
+        max_bytes: int | None = DEFAULT_MAX_BYTES,
+        base_dir: str | Path | None = None,
+        **kwargs,
+    ) -> TranscriptBlob:
         path = Path(file_path)
-        with open(path) as f:
+        safe_path = resolve_within(file_path, base_dir)
+        check_size(safe_path, max_bytes)
+        with open(safe_path) as f:
             data = json.load(f)
 
         rows: list[dict] = []