diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index ca7884482f6..a0003bc1009 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/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. diff --git a/invokeai/app/services/download/download_default.py b/invokeai/app/services/download/download_default.py index 94154c3b493..e61e2f9830a 100644 --- a/invokeai/app/services/download/download_default.py +++ b/invokeai/app/services/download/download_default.py @@ -51,6 +51,7 @@ 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. @@ -58,6 +59,7 @@ def __init__( :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 diff --git a/tests/app/services/download/test_download_queue.py b/tests/app/services/download/test_download_queue.py index ab6d1dd42fd..dcadaf02c37 100644 --- a/tests/app/services/download/test_download_queue.py +++ b/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) diff --git a/tests/backend/model_manager/model_manager_fixtures.py b/tests/backend/model_manager/model_manager_fixtures.py index 9ab4570dded..4f8da686d0f 100644 --- a/tests/backend/model_manager/model_manager_fixtures.py +++ b/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()