fix(download): require trust for injected sessions

docs/src/content/docs/development/Architecture/model-manager.mdx

@@ -736,6 +736,7 @@ The `DownloadQueueService` constructor takes the following arguments:
 | `event_handlers` | List[DownloadEventHandler] | [] | Event handlers |
 | `max_parallel_dl` | int | 5 | Maximum number of simultaneous downloads allowed |
 | `requests_session` | requests.sessions.Session | None | An alternative requests Session object to use for the download |
+| `requests_session_is_trusted` | bool | False | Explicitly trust a caller-supplied Session that bypasses the SSRF socket guard |
 | `quiet` | bool | False | Do work quietly without issuing log messages |
 
 A typical initialization sequence will look like:
@@ -755,7 +756,7 @@ Event handlers can be provided to the queue at initialization time as shown in t
 
 `max_parallel_dl` sets the number of simultaneous active downloads that are allowed. The default of five has not been benchmarked in any way, but seems to give acceptable performance.
 
-`requests_session` can be used to provide a `requests` module Session object that will be used to stream remote URLs to disk. This facility was added for use in the module's unit tests to simulate a remote web server, but may be useful in other contexts.
+`requests_session` can be used to provide a `requests` module Session object that will be used to stream remote URLs to disk. This facility was added for use in the module's unit tests to simulate a remote web server, but may be useful in other contexts. A caller-supplied Session does not have Invoke's socket-level SSRF guard, so it is rejected while the private-address policy is enabled unless `requests_session_is_trusted=True` is also passed. Only opt in for a Session whose destination policy you trust.
 
 `quiet` will prevent the queue from issuing any log messages at the INFO or higher levels.
 

invokeai/app/services/download/download_default.py

@@ -51,13 +51,15 @@ def __init__(
         app_config: Optional[InvokeAIAppConfig] = None,
         event_bus: Optional["EventServiceBase"] = None,
         requests_session: Optional[requests.sessions.Session] = None,
+        requests_session_is_trusted: bool = False,
     ):
         """
         Initialize DownloadQueue.
 
         :param app_config: InvokeAIAppConfig object
         :param max_parallel_dl: Number of simultaneous downloads allowed [5].
         :param requests_session: Optional requests.sessions.Session object, for unit tests.
+        :param requests_session_is_trusted: Accept a caller-supplied session without the SSRF socket guard.
         """
         self._app_config = app_config or get_config()
         self._jobs: Dict[int, DownloadJob] = {}
@@ -73,11 +75,18 @@ def __init__(
         self._logger = InvokeAILogger.get_logger("DownloadQueueService")
         self._event_bus = event_bus
         # A caller-supplied session is left exactly as given (the tests inject mock
-        # transports). Sessions we build ourselves refuse to connect to a non-public
-        # address, which is the check that holds against DNS rebinding and against host
-        # spellings that `requests` decodes differently from us.
+        # transports), so accepting one requires an explicit trust decision while the
+        # private-address policy is enabled. Sessions we build ourselves refuse to
+        # connect to a non-public address, which is the check that holds against DNS
+        # rebinding and host spellings that `requests` decodes differently from us.
         self._request_proxies: Optional[Dict[str, str]] = None
         if requests_session is not None:
+            if not self._app_config.allow_private_download_urls and not requests_session_is_trusted:
+                raise ValueError(
+                    "A caller-supplied requests_session bypasses the SSRF socket guard. "
+                    "Pass requests_session_is_trusted=True only for a trusted session, or use "
+                    "allow_private_download_urls to opt out of the private-address policy."
+                )
             self._requests = requests_session
         elif self._app_config.allow_private_download_urls:
             # The operator has opted out of the address policy, so ambient proxy variables

tests/app/services/download/test_download_queue.py

@@ -41,6 +41,7 @@ def event_handler(job: DownloadJob, excp: Optional[Exception] = None) -> None:
 
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
     job = queue.download(
@@ -67,6 +68,7 @@ def event_handler(job: DownloadJob, excp: Optional[Exception] = None) -> None:
 def test_errors(tmp_path: Path, mm2_session: Session) -> None:
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
 
@@ -99,7 +101,7 @@ def test_completed_resume_with_416_promotes_in_progress_file(tmp_path: Path) ->
         TestAdapter(b"", status=416, headers={"Content-Range": f"bytes */{len(content)}"}),
     )
     completed_files: list[bool] = []
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = queue.download(
@@ -127,7 +129,7 @@ def test_headerless_416_falls_back_to_recorded_size(tmp_path: Path) -> None:
 
     session = TestSession()
     session.mount(str(source), TestAdapter(b"", status=416))
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = DownloadJob(source=source, dest=destination, expected_total_bytes=len(content))
@@ -150,7 +152,7 @@ def test_headerless_416_without_recorded_size_pauses(tmp_path: Path) -> None:
 
     session = TestSession()
     session.mount(str(source), TestAdapter(b"", status=416))
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = queue.download(source=source, dest=destination)
@@ -176,7 +178,7 @@ def test_mismatched_416_resume_keeps_in_progress_file(tmp_path: Path) -> None:
         str(source),
         TestAdapter(b"", status=416, headers={"Content-Range": "bytes */8"}),
     )
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = queue.download(source=source, dest=destination)
@@ -194,7 +196,7 @@ def test_mismatched_416_resume_keeps_in_progress_file(tmp_path: Path) -> None:
 def test_event_bus(tmp_path: Path, mm2_session: Session) -> None:
     event_bus = TestEventService()
 
-    queue = DownloadQueueService(requests_session=mm2_session, event_bus=event_bus)
+    queue = DownloadQueueService(requests_session=mm2_session, requests_session_is_trusted=True, event_bus=event_bus)
     queue.start()
     queue.download(
         source=AnyHttpUrl("http://www.civitai.com/models/12345"),
@@ -230,6 +232,7 @@ def test_event_bus(tmp_path: Path, mm2_session: Session) -> None:
 def test_broken_callbacks(tmp_path: Path, mm2_session: Session, capsys) -> None:
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
 
@@ -262,7 +265,7 @@ def broken_callback(job: DownloadJob) -> None:
 def test_cancel(tmp_path: Path, mm2_session: Session) -> None:
     event_bus = TestEventService()
 
-    queue = DownloadQueueService(requests_session=mm2_session, event_bus=event_bus)
+    queue = DownloadQueueService(requests_session=mm2_session, requests_session_is_trusted=True, event_bus=event_bus)
     queue.start()
 
     cancelled = False
@@ -303,6 +306,7 @@ def event_handler(job: DownloadJob | MultiFileDownloadJob, excp: Optional[Except
 
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
     job = queue.multifile_download(
@@ -343,6 +347,7 @@ def event_handler(job: DownloadJob | MultiFileDownloadJob, excp: Optional[Except
 
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
     files = metadata.download_urls(session=mm2_session)
@@ -369,7 +374,7 @@ def event_handler(job: DownloadJob | MultiFileDownloadJob, excp: Optional[Except
 def test_multifile_cancel(tmp_path: Path, mm2_session: Session, monkeypatch: Any) -> None:
     event_bus = TestEventService()
 
-    queue = DownloadQueueService(requests_session=mm2_session, event_bus=event_bus)
+    queue = DownloadQueueService(requests_session=mm2_session, requests_session_is_trusted=True, event_bus=event_bus)
     queue.start()
 
     cancelled = False
@@ -400,6 +405,7 @@ def cancelled_callback(job: DownloadJob) -> None:
 def test_multifile_onefile(tmp_path: Path, mm2_session: Session) -> None:
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
     job = queue.multifile_download(
@@ -424,6 +430,7 @@ def test_multifile_download_with_relative_dest(tmp_path: Path, mm2_session: Sess
     monkeypatch.chdir(tmp_path)
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
     queue.start()
     job = queue.multifile_download(
@@ -448,6 +455,7 @@ def test_multifile_download_with_relative_dest(tmp_path: Path, mm2_session: Sess
 def test_multifile_no_rel_paths(tmp_path: Path, mm2_session: Session) -> None:
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
 
     with pytest.raises(ValueError) as error:
@@ -461,6 +469,7 @@ def test_multifile_no_rel_paths(tmp_path: Path, mm2_session: Session) -> None:
 def test_multifile_no_parent_traversal_paths(tmp_path: Path, mm2_session: Session) -> None:
     queue = DownloadQueueService(
         requests_session=mm2_session,
+        requests_session_is_trusted=True,
     )
 
     with pytest.raises(ValueError) as error:
@@ -488,7 +497,7 @@ def test_tokens(tmp_path: Path, mm2_session: Session):
     with clear_config():
         config = get_config()
         config.remote_api_tokens = [URLRegexTokenPair(url_regex="civitai", token="cv_12345")]
-        queue = DownloadQueueService(requests_session=mm2_session)
+        queue = DownloadQueueService(requests_session=mm2_session, requests_session_is_trusted=True)
         queue.start()
         # this one has an access token assigned
         job1 = queue.download(
@@ -522,6 +531,25 @@ def test_production_queue_uses_guarded_session_by_default() -> None:
         queue._requests.close()
 
 
+def test_caller_supplied_session_requires_explicit_trust() -> None:
+    session = Session()
+    with pytest.raises(ValueError, match="requests_session_is_trusted=True"):
+        DownloadQueueService(
+            app_config=InvokeAIAppConfig(allow_private_download_urls=False),
+            requests_session=session,
+        )
+
+
+def test_caller_supplied_session_accepts_explicit_trust() -> None:
+    session = Session()
+    queue = DownloadQueueService(
+        app_config=InvokeAIAppConfig(allow_private_download_urls=False),
+        requests_session=session,
+        requests_session_is_trusted=True,
+    )
+    assert queue._requests is session
+
+
 def test_production_queue_allows_explicit_private_download_opt_in() -> None:
     queue = DownloadQueueService(app_config=InvokeAIAppConfig(allow_private_download_urls=True))
     try:
@@ -580,7 +608,7 @@ def test_download_refuses_non_public_source(tmp_path: Path) -> None:
     session = TestSession()
     session.mount(str(source), TestAdapter(b"secret", status=200))
 
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = queue.download(source=source, dest=tmp_path)
@@ -604,7 +632,7 @@ def test_download_refuses_redirect_to_non_public_address(tmp_path: Path) -> None
     )
     session.mount("http://169.254.169.254/", TestAdapter(b"cloud-credentials", status=200))
 
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = queue.download(source=source, dest=tmp_path)
@@ -618,7 +646,7 @@ def test_download_refuses_redirect_to_non_public_address(tmp_path: Path) -> None
 
 
 def test_rejected_redirect_closes_streamed_response() -> None:
-    queue = DownloadQueueService(requests_session=TestSession())
+    queue = DownloadQueueService(requests_session=TestSession(), requests_session_is_trusted=True)
     response = Response()
     response.status_code = 302
     response.url = "https://public.example/redirect"
@@ -646,7 +674,7 @@ def test_download_refuses_multi_hop_redirect_to_non_public_address(tmp_path: Pat
         TestAdapter(b"", status=302, headers={"Location": "http://169.254.169.254/latest/meta-data/"}),
     )
 
-    queue = DownloadQueueService(requests_session=session)
+    queue = DownloadQueueService(requests_session=session, requests_session_is_trusted=True)
     queue.start()
     try:
         job = queue.download(source=source, dest=tmp_path)

tests/backend/model_manager/model_manager_fixtures.py

@@ -84,7 +84,7 @@ def mm2_app_config(mm2_root_dir: Path) -> InvokeAIAppConfig:
 
 @pytest.fixture
 def mm2_download_queue(mm2_session: Session) -> DownloadQueueServiceBase:
-    download_queue = DownloadQueueService(requests_session=mm2_session)
+    download_queue = DownloadQueueService(requests_session=mm2_session, requests_session_is_trusted=True)
     download_queue.start()
     yield download_queue
     download_queue.stop()