diff --git a/app.py b/app.py index 5215966fb..60441baf1 100644 --- a/app.py +++ b/app.py @@ -516,28 +516,69 @@ app.mount("/static", _RevalidatingStatic(directory=STATIC_DIR), name="static") @app.get("/api/generated-image/{filename}") async def serve_generated_image(filename: str, request: Request): """Serve generated images from the data directory.""" - img_path = resolve_generated_image_path(filename) - # SECURITY: filename is the only key, so anyone who knows / guesses a - # 12-hex content hash could pull another user's image bytes. Require - # auth and verify ownership via the gallery row (when one exists). - try: - from src.auth_helpers import get_current_user - from core.database import SessionLocal as _SL, GalleryImage as _GI - _user = get_current_user(request) - if _user: - _db = _SL() + # 1. Validate filename syntax upfront, rejecting malformed requests (400) + # without probing the filesystem or database. + from src.generated_images import GENERATED_IMAGE_RE + if not isinstance(filename, str) or not GENERATED_IMAGE_RE.fullmatch(filename): + raise HTTPException(status_code=400, detail="Invalid filename") + + # 2. Ownership verification: + # If auth is enabled, verify caller identity and ownership BEFORE probing disk, + # preventing existence oracles during outages or unauthorized requests. + from src.auth_helpers import auth_disabled, effective_user + if not auth_disabled(): + _user = effective_user(request) + if not _user: + raise HTTPException(status_code=401, detail="Not authenticated") + + from src.generated_images import get_registered_image_owner + _reg_owner = get_registered_image_owner(filename) + if _reg_owner is not None: + if _reg_owner != _user: + raise HTTPException(status_code=404, detail="Image not found") + else: try: - _row = _db.query(_GI).filter(_GI.filename == filename).first() - # Generated-but-not-yet-imported images have no row → allow. - # Row exists with a different owner → 404 (don't confirm existence). - if _row is not None and _row.owner and _row.owner != _user: - raise HTTPException(status_code=404, detail="Image not found") - finally: - _db.close() - except HTTPException: - raise - except Exception as _e: - logger.warning("Image ownership verification failed for %r", filename, exc_info=_e) + from core.database import ( + SessionLocal as _SL, + GalleryImage as _GI, + Session as _DbSession, + ChatMessage as _CM, + ) + _db = _SL() + try: + _row = _db.query(_GI).filter(_GI.filename == filename).first() + if _row is not None: + if _row.owner: + if _row.owner != _user: + raise HTTPException(status_code=404, detail="Image not found") + elif _row.session_id: + _sess = _db.query(_DbSession).filter(_DbSession.id == _row.session_id).first() + if not _sess or not _sess.owner or _sess.owner != _user: + raise HTTPException(status_code=404, detail="Image not found") + else: + # Gallery record has no owner or session link -> fail closed + raise HTTPException(status_code=404, detail="Image not found") + else: + # Rowless image: check if referenced in a chat session owned by this user + _chat_ref = ( + _db.query(_CM.id) + .join(_DbSession, _CM.session_id == _DbSession.id) + .filter(_DbSession.owner == _user, _CM.meta_data.like(f"%{filename}%")) + .first() + ) + if not _chat_ref: + raise HTTPException(status_code=404, detail="Image not found") + finally: + _db.close() + except HTTPException: + raise + except Exception as _e: + logger.warning("Image ownership verification failed for %r", filename, exc_info=_e) + raise HTTPException(status_code=503, detail="Ownership verification unavailable") + + # 3. Resolve confined path and verify existence on disk + img_path = resolve_generated_image_path(filename) + ext = filename.rsplit('.', 1)[-1].lower() mime = { "png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", @@ -545,10 +586,6 @@ async def serve_generated_image(filename: str, request: Request): "mp4": "video/mp4", "mov": "video/quicktime", "webm": "video/webm", "mkv": "video/x-matroska", "m4v": "video/mp4", }.get(ext, "application/octet-stream") - # Generated-image filenames are content hashes → the bytes for a given - # filename never change. Cache them hard so the gallery doesn't - # re-download every full-size image each time it's opened. `immutable` - # tells the browser it never needs to revalidate within the max-age. return FileResponse( str(img_path), media_type=mime, diff --git a/mcp_servers/image_gen_server.py b/mcp_servers/image_gen_server.py index 6b68a27f8..697522e3d 100644 --- a/mcp_servers/image_gen_server.py +++ b/mcp_servers/image_gen_server.py @@ -139,6 +139,7 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]: # Save to gallery try: from src.database import SessionLocal, GalleryImage + owner = str(arguments.get("_odysseus_owner") or "").strip() or None db = SessionLocal() db.add(GalleryImage( id=str(uuid.uuid4()), @@ -147,9 +148,16 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]: model=model_id, size=size, quality=payload.get("quality", "medium"), + owner=owner, )) db.commit() db.close() + if owner: + try: + from src.generated_images import register_generated_image + register_generated_image(filename, owner) + except Exception: + pass except Exception: pass diff --git a/src/ai_interaction.py b/src/ai_interaction.py index 9890ed549..c8f7979ca 100644 --- a/src/ai_interaction.py +++ b/src/ai_interaction.py @@ -1372,6 +1372,9 @@ async def do_generate_image(content: str, session_id: Optional[str] = None, owne filename = f"{uuid.uuid4().hex[:12]}.png" img_path = img_dir / filename img_path.write_bytes(base64.b64decode(img.get("b64_json"))) + if owner: + from src.generated_images import register_generated_image + register_generated_image(filename, owner) image_url = f"/api/generated-image/{filename}" image_id = _save_to_gallery(filename) @@ -1392,6 +1395,9 @@ async def do_generate_image(content: str, session_id: Optional[str] = None, owne filename = f"{uuid.uuid4().hex[:12]}.png" img_path = img_dir / filename img_path.write_bytes(dl_resp.content) + if owner: + from src.generated_images import register_generated_image + register_generated_image(filename, owner) image_url = f"/api/generated-image/{filename}" image_id = _save_to_gallery(filename) else: @@ -1537,6 +1543,9 @@ async def do_edit_image( img_dir.mkdir(parents=True, exist_ok=True) filename = f"{uuid.uuid4().hex[:12]}{suffix}" (img_dir / filename).write_bytes(image_bytes) + if owner: + from src.generated_images import register_generated_image + register_generated_image(filename, owner) return f"/api/generated-image/{filename}", _save_edited_image_to_gallery(filename) async def _try_local_img2img_fallback(client: httpx.AsyncClient) -> Optional[Dict[str, Any]]: diff --git a/src/generated_images.py b/src/generated_images.py index 8c312cd99..012c3c568 100644 --- a/src/generated_images.py +++ b/src/generated_images.py @@ -1,5 +1,8 @@ import re +import threading +import time from pathlib import Path +from typing import Optional from fastapi import HTTPException @@ -12,10 +15,47 @@ GENERATED_IMAGE_RE = re.compile( r"^[a-f0-9]{8,64}\.(png|jpg|jpeg|webp|gif|mp4|mov|webm|mkv|m4v)$" ) GENERATED_IMAGE_HEADERS = { - "Cache-Control": "public, max-age=31536000, immutable", + "Cache-Control": "private, no-cache, must-revalidate", "X-Content-Type-Options": "nosniff", } +_RECENT_GENERATED_IMAGES: dict[str, tuple[str, float]] = {} +_RECENT_IMAGES_LOCK = threading.Lock() +_RECENT_IMAGES_TTL = 3600 # 1 hour TTL + + +def register_generated_image(filename: str, owner: Optional[str]) -> None: + """Record ownership of a freshly generated image. + + Ensures legitimate access to freshly generated images even before + or in case Gallery database records are created/committed. + """ + if not filename or not owner: + return + now = time.time() + with _RECENT_IMAGES_LOCK: + if len(_RECENT_GENERATED_IMAGES) > 200: + expired = [k for k, (_, exp) in _RECENT_GENERATED_IMAGES.items() if exp < now] + for k in expired: + _RECENT_GENERATED_IMAGES.pop(k, None) + _RECENT_GENERATED_IMAGES[filename] = (owner, now + _RECENT_IMAGES_TTL) + + +def get_registered_image_owner(filename: str) -> Optional[str]: + """Return the owner of a freshly generated image if still valid in the registry.""" + if not filename: + return None + now = time.time() + with _RECENT_IMAGES_LOCK: + entry = _RECENT_GENERATED_IMAGES.get(filename) + if not entry: + return None + owner, expiry = entry + if now > expiry: + _RECENT_GENERATED_IMAGES.pop(filename, None) + return None + return owner + def resolve_generated_image_path(filename: str) -> Path: if not isinstance(filename, str) or not GENERATED_IMAGE_RE.fullmatch(filename): @@ -27,3 +67,4 @@ def resolve_generated_image_path(filename: str) -> Path: if not path.exists(): raise HTTPException(status_code=404, detail="Image not found") return path + diff --git a/src/tool_execution.py b/src/tool_execution.py index a83727422..599eade99 100644 --- a/src/tool_execution.py +++ b/src/tool_execution.py @@ -1317,7 +1317,8 @@ def _parse_qualified_mcp_args(tool: str, content: str) -> tuple[Dict, Optional[s if tool.startswith("mcp__email__"): return {}, "Email MCP tool arguments must be a JSON object." return {}, None - return parsed, None + sanitized = {k: v for k, v in parsed.items() if not k.startswith("_odysseus_")} + return sanitized, None def _parse_generate_image(content: str) -> Dict: @@ -1402,9 +1403,12 @@ def _build_mcp_args(tool: str, content: str) -> Dict: except (json.JSONDecodeError, TypeError): decoded = None if isinstance(decoded, dict) and any(k in decoded for k in primaries): - return decoded + return {k: v for k, v in decoded.items() if not k.startswith("_odysseus_")} parser = _MCP_ARG_PARSERS.get(tool) - return parser(content) if parser else {} + res = parser(content) if parser else {} + if isinstance(res, dict): + return {k: v for k, v in res.items() if not k.startswith("_odysseus_")} + return res def _normalize_mcp_text_error(result: Dict) -> Dict: @@ -1421,20 +1425,25 @@ async def _call_mcp_tool( tool: str, content: str, progress_cb: Optional[Callable[[Dict], Awaitable[None]]] = None, + owner: Optional[str] = None, ) -> Dict: """Route a legacy tool call through the MCP manager, with direct fallbacks.""" bound = active_backend_operation() if bound is not None and isinstance(bound.resource, NativeBackendResource): - return await _direct_fallback(tool, content, progress_cb=progress_cb) or {"error": f"Native tool '{tool}' unavailable", "exit_code": 1} + return await _direct_fallback(tool, content, progress_cb=progress_cb, owner=owner) or {"error": f"Native tool '{tool}' unavailable", "exit_code": 1} mcp = get_mcp_manager() if not mcp: if bound is not None: raise ResourceIdentityError("Pinned MCP backend is unavailable") - return await _direct_fallback(tool, content, progress_cb=progress_cb) or {"error": f"MCP manager not available for tool '{tool}'", "exit_code": 1} + return await _direct_fallback(tool, content, progress_cb=progress_cb, owner=owner) or {"error": f"MCP manager not available for tool '{tool}'", "exit_code": 1} server_id, tool_name = _MCP_TOOL_MAP[tool] qualified = f"mcp__{server_id}__{tool_name}" args = _build_mcp_args(tool, content) + if tool == "generate_image": + args.pop("_odysseus_owner", None) + if owner: + args["_odysseus_owner"] = owner result = await mcp.call_tool(qualified, args) # Stdio MCP servers can only return TextContent, so an explicit @@ -1444,7 +1453,7 @@ async def _call_mcp_tool( # If MCP server not connected, try direct fallback if bound is None and isinstance(result, dict) and result.get("exit_code") == 1 and "not connected" in result.get("error", ""): - fallback = await _direct_fallback(tool, content, progress_cb=progress_cb) + fallback = await _direct_fallback(tool, content, progress_cb=progress_cb, owner=owner) if fallback: return fallback @@ -2193,7 +2202,7 @@ async def _execute_tool_block_impl( elif tool in _MCP_TOOL_MAP: first_line = content.split(chr(10))[0][:80] desc = f"{tool}: {first_line}" - result = await dispatched(_call_mcp_tool(tool, content, progress_cb=progress_cb)) + result = await dispatched(_call_mcp_tool(tool, content, progress_cb=progress_cb, owner=owner)) elif tool in ("grep", "glob", "ls", "get_workspace", "host_shell"): # Code-navigation tools — no MCP server; run the direct implementation. first_line = content.split(chr(10))[0][:80] @@ -2379,7 +2388,7 @@ async def _execute_tool_block_impl( ) else: if isinstance(parsed, dict): - args = parsed + args = {k: v for k, v in parsed.items() if not k.startswith("_odysseus_")} else: _args_error = ( f"'{tool}' arguments must be a JSON object, " @@ -2388,11 +2397,12 @@ async def _execute_tool_block_impl( if _args_error is not None: result = {"error": _args_error, "exit_code": 1} else: + args = dict(args) + args.pop(_EMAIL_MCP_OWNER_ARG, None) + args.pop(_EMAIL_MCP_SESSION_ARG, None) if owner: - args = dict(args) args[_EMAIL_MCP_OWNER_ARG] = owner if session_id: - args = dict(args) args[_EMAIL_MCP_SESSION_ARG] = session_id result = await dispatched(mcp.call_tool(qualified, args)) else: @@ -2406,14 +2416,20 @@ async def _execute_tool_block_impl( if parse_error: result = {"error": parse_error, "exit_code": 1} else: + args = dict(args) + args.pop("_odysseus_owner", None) + args.pop("_odysseus_session_id", None) if tool.startswith("mcp__email__"): if owner: - args = dict(args) args[_EMAIL_MCP_OWNER_ARG] = owner if session_id: - args = dict(args) args[_EMAIL_MCP_SESSION_ARG] = session_id + elif tool == "mcp__image_gen__generate_image" or tool.startswith("mcp__image_gen__"): + if owner: + args["_odysseus_owner"] = owner result = _normalize_mcp_text_error(await dispatched(mcp.call_tool(tool, args))) + if tool == "mcp__image_gen__generate_image" or tool.startswith("mcp__image_gen__"): + _promote_image_fields(result) else: desc = f"mcp: {tool}" result = {"error": "MCP manager not available", "exit_code": 1} diff --git a/tests/test_generated_image_confinement.py b/tests/test_generated_image_confinement.py index 5628706cb..aec8effff 100644 --- a/tests/test_generated_image_confinement.py +++ b/tests/test_generated_image_confinement.py @@ -60,7 +60,7 @@ def test_generated_image_headers_include_nosniff(): assert generated_images.GENERATED_IMAGE_HEADERS["X-Content-Type-Options"] == "nosniff" assert ( generated_images.GENERATED_IMAGE_HEADERS["Cache-Control"] - == "public, max-age=31536000, immutable" + == "private, no-cache, must-revalidate" ) diff --git a/tests/test_generated_image_ownership.py b/tests/test_generated_image_ownership.py new file mode 100644 index 000000000..7075366b9 --- /dev/null +++ b/tests/test_generated_image_ownership.py @@ -0,0 +1,493 @@ +import os +from unittest.mock import MagicMock, patch +import pytest +from fastapi import HTTPException +from starlette.requests import Request + +from app import serve_generated_image +from src.generated_images import ( + GENERATED_IMAGE_HEADERS, + register_generated_image, +) + + +class DummyState: + def __init__(self, current_user=None, api_token=False, api_token_owner=None): + self.current_user = current_user + self.api_token = api_token + self.api_token_owner = api_token_owner + + +def _make_request(current_user=None, api_token=False, api_token_owner=None): + req = MagicMock(spec=Request) + req.state = DummyState(current_user, api_token, api_token_owner) + return req + + +@pytest.fixture +def sample_image(tmp_path, monkeypatch): + from src import generated_images + image_dir = tmp_path / "generated_images" + image_dir.mkdir(parents=True, exist_ok=True) + filename = "0123456789ab.png" + img_file = image_dir / filename + img_file.write_bytes(b"\x89PNG\r\n\x1a\nfake-image-bytes") + monkeypatch.setattr(generated_images, "GENERATED_IMAGE_DIR", image_dir) + return filename + + +@pytest.mark.asyncio +async def test_authorized_cookie_session_owner(sample_image): + """Cookie session owner receives 200 for their owned image.""" + req = _make_request(current_user="alice") + mock_row = MagicMock() + mock_row.owner = "alice" + mock_row.filename = sample_image + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_row + mock_sl.return_value = mock_db + + resp = await serve_generated_image(sample_image, req) + assert resp.status_code == 200 + assert resp.headers["Cache-Control"] == "private, no-cache, must-revalidate" + assert resp.headers["X-Content-Type-Options"] == "nosniff" + + +@pytest.mark.asyncio +async def test_authorized_api_token_owner(sample_image): + """API token request matching image owner receives 200.""" + req = _make_request(current_user="api", api_token=True, api_token_owner="alice") + mock_row = MagicMock() + mock_row.owner = "alice" + mock_row.filename = sample_image + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_row + mock_sl.return_value = mock_db + + resp = await serve_generated_image(sample_image, req) + assert resp.status_code == 200 + + +@pytest.mark.asyncio +async def test_cross_user_denial_returns_404(sample_image): + """Bob attempting to access Alice's image receives 404 (preventing existence leak).""" + req = _make_request(current_user="bob") + mock_row = MagicMock() + mock_row.owner = "alice" + mock_row.filename = sample_image + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_row + mock_sl.return_value = mock_db + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req) + assert exc_info.value.status_code == 404 + assert exc_info.value.detail == "Image not found" + + +@pytest.mark.asyncio +async def test_cross_user_api_token_denial_returns_404(sample_image): + """Bob's API token attempting to access Alice's image receives 404.""" + req = _make_request(current_user="api", api_token=True, api_token_owner="bob") + mock_row = MagicMock() + mock_row.owner = "alice" + mock_row.filename = sample_image + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_row + mock_sl.return_value = mock_db + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_unauthenticated_request_rejected(sample_image): + """Unauthenticated caller (no user identity) is rejected with 401 when auth is enabled.""" + req = _make_request(current_user=None) + + with patch("src.auth_helpers.auth_disabled", return_value=False): + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail == "Not authenticated" + + +@pytest.mark.asyncio +async def test_missing_gallery_row_fails_closed_404(sample_image): + """Image on disk with no gallery row, no registry, and no chat reference returns 404.""" + req = _make_request(current_user="bob") + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + # No gallery row + mock_db.query.return_value.filter.return_value.first.return_value = None + # No chat message reference + mock_db.query.return_value.join.return_value.filter.return_value.first.return_value = None + mock_sl.return_value = mock_db + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req) + assert exc_info.value.status_code == 404 + assert exc_info.value.detail == "Image not found" + + +@pytest.mark.asyncio +async def test_null_owner_metadata_fails_closed_404(sample_image): + """Gallery row with owner=None and no session link returns 404.""" + req = _make_request(current_user="bob") + mock_row = MagicMock() + mock_row.owner = None + mock_row.session_id = None + mock_row.filename = sample_image + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.return_value = mock_row + mock_sl.return_value = mock_db + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_session_linked_null_owner_authorized(sample_image): + """Gallery row with owner=None resolves owner through session_id.""" + req_alice = _make_request(current_user="alice") + req_bob = _make_request(current_user="bob") + + mock_row = MagicMock() + mock_row.owner = None + mock_row.session_id = "sess-123" + mock_row.filename = sample_image + + mock_sess = MagicMock() + mock_sess.owner = "alice" + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.return_value.filter.return_value.first.side_effect = [ + mock_row, mock_sess, # for Alice + mock_row, mock_sess, # for Bob + ] + mock_sl.return_value = mock_db + + # Alice is the session owner -> 200 + resp = await serve_generated_image(sample_image, req_alice) + assert resp.status_code == 200 + + # Bob is not the session owner -> 404 + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req_bob) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_chat_referenced_rowless_image_authorized(sample_image): + """Rowless image referenced in a chat session owned by user is authorized.""" + req_alice = _make_request(current_user="alice") + req_bob = _make_request(current_user="bob") + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + # Row is None + mock_db.query.return_value.filter.return_value.first.return_value = None + # Alice's chat query returns a match, Bob's does not + mock_db.query.return_value.join.return_value.filter.return_value.first.side_effect = [ + MagicMock(), # Alice has a match + None, # Bob has no match + ] + mock_sl.return_value = mock_db + + resp = await serve_generated_image(sample_image, req_alice) + assert resp.status_code == 200 + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req_bob) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_freshly_generated_registry_preserves_access(sample_image): + """Freshly generated image registered to Alice is accessible before DB commit; Bob is denied.""" + req_alice = _make_request(current_user="alice") + req_bob = _make_request(current_user="bob") + + fresh_filename = "fedcba987654.png" + from src import generated_images + img_file = generated_images.GENERATED_IMAGE_DIR / fresh_filename + img_file.write_bytes(b"\x89PNG\r\n\x1a\nfake-image-bytes") + + register_generated_image(fresh_filename, "alice") + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_sl.side_effect = AssertionError("DB should not be reached for registered owner") + + # Alice retrieves without hitting DB + resp = await serve_generated_image(fresh_filename, req_alice) + assert resp.status_code == 200 + + # Bob is denied with 404 + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(fresh_filename, req_bob) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_database_failure_fails_closed_503(sample_image): + """Database exception during ownership check raises HTTP 503.""" + req = _make_request(current_user="alice") + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.side_effect = RuntimeError("SQLite database locked or corrupted") + mock_sl.return_value = mock_db + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image(sample_image, req) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == "Ownership verification unavailable" + + +@pytest.mark.asyncio +async def test_no_login_compatibility_auth_disabled(sample_image): + """When AUTH_ENABLED=false, existing images are served with 200, missing images return 404.""" + req = _make_request(current_user=None) + + with patch("src.auth_helpers.auth_disabled", return_value=True): + resp = await serve_generated_image(sample_image, req) + assert resp.status_code == 200 + + with pytest.raises(HTTPException) as exc_info: + await serve_generated_image("aaaaaaaaaaaa.png", req) + assert exc_info.value.status_code == 404 + + +def test_cache_headers_private_no_cache(): + """Verify generated image response headers use private, no-cache directives.""" + assert GENERATED_IMAGE_HEADERS["Cache-Control"] == "private, no-cache, must-revalidate" + assert GENERATED_IMAGE_HEADERS["X-Content-Type-Options"] == "nosniff" + + +@pytest.mark.asyncio +async def test_existence_oracle_eliminated_during_outage(sample_image): + """During a DB outage, both existing and non-existent files return 503 without revealing existence.""" + req = _make_request(current_user="alice") + + with patch("src.auth_helpers.auth_disabled", return_value=False), \ + patch("core.database.SessionLocal") as mock_sl: + mock_db = MagicMock() + mock_db.query.side_effect = RuntimeError("DB down") + mock_sl.return_value = mock_db + + # Existing file -> 503 + with pytest.raises(HTTPException) as exc_existing: + await serve_generated_image(sample_image, req) + assert exc_existing.value.status_code == 503 + + # Non-existent file -> 503 (NOT 404) + with pytest.raises(HTTPException) as exc_missing: + await serve_generated_image("111122223333.png", req) + assert exc_missing.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_mcp_image_generation_passes_owner_arg(): + """_call_mcp_tool passes _odysseus_owner to MCP arguments for generate_image.""" + from src import tool_execution + + mock_mcp = MagicMock() + mock_mcp.call_tool = MagicMock() + + async def fake_call(qualified, args): + return {"stdout": "Generated image for: cat\nhttps://host/api/generated-image/1234567890ab.png", "exit_code": 0} + + mock_mcp.call_tool.side_effect = fake_call + + with patch.object(tool_execution, "get_mcp_manager", return_value=mock_mcp), \ + patch.object(tool_execution, "active_backend_operation", return_value=None): + res = await tool_execution._call_mcp_tool("generate_image", "cat", owner="alice") + assert res.get("image_url") == "https://host/api/generated-image/1234567890ab.png" + call_args = mock_mcp.call_tool.call_args[0] + assert call_args[1].get("_odysseus_owner") == "alice" + + +@pytest.mark.asyncio +async def test_mcp_fallback_propagates_owner(): + """When MCP manager is unavailable, _direct_fallback receives owner.""" + from src import tool_execution + + with patch.object(tool_execution, "get_mcp_manager", return_value=None), \ + patch.object(tool_execution, "active_backend_operation", return_value=None), \ + patch.object(tool_execution, "_direct_fallback", return_value={"output": "ok", "exit_code": 0}) as mock_fallback: + res = await tool_execution._call_mcp_tool("web_search", "query", owner="alice") + assert res == {"output": "ok", "exit_code": 0} + mock_fallback.assert_called_once_with("web_search", "query", progress_cb=None, owner="alice") + + +@pytest.mark.asyncio +async def test_trusted_owner_overrides_forged_owner_in_tool_args(): + """Trusted owner overrides any forged _odysseus_owner in tool block content.""" + import json + from src import tool_execution + + mock_mcp = MagicMock() + mock_mcp.call_tool = MagicMock() + + async def fake_call(qualified, args): + return {"stdout": "Generated image for: cat\n/api/generated-image/1234567890ab.png", "exit_code": 0} + + mock_mcp.call_tool.side_effect = fake_call + forged_content = json.dumps({"prompt": "cat", "_odysseus_owner": "bob"}) + + with patch.object(tool_execution, "get_mcp_manager", return_value=mock_mcp), \ + patch.object(tool_execution, "active_backend_operation", return_value=None): + await tool_execution._call_mcp_tool("generate_image", forged_content, owner="alice") + call_args = mock_mcp.call_tool.call_args[0] + assert call_args[1].get("_odysseus_owner") == "alice" + + +@pytest.mark.asyncio +async def test_missing_or_empty_owner_removes_forged_owner_claim(): + """When caller has no trusted owner (None or empty), forged _odysseus_owner is completely removed.""" + import json + from src import tool_execution + + mock_mcp = MagicMock() + mock_mcp.call_tool = MagicMock() + + async def fake_call(qualified, args): + return {"stdout": "Generated image for: cat\n/api/generated-image/1234567890ab.png", "exit_code": 0} + + mock_mcp.call_tool.side_effect = fake_call + forged_content = json.dumps({"prompt": "cat", "_odysseus_owner": "bob"}) + + with patch.object(tool_execution, "get_mcp_manager", return_value=mock_mcp), \ + patch.object(tool_execution, "active_backend_operation", return_value=None): + # Case 1: owner=None + await tool_execution._call_mcp_tool("generate_image", forged_content, owner=None) + call_args = mock_mcp.call_tool.call_args[0] + assert "_odysseus_owner" not in call_args[1] + + # Case 2: owner="" + await tool_execution._call_mcp_tool("generate_image", forged_content, owner="") + call_args = mock_mcp.call_tool.call_args[0] + assert "_odysseus_owner" not in call_args[1] + + +@pytest.mark.asyncio +async def test_qualified_mcp_image_gen_sanitizes_forged_owner(): + """Qualified mcp__image_gen__generate_image tool dispatches with sanitized owner.""" + import json + from unittest.mock import AsyncMock + from src import tool_execution + from src.tool_types import ToolBlock + + mock_mcp = MagicMock() + mock_mcp.call_tool = AsyncMock(return_value={"stdout": "ok", "exit_code": 0}) + forged_content = json.dumps({"prompt": "cat", "_odysseus_owner": "bob"}) + + with patch.object(tool_execution, "get_mcp_manager", return_value=mock_mcp), \ + patch.object(tool_execution, "_owner_is_admin", return_value=True): + # With trusted owner 'alice' -> sets 'alice' + await tool_execution._execute_tool_block_impl( + ToolBlock("mcp__image_gen__generate_image", forged_content), + owner="alice", + ) + call_args = mock_mcp.call_tool.call_args[0] + assert call_args[1].get("_odysseus_owner") == "alice" + + # Without owner -> removes forged owner + await tool_execution._execute_tool_block_impl( + ToolBlock("mcp__image_gen__generate_image", forged_content), + owner=None, + ) + call_args = mock_mcp.call_tool.call_args[0] + assert "_odysseus_owner" not in call_args[1] + + +@pytest.mark.asyncio +async def test_mcp_disconnected_fallback_does_not_resurrect_forged_owner(): + """When MCP is disconnected, direct fallback receives verified owner context, not forged arguments.""" + import json + from src import tool_execution + + mock_mcp = MagicMock() + mock_mcp.call_tool = MagicMock() + + async def fake_call(qualified, args): + return {"error": "Server not connected", "exit_code": 1} + + mock_mcp.call_tool.side_effect = fake_call + forged_content = json.dumps({"query": "python", "_odysseus_owner": "bob"}) + + with patch.object(tool_execution, "get_mcp_manager", return_value=mock_mcp), \ + patch.object(tool_execution, "active_backend_operation", return_value=None), \ + patch.object(tool_execution, "_direct_fallback", return_value={"output": "ok", "exit_code": 0}) as mock_fallback: + # Trusted owner=None: mock_fallback receives owner=None + await tool_execution._call_mcp_tool("web_search", forged_content, owner=None) + mock_fallback.assert_called_with("web_search", forged_content, progress_cb=None, owner=None) + + # Trusted owner="alice": mock_fallback receives owner="alice" + await tool_execution._call_mcp_tool("web_search", forged_content, owner="alice") + mock_fallback.assert_called_with("web_search", forged_content, progress_cb=None, owner="alice") + + +@pytest.mark.asyncio +async def test_image_gen_server_owner_sanitization(): + """mcp_servers.image_gen_server sanitizes whitespace, empty, and None owner claims.""" + from unittest.mock import AsyncMock + from mcp_servers import image_gen_server + + with patch("src.database.SessionLocal") as mock_sl, \ + patch("httpx.AsyncClient") as mock_http, \ + patch("src.ai_interaction._resolve_model", return_value=("https://host/v1/chat/completions", "gpt-image-1", {})), \ + patch("src.settings.get_setting", return_value="dummy_key"), \ + patch("src.settings.load_settings", return_value={"image_model": "gpt-image-1"}), \ + patch("pathlib.Path.write_bytes"): + mock_db = MagicMock() + mock_sl.return_value = mock_db + + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.json.return_value = { + "data": [{"b64_json": "aGVsbG8="}] + } + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_resp) + mock_http.return_value.__aenter__.return_value = mock_client + + # Call with whitespace owner: " " -> owner becomes None + await image_gen_server.call_tool( + "generate_image", + {"prompt": "cat", "_odysseus_owner": " "}, + ) + added = mock_db.add.call_args[0][0] + assert added.owner is None + + # Call with valid owner: "alice" -> owner becomes "alice" + await image_gen_server.call_tool( + "generate_image", + {"prompt": "cat", "_odysseus_owner": "alice"}, + ) + added = mock_db.add.call_args[0][0] + assert added.owner == "alice"