diff --git a/taosmd/loaders/__init__.py b/taosmd/loaders/__init__.py index 11f7872f..547e9e87 100644 --- a/taosmd/loaders/__init__.py +++ b/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", ] diff --git a/taosmd/loaders/_safety.py b/taosmd/loaders/_safety.py new file mode 100644 index 00000000..2c454acb --- /dev/null +++ b/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 diff --git a/taosmd/loaders/chat_loader.py b/taosmd/loaders/chat_loader.py index fd3ac133..d244b382 100644 --- a/taosmd/loaders/chat_loader.py +++ b/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"] diff --git a/taosmd/loaders/doc_loader.py b/taosmd/loaders/doc_loader.py index aad34cc4..2b553bfe 100644 --- a/taosmd/loaders/doc_loader.py +++ b/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 = "" diff --git a/taosmd/loaders/email_loader.py b/taosmd/loaders/email_loader.py index 76e46c93..e8203ac5 100644 --- a/taosmd/loaders/email_loader.py +++ b/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. diff --git a/taosmd/loaders/interface.py b/taosmd/loaders/interface.py index fdd8d332..74b73b80 100644 --- a/taosmd/loaders/interface.py +++ b/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,7 +48,14 @@ 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 @@ -55,5 +63,19 @@ async def load(self, file_path: str | Path, **kwargs) -> Blob: 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 diff --git a/taosmd/loaders/transcript_loader.py b/taosmd/loaders/transcript_loader.py index ba28780a..0bf6ee7c 100644 --- a/taosmd/loaders/transcript_loader.py +++ b/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] = [] diff --git a/tests/test_loaders.py b/tests/test_loaders.py index b67339d2..4b81abb7 100644 --- a/tests/test_loaders.py +++ b/tests/test_loaders.py @@ -13,14 +13,17 @@ BlobType, ChatBlob, ChatLoader, + DEFAULT_MAX_BYTES, DocBlob, DocLoader, EmailBlob, EmailLoader, TranscriptBlob, TranscriptLoader, + check_size, pick_loader, register_loader, + resolve_within, REGISTRY, ) from taosmd.loaders.interface import LoaderInterface @@ -301,3 +304,136 @@ def test_pick_and_load_e2e(tmp_path): assert isinstance(blob, ChatBlob) assert blob.messages[0].content == "Hi." + + +# --------------------------------------------------------------------------- +# Opt-in safety guards (#112 size limit, #113 path containment) +# --------------------------------------------------------------------------- + + +def test_default_max_bytes_is_generous(): + # 100 MB — large enough that no normal loaded file trips it. + assert DEFAULT_MAX_BYTES == 100 * 1024 * 1024 + + +# --- check_size ------------------------------------------------------------ + + +def test_check_size_allows_small_file(tmp_path): + p = tmp_path / "small.txt" + p.write_text("tiny") + # Default generous cap and an explicit cap both pass silently. + check_size(p) + check_size(p, max_bytes=1024) + + +def test_check_size_rejects_oversized_file(tmp_path): + p = tmp_path / "big.txt" + p.write_bytes(b"x" * 2048) + with pytest.raises(ValueError, match="exceeds the loader size"): + check_size(p, max_bytes=1024) + + +def test_check_size_none_disables_check(tmp_path): + p = tmp_path / "big.txt" + p.write_bytes(b"x" * 2048) + # max_bytes=None is an explicit opt-out — no error even when large. + check_size(p, max_bytes=None) + + +# --- resolve_within -------------------------------------------------------- + + +def test_resolve_within_no_base_dir_allows_any_path(tmp_path): + # With base_dir=None (default) nothing is restricted — standalone + # use of any absolute path keeps working. + p = tmp_path / "anywhere.txt" + p.write_text("ok") + assert resolve_within(p) == p.resolve() + # A path well outside the tree resolves fine too. + assert resolve_within("/etc/hosts") == Path("/etc/hosts").resolve() + + +def test_resolve_within_allows_path_inside_base(tmp_path): + sub = tmp_path / "data" + sub.mkdir() + p = sub / "ok.txt" + p.write_text("ok") + assert resolve_within(p, base_dir=sub) == p.resolve() + # Nested deeper is fine too. + nested = sub / "a" / "b" + nested.mkdir(parents=True) + deep = nested / "deep.txt" + deep.write_text("ok") + assert resolve_within(deep, base_dir=sub) == deep.resolve() + + +def test_resolve_within_blocks_traversal_escape(tmp_path): + sub = tmp_path / "data" + sub.mkdir() + escape = sub / ".." / ".." / "etc" / "passwd" + with pytest.raises(ValueError, match="outside the allowed base"): + resolve_within(escape, base_dir=sub) + + +def test_resolve_within_blocks_absolute_escape(tmp_path): + sub = tmp_path / "data" + sub.mkdir() + with pytest.raises(ValueError, match="outside the allowed base"): + resolve_within("/etc/passwd", base_dir=sub) + + +def test_resolve_within_blocks_symlink_escape(tmp_path): + sub = tmp_path / "data" + sub.mkdir() + outside = tmp_path / "secret.txt" + outside.write_text("secret") + link = sub / "link.txt" + link.symlink_to(outside) + with pytest.raises(ValueError, match="outside the allowed base"): + resolve_within(link, base_dir=sub) + + +# --- wired into loader.load() --------------------------------------------- + + +def test_loader_rejects_oversized_file(tmp_path): + src = tmp_path / "huge.txt" + src.write_bytes(b"x" * 4096) + with pytest.raises(ValueError, match="exceeds the loader size"): + _run(DocLoader().load(src, max_bytes=1024)) + + +def test_loader_base_dir_blocks_traversal(tmp_path): + base = tmp_path / "allowed" + base.mkdir() + outside = tmp_path / "outside.txt" + outside.write_text("nope") + escape = base / ".." / "outside.txt" + with pytest.raises(ValueError, match="outside the allowed base"): + _run(DocLoader().load(escape, base_dir=base)) + + +def test_loader_base_dir_allows_normal_path(tmp_path): + base = tmp_path / "allowed" + base.mkdir() + src = base / "notes.txt" + src.write_text("# Hello\n\nbody") + blob = _run(DocLoader().load(src, base_dir=base)) + assert blob.title == "Hello" + + +def test_loader_no_base_dir_allows_any_path(tmp_path): + # Default call (no base_dir) loads a file regardless of location. + src = tmp_path / "notes.txt" + src.write_text("plain content") + blob = _run(DocLoader().load(src)) + assert blob.content == "plain content" + + +def test_loader_default_call_unchanged(tmp_path): + # A normal small file loads fine with no safety args at all. + src = tmp_path / "session.chat.json" + src.write_text(json.dumps([{"role": "user", "content": "Hi."}])) + blob = _run(ChatLoader().load(src)) + assert blob.messages[0].content == "Hi."