diff --git a/src/docgen/ai_client.py b/src/docgen/ai_client.py index 3a4cd10..22b24d2 100644 --- a/src/docgen/ai_client.py +++ b/src/docgen/ai_client.py @@ -35,6 +35,7 @@ import random import time import urllib.error +import urllib.parse import urllib.request import uuid from dataclasses import dataclass @@ -609,8 +610,17 @@ def _call() -> Any: return _openai_stt_payload(result) +def _require_http_url(url: str) -> None: + """Image fetches are http(s) only. ``file:`` and other schemes are refused.""" + scheme = urllib.parse.urlparse(url).scheme.lower() + if scheme not in ("http", "https"): + shown = scheme or "an empty scheme" + raise AIError(f"image URL must use http or https, not {shown!r}") + + def fetch_url_bytes(url: str, *, timeout: int = 120) -> bytes: """GET ``url`` and return the response body (image CDN / signed URLs).""" + _require_http_url(url) last_exc: BaseException | None = None for attempt in range(_MAX_HTTP_ATTEMPTS): req = urllib.request.Request(url, method="GET") diff --git a/tests/test_ai_client.py b/tests/test_ai_client.py index a1427a6..836d17d 100644 --- a/tests/test_ai_client.py +++ b/tests/test_ai_client.py @@ -931,6 +931,15 @@ def _urlopen(req: object, timeout: int = 120) -> _Resp: assert calls["n"] == 2 +def test_fetch_url_bytes_refuses_file_scheme() -> None: + from docgen.ai_client import fetch_url_bytes + + with patch("docgen.ai_client.urllib.request.urlopen") as urlopen: + with pytest.raises(AIError, match="http or https"): + fetch_url_bytes("file:///etc/passwd") + urlopen.assert_not_called() + + def test_chat_retries_rate_limit( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, clear_ai_env: None ) -> None: