Spaces:
Sleeping
Sleeping
File size: 6,397 Bytes
6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 357ee36 6303ae6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """File utilities used across services.
These helpers harden the dossier walk against the cases the project's
privacy and robustness guarantees depend on:
* **Symlink containment.** ``Path.is_file()`` follows symlinks, so a
symlink dropped inside the chosen dossier folder (e.g.
``evidence.txt -> ~/.ssh/id_rsa``) would otherwise be read and its
content shipped to the LLM, breaking "original files never leave the
machine". :func:`iter_dossier_files` skips symlinked entries and any
real path that resolves outside the chosen folder.
* **Size / count caps.** A single huge file or an enormous folder would
otherwise be read fully into memory and hang the app. The walk is
bounded by :data:`MAX_DOSSIER_FILES` and :data:`MAX_DOSSIER_WALK_ENTRIES`
and individual oversized files are reported via
:func:`is_within_size_limit` so callers can degrade gracefully instead
of reading them.
"""
from __future__ import annotations
import os
from itertools import islice
from pathlib import Path
from typing import Iterator
from app.config import SUPPORTED_DOSSIER_EXTENSIONS
# Hard caps for the dossier walk. These bound a self-inflicted memory/time
# blow-up on a single-user local tool; they are deliberately generous.
MAX_DOSSIER_FILES = 2000
MAX_DOSSIER_FILE_BYTES = 25 * 1024 * 1024 # 25 MB per file
# Directory entries (files and folders) the walk reads in total, so a home
# directory or a project with node_modules stops the walk instead of
# hanging the app.
MAX_DOSSIER_WALK_ENTRIES = 20_000
def is_within_folder(path: Path, folder: Path) -> bool:
"""Return True if ``path`` resolves to a location inside ``folder``.
Both paths are fully resolved (following symlinks) before comparison,
so a symlink whose target escapes ``folder`` is rejected.
"""
try:
real = path.resolve()
folder_real = folder.resolve()
except OSError:
return False
if real == folder_real:
return True
return folder_real in real.parents
def is_within_size_limit(path: Path, max_bytes: int = MAX_DOSSIER_FILE_BYTES) -> bool:
"""Return True if ``path`` is at or below ``max_bytes``.
Returns False for files whose size cannot be determined so an
unreadable/odd entry is treated as "skip", not "read unbounded".
"""
try:
return path.stat().st_size <= max_bytes
except OSError:
return False
def iter_contained_files(
folder: str | Path,
*,
max_files: int = MAX_DOSSIER_FILES,
max_entries: int = MAX_DOSSIER_WALK_ENTRIES,
stats: dict | None = None,
) -> Iterator[Path]:
"""Yield regular files under ``folder`` that are safe to consider.
This is the security primitive the dossier walk is built on:
* **Symlinks are skipped** — both symlinked files and any entry whose
resolved real path falls outside ``folder`` (containment check), so
content from elsewhere on disk can never enter the dossier.
* **The walk is lazy and bounded** so an enormous folder cannot hang
the app. Folders are read one at a time, each in name order (the
order ``sorted(rglob("*"))`` gave), so files come out without the
whole tree being listed first. At most ``max_entries`` directory
entries are read, and the walk stops after ``max_files`` supported
dossier files. Callers are responsible for surfacing truncation: pass
a ``stats`` dict and ``stats["truncated"]`` says whether either limit
cut the walk short.
No extension filtering is applied here, so callers that need to report
unsupported files (e.g. the folder validator) still see them. They do
not count toward ``max_files``, so a folder full of images, code or
build output cannot crowd out the real dossier files.
Symlinked *directories* are never descended into, so a directory-symlink
loop is not reachable through this walk.
"""
folder = Path(folder).expanduser()
if stats is not None:
stats["truncated"] = False
if not folder.exists() or not folder.is_dir():
return
entries_left = max_entries
def truncate() -> None:
if stats is not None:
stats["truncated"] = True
def sorted_entries(directory: Path) -> Iterator[os.DirEntry]:
# Read no more than the entry budget allows (one extra tells us
# whether anything was left unread), then sort by name.
nonlocal entries_left
try:
with os.scandir(directory) as it:
entries = list(islice(it, entries_left + 1))
except OSError:
# An unreadable sub-folder is skipped, as rglob did.
return iter(())
if len(entries) > entries_left:
entries = entries[:entries_left]
truncate()
entries_left -= len(entries)
return iter(sorted(entries, key=lambda entry: entry.name))
yielded = 0
stack = [(folder, sorted_entries(folder))]
while stack:
directory, entries = stack[-1]
entry = next(entries, None)
if entry is None:
stack.pop()
continue
path = directory / entry.name
# Skip symlinks outright (both files and dirs) and anything that
# resolves outside the chosen folder.
if entry.is_symlink():
continue
if entry.is_dir(follow_symlinks=False):
stack.append((path, sorted_entries(path)))
continue
if not entry.is_file(follow_symlinks=False):
continue
if not is_within_folder(path, folder):
continue
if path.suffix.lower() in SUPPORTED_DOSSIER_EXTENSIONS:
if yielded >= max_files:
truncate()
return
yielded += 1
yield path
def iter_dossier_files(
folder: str | Path,
*,
max_files: int = MAX_DOSSIER_FILES,
) -> Iterator[Path]:
"""Yield contained, supported-extension dossier files under ``folder``."""
for path in iter_contained_files(folder, max_files=max_files):
if path.suffix.lower() in SUPPORTED_DOSSIER_EXTENSIONS:
yield path
def iter_supported_files(folder: str | Path):
"""Backwards-compatible alias for :func:`iter_dossier_files`.
Retained so existing imports keep working; new code should call
:func:`iter_dossier_files` directly.
"""
yield from iter_dossier_files(folder)
|