Skip to content

Commit dcfe920

Browse files
committed
Fix API local media path containment
Reuse the media route's canonical containment check for preview and message inputs. Validate original and converted media, reject incomplete media references, and keep URLs, uploads, and converter outputs unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 59fd033b-89ce-4424-b418-510b8144b8c8
1 parent 3f729ac commit dcfe920

7 files changed

Lines changed: 363 additions & 113 deletions

File tree

‎pyrit/backend/routes/media.py‎

Lines changed: 8 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -16,22 +16,21 @@
1616
opaque bytes.
1717
"""
1818

19+
import asyncio
1920
import logging
2021
import mimetypes
2122
from pathlib import Path
2223

2324
from fastapi import APIRouter, HTTPException, Query
2425
from fastapi.responses import FileResponse
2526

27+
from pyrit.backend.services.media_persistence import validate_media_path
2628
from pyrit.memory import CentralMemory
2729

2830
logger = logging.getLogger(__name__)
2931

3032
router = APIRouter()
3133

32-
# Only serve files from known media subdirectories under results_path.
33-
_ALLOWED_SUBDIRECTORIES = {"prompt-memory-entries", "seed-prompt-entries"}
34-
3534
# Only these known-safe media types render inline. Every other extension is
3635
# served as an application/octet-stream attachment.
3736
_INLINE_EXTENSIONS = {
@@ -60,41 +59,6 @@
6059
}
6160

6261

63-
def _validate_media_path(*, path: str, allowed_root: Path) -> Path:
64-
"""
65-
Validate and sanitize a user-provided file path against an allowed root directory.
66-
67-
Uses ``Path.resolve()`` to resolve symlinks and ``..`` components, then
68-
verifies the canonical path is under the allowed root. This is the standard
69-
sanitization pattern recognized by static analysis tools (e.g. CodeQL
70-
``py/path-injection``).
71-
72-
Args:
73-
path: The user-provided file path to validate.
74-
allowed_root: The canonical (``resolve``-d) allowed root directory.
75-
76-
Returns:
77-
The canonical, validated file path.
78-
79-
Raises:
80-
HTTPException 403: If the path fails any validation check.
81-
"""
82-
real_path = Path(path).resolve(strict=False)
83-
84-
try:
85-
relative_parts = real_path.relative_to(allowed_root).parts
86-
except ValueError as exc:
87-
raise HTTPException(
88-
status_code=403, detail="Access denied: path is outside the allowed results directory."
89-
) from exc
90-
91-
# Restrict to known media subdirectories (e.g. prompt-memory-entries/)
92-
if not relative_parts or relative_parts[0] not in _ALLOWED_SUBDIRECTORIES:
93-
raise HTTPException(status_code=403, detail="Access denied: path is not in a media subdirectory.")
94-
95-
return real_path
96-
97-
9862
@router.get("/media")
9963
async def serve_media_async(
10064
path: str = Query(..., description="Absolute path to the local media file to serve."),
@@ -126,13 +90,16 @@ async def serve_media_async(
12690
memory = CentralMemory.get_memory_instance()
12791
if not memory.results_path:
12892
raise HTTPException(status_code=500, detail="Memory results_path is not configured.")
129-
allowed_root = Path(memory.results_path).resolve(strict=False)
93+
allowed_root = await asyncio.to_thread(Path(memory.results_path).resolve, strict=False)
13094
except Exception as exc:
13195
raise HTTPException(status_code=500, detail="Memory not initialized; cannot determine results path.") from exc
13296

133-
validated_path = _validate_media_path(path=path, allowed_root=allowed_root)
97+
try:
98+
validated_path = await asyncio.to_thread(validate_media_path, path=path, allowed_root=allowed_root)
99+
except ValueError as exc:
100+
raise HTTPException(status_code=403, detail=str(exc)) from exc
134101

135-
if not validated_path.is_file():
102+
if not await asyncio.to_thread(validated_path.is_file):
136103
raise HTTPException(status_code=404, detail="File not found.")
137104

138105
extension = validated_path.suffix.lower()

‎pyrit/backend/services/media_persistence.py‎

Lines changed: 55 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
from urllib.parse import parse_qs, urlparse
1818

1919
from pyrit.backend.models import DEFAULT_MEDIA_EXTENSIONS
20-
from pyrit.memory import data_serializer_factory
20+
from pyrit.memory import CentralMemory, data_serializer_factory
2121
from pyrit.models import MEDIA_PATH_DATA_TYPES
2222

2323
if TYPE_CHECKING:
@@ -50,6 +50,36 @@ class MediaPersistenceResult:
5050
SerializerFactory = Callable[..., Any]
5151

5252

53+
def validate_media_path(*, path: str, allowed_root: Path) -> Path:
54+
"""
55+
Resolve symlinks and parent components before checking results-directory containment.
56+
57+
Returns:
58+
The canonical path in an allowed media subdirectory.
59+
60+
Raises:
61+
ValueError: If the path is outside the allowed media directories.
62+
"""
63+
real_path = Path(path).resolve(strict=False)
64+
try:
65+
relative_parts = real_path.relative_to(allowed_root.resolve(strict=False)).parts
66+
except ValueError as exc:
67+
raise ValueError("Access denied: path is outside the allowed results directory.") from exc
68+
69+
if not relative_parts or relative_parts[0] not in {"prompt-memory-entries", "seed-prompt-entries"}:
70+
raise ValueError("Access denied: path is not in a media subdirectory.")
71+
72+
return real_path
73+
74+
75+
async def _validate_local_media_path_async(path: str) -> str:
76+
memory = CentralMemory.get_memory_instance()
77+
if not memory.results_path:
78+
raise ValueError("Memory results_path is not configured.")
79+
validated_path = await asyncio.to_thread(validate_media_path, path=path, allowed_root=Path(memory.results_path))
80+
return str(validated_path)
81+
82+
5383
def _is_raw_base64(value: str) -> bool:
5484
"""Return whether *value* is syntactically valid raw base64."""
5585
try:
@@ -102,6 +132,9 @@ async def persist_media_value_async(
102132
attack ingestion and converter preview while keeping origin detection,
103133
extension resolution, and persistence in one component.
104134
135+
Local paths and ``/api/media`` references must be in a media subdirectory
136+
of the configured results directory, as required by ``GET /api/media``.
137+
105138
Returns:
106139
A typed result containing the resolved value and persistence metadata.
107140
"""
@@ -116,14 +149,17 @@ async def persist_media_value_async(
116149

117150
if value.startswith("/api/media"):
118151
parsed = urlparse(value)
119-
file_path = parse_qs(parsed.query).get("path", [None])[0]
120-
return MediaPersistenceResult(
121-
value=file_path or value,
122-
origin=MediaOrigin.MEDIA_REFERENCE,
123-
persisted=False,
124-
resolved=file_path is not None,
125-
mime_type=mime_type,
126-
)
152+
if parsed.path == "/api/media":
153+
file_path = parse_qs(parsed.query).get("path", [None])[0]
154+
if file_path is None:
155+
raise ValueError("Media reference must include a path.")
156+
return MediaPersistenceResult(
157+
value=await _validate_local_media_path_async(file_path),
158+
origin=MediaOrigin.MEDIA_REFERENCE,
159+
persisted=False,
160+
resolved=True,
161+
mime_type=mime_type,
162+
)
127163

128164
data_uri_mime_type: str | None = None
129165
payload = value
@@ -133,17 +169,19 @@ async def persist_media_value_async(
133169
origin = MediaOrigin.DATA_URI
134170
else:
135171
try:
136-
if await asyncio.to_thread(Path(value).is_file):
137-
return MediaPersistenceResult(
138-
value=value,
139-
origin=MediaOrigin.LOCAL_PATH,
140-
persisted=False,
141-
resolved=True,
142-
mime_type=mime_type,
143-
)
172+
is_file = await asyncio.to_thread(Path(value).is_file)
144173
except (OSError, ValueError):
145174
if require_valid_base64_after_path_error and not _is_raw_base64(value):
146175
raise
176+
is_file = False
177+
if is_file:
178+
return MediaPersistenceResult(
179+
value=await _validate_local_media_path_async(value),
180+
origin=MediaOrigin.LOCAL_PATH,
181+
persisted=False,
182+
resolved=True,
183+
mime_type=mime_type,
184+
)
147185

148186
extension = _resolve_extension(
149187
data_type=data_type,

‎tests/unit/backend/conftest.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,8 @@
44
"""Backend compatibility fixtures independent of a packaged workspace stamp."""
55

66
from collections.abc import Iterator
7-
from unittest.mock import patch
7+
from pathlib import Path
8+
from unittest.mock import MagicMock, patch
89

910
import pytest
1011

@@ -13,6 +14,17 @@
1314
from pyrit.backend.services.attack_service import get_attack_service
1415
from pyrit.backend.services.manual_send_scheduler import get_manual_send_scheduler
1516
from pyrit.backend.services.message_send_service import get_message_send_service
17+
from pyrit.memory import SQLiteMemory
18+
19+
20+
@pytest.fixture
21+
def managed_media_path(*, sqlite_instance: SQLiteMemory, patch_central_database: MagicMock) -> Path:
22+
"""Create a file in the isolated memory's media directory."""
23+
assert sqlite_instance.results_path is not None
24+
path = Path(sqlite_instance.results_path) / "prompt-memory-entries" / "image.png"
25+
path.parent.mkdir(parents=True, exist_ok=True)
26+
path.write_bytes(b"image")
27+
return path
1628

1729

1830
@pytest.fixture(autouse=True)

‎tests/unit/backend/test_attack_service.py‎

Lines changed: 41 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from pathlib import Path
1616
from typing import Any
1717
from unittest.mock import AsyncMock, MagicMock, patch
18+
from urllib.parse import urlencode
1819

1920
import pytest
2021
from sqlalchemy import select
@@ -813,6 +814,33 @@ async def test_get_conversation_messages_raises_for_unrelated_conversation(
813814
class TestCreateAttack:
814815
"""Tests for create_attack method."""
815816

817+
@pytest.mark.parametrize("reference", [False, True])
818+
@pytest.mark.parametrize("field", ["original_value", "converted_value"])
819+
async def test_rejects_unmanaged_prepended_media_async(
820+
self, *, managed_media_path: Path, tmp_path: Path, reference: bool, field: str
821+
) -> None:
822+
path = tmp_path / "outside.png"
823+
path.write_bytes(b"image")
824+
value = f"/api/media?{urlencode({'path': str(path)})}" if reference else str(path)
825+
piece = MessagePieceRequest(
826+
data_type="image_path",
827+
original_value=value if field == "original_value" else str(managed_media_path),
828+
converted_value=value if field == "converted_value" else str(managed_media_path),
829+
)
830+
service = AttackService()
831+
with (
832+
patch.object(service, "_get_save_target_async", new=AsyncMock(return_value=None)),
833+
patch.object(service._memory, "add_conversation_branches_to_attack_async", new_callable=AsyncMock) as store,
834+
pytest.raises(ValueError, match="outside the allowed results directory"),
835+
):
836+
await service.create_attack_async(
837+
request=CreateAttackRequest(
838+
prepended_conversation=[PrependedMessageRequest(role="user", pieces=[piece])]
839+
)
840+
)
841+
842+
store.assert_not_awaited()
843+
816844
@pytest.mark.parametrize("copy_history", [False, True])
817845
async def test_manual_creation_registers_ownership_async(
818846
self, *, sqlite_instance: SQLiteMemory, copy_history: bool
@@ -2110,17 +2138,14 @@ async def test_persists_original_and_converted_media_independently_async(
21102138
for serializer in serializers:
21112139
serializer.save_b64_image_async.assert_awaited_once()
21122140

2113-
@pytest.mark.parametrize(
2114-
("converted_value", "expected_value"),
2115-
[
2116-
("/api/media?path=preview.png", "preview.png"),
2117-
("https://example.com/preview.png?token=example", "https://example.com/preview.png?token=example"),
2118-
("preview.png", "preview.png"),
2119-
],
2120-
)
2141+
@pytest.mark.parametrize("origin", ["reference", "url", "local"])
21212142
async def test_converted_media_references_are_not_repersisted_async(
2122-
self, *, converted_value: str, expected_value: str
2143+
self, *, origin: str, managed_media_path: Path
21232144
) -> None:
2145+
expected_value = "https://example.com/preview.png?token=example" if origin == "url" else str(managed_media_path)
2146+
converted_value = (
2147+
f"/api/media?{urlencode({'path': expected_value})}" if origin == "reference" else expected_value
2148+
)
21242149
request = AddMessageRequest(
21252150
pieces=[
21262151
MessagePieceRequest(
@@ -2132,10 +2157,7 @@ async def test_converted_media_references_are_not_repersisted_async(
21322157
send=False,
21332158
target_conversation_id="test-id",
21342159
)
2135-
with (
2136-
patch("pyrit.backend.services.media_persistence.Path.is_file", return_value=True),
2137-
patch("pyrit.backend.services.attack_service.data_serializer_factory") as factory,
2138-
):
2160+
with patch("pyrit.backend.services.attack_service.data_serializer_factory") as factory:
21392161
await AttackService._persist_base64_pieces_async(pieces=request.pieces)
21402162

21412163
assert request.pieces[0].original_value == "source"
@@ -2401,14 +2423,14 @@ async def test_http_url_is_kept_as_is(self, attack_service) -> None:
24012423
assert request.pieces[0].original_value == ("https://myblob.blob.core.windows.net/images/photo.png?sv=2024")
24022424
assert request.pieces[0].converted_value == request.pieces[0].original_value
24032425

2404-
async def test_media_reference_is_resolved_without_persistence(self, attack_service) -> None:
2426+
async def test_media_reference_is_resolved_without_persistence(self, *, managed_media_path: Path) -> None:
24052427
"""Local media URLs are converted back to their decoded file paths."""
24062428
request = AddMessageRequest(
24072429
role="user",
24082430
pieces=[
24092431
MessagePieceRequest(
24102432
data_type="image_path",
2411-
original_value="/api/media?path=%2Ftmp%2Fimage.png",
2433+
original_value=f"/api/media?{urlencode({'path': str(managed_media_path)})}",
24122434
),
24132435
],
24142436
send=False,
@@ -2418,14 +2440,13 @@ async def test_media_reference_is_resolved_without_persistence(self, attack_serv
24182440
with patch("pyrit.backend.services.attack_service.data_serializer_factory") as factory:
24192441
await AttackService._persist_base64_pieces_async(pieces=request.pieces)
24202442

2421-
assert request.pieces[0].original_value == "/tmp/image.png"
2422-
assert request.pieces[0].converted_value == "/tmp/image.png"
2443+
assert request.pieces[0].original_value == str(managed_media_path)
2444+
assert request.pieces[0].converted_value == str(managed_media_path)
24232445
factory.assert_not_called()
24242446

2425-
async def test_existing_file_is_kept_without_persistence(self, attack_service, tmp_path: Path) -> None:
2447+
async def test_existing_file_is_kept_without_persistence(self, *, managed_media_path: Path) -> None:
24262448
"""An existing path remains the canonical original and converted value."""
2427-
media_path = tmp_path / "image.png"
2428-
media_path.write_bytes(b"image")
2449+
media_path = managed_media_path
24292450
request = AddMessageRequest(
24302451
role="user",
24312452
pieces=[MessagePieceRequest(data_type="image_path", original_value=str(media_path))],

0 commit comments

Comments
 (0)