diff --git a/src/specify_cli/workflows/_commands.py b/src/specify_cli/workflows/_commands.py index c3ba89157..e075de077 100644 --- a/src/specify_cli/workflows/_commands.py +++ b/src/specify_cli/workflows/_commands.py @@ -154,14 +154,21 @@ _MAX_WORKFLOW_YAML_BYTES = 5 * 1024 * 1024 # 5 MiB _DOWNLOAD_CHUNK_SIZE = 65536 -def _read_response_within_limit(response, max_bytes: int = _MAX_WORKFLOW_YAML_BYTES) -> bytes: +def _read_response_within_limit(response, max_bytes: int | None = None) -> bytes: """Read *response* fully, enforcing *max_bytes* via bounded streaming. A ``Content-Length`` header is checked up front to fail fast, but it is never trusted alone: the actual bytes read are also counted as they stream in, so a chunked or ``Content-Length``-less response that lies about (or omits) its size still cannot exceed the limit. + + ``max_bytes`` defaults to ``None`` (resolved to the module-level + ``_MAX_WORKFLOW_YAML_BYTES`` at call time, not at function-definition + time) so tests can override the effective limit via monkeypatching the + module attribute. """ + if max_bytes is None: + max_bytes = _MAX_WORKFLOW_YAML_BYTES content_length = None getheader = getattr(response, "getheader", None) if callable(getheader): @@ -910,6 +917,7 @@ def workflow_add( _wf_url_extra_headers = {"Accept": "application/octet-stream"} import tempfile + tmp_path: Path | None = None try: with _open_url( download_url, @@ -933,11 +941,17 @@ def workflow_add( ) raise typer.Exit(1) with tempfile.NamedTemporaryFile(suffix=".yml", delete=False) as tmp: - tmp.write(_read_response_within_limit(resp)) + # Assign tmp_path immediately: NamedTemporaryFile(delete=False) + # creates the file on disk right away, before any bytes are + # written, so a failure in the size-limited read below must + # still be able to find and remove it. tmp_path = Path(tmp.name) + tmp.write(_read_response_within_limit(resp)) except typer.Exit: raise except Exception as exc: + if tmp_path is not None: + tmp_path.unlink(missing_ok=True) console.print(f"[red]Error:[/red] Failed to download workflow: {_escape_markup(str(exc))}") raise typer.Exit(1) try: diff --git a/tests/test_workflows.py b/tests/test_workflows.py index 4242146ca..85339b95b 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -7675,7 +7675,7 @@ steps: ) assert result.exit_code != 0 assert result.exception is None or isinstance(result.exception, SystemExit) - assert result.output.strip() != "" + assert "exceedingthe100-byteworkflowsizelimit" in "".join(result.output.split()) def test_add_from_url_rejects_oversized_streamed_body_without_content_length( self, project_dir, monkeypatch @@ -7705,7 +7705,77 @@ steps: ) assert result.exit_code != 0 assert result.exception is None or isinstance(result.exception, SystemExit) - assert result.output.strip() != "" + assert "exceedsthe100-byteworkflowsizelimit" in "".join(result.output.split()) + + def test_add_from_url_oversized_streamed_body_leaves_no_temp_file( + self, project_dir, monkeypatch, tmp_path + ): + """A rejected --from download (oversized streamed body, no + Content-Length) must not leave the 0-byte NamedTemporaryFile behind: + the file is created on disk as soon as it is opened (delete=False), + before any bytes are written, so a failure inside the size-limit + check must still clean it up rather than merely erroring out.""" + import tempfile as tempfile_mod + from unittest.mock import patch + from typer.testing import CliRunner + from specify_cli import app + from specify_cli.workflows import _commands as wf_commands + + monkeypatch.chdir(project_dir) + monkeypatch.setattr(wf_commands, "_MAX_WORKFLOW_YAML_BYTES", 100) + scratch_tmp = tmp_path / "scratch-tmp" + scratch_tmp.mkdir() + monkeypatch.setattr(tempfile_mod, "tempdir", str(scratch_tmp)) + oversized_body = b"x" * 500 # no Content-Length header at all + runner = CliRunner() + with patch( + "specify_cli.authentication.http.open_url", + side_effect=lambda url, timeout=None, extra_headers=None, redirect_validator=None: self._FakeResponse( + oversized_body, url + ), + ): + result = runner.invoke( + app, + ["workflow", "add", "align-wf", "--from", "https://example.com/workflow.yml"], + ) + assert result.exit_code != 0 + assert "exceedsthe100-byteworkflowsizelimit" in "".join(result.output.split()) + leaked = list(scratch_tmp.glob("*.yml")) + assert leaked == [], f"leaked temp files: {leaked}" + + def test_add_from_url_oversized_content_length_leaves_no_temp_file( + self, project_dir, monkeypatch, tmp_path + ): + """Same guarantee for the fail-fast Content-Length rejection path: + it must not even leave a 0-byte temp file behind.""" + import tempfile as tempfile_mod + from unittest.mock import patch + from typer.testing import CliRunner + from specify_cli import app + from specify_cli.workflows import _commands as wf_commands + + monkeypatch.chdir(project_dir) + monkeypatch.setattr(wf_commands, "_MAX_WORKFLOW_YAML_BYTES", 100) + scratch_tmp = tmp_path / "scratch-tmp" + scratch_tmp.mkdir() + monkeypatch.setattr(tempfile_mod, "tempdir", str(scratch_tmp)) + small_body = b"id: align-wf\n" # small actual body; Content-Length lies + runner = CliRunner() + with patch( + "specify_cli.authentication.http.open_url", + side_effect=lambda url, timeout=None, extra_headers=None, redirect_validator=None: self._FakeResponse( + small_body, url, headers={"Content-Length": "1000"} + ), + ): + result = runner.invoke( + app, + ["workflow", "add", "align-wf", "--from", "https://example.com/workflow.yml"], + ) + assert result.exit_code != 0 + assert "exceedingthe100-byteworkflowsizelimit" in "".join(result.output.split()) + leaked = list(scratch_tmp.glob("*.yml")) + assert leaked == [], f"leaked temp files: {leaked}" + def test_add_from_url_installs(self, project_dir, monkeypatch): from unittest.mock import patch @@ -8207,7 +8277,7 @@ steps: assert result.exit_code != 0 assert result.exception is None or isinstance(result.exception, SystemExit) - assert result.output.strip() != "" + assert "exceedingthe100-byteworkflowsizelimit" in "".join(result.output.split()) dest_dir = project_dir / ".specify" / "workflows" / "align-wf" assert not dest_dir.exists() assert not WorkflowRegistry(project_dir).is_installed("align-wf") @@ -8250,7 +8320,7 @@ steps: assert result.exit_code != 0 assert result.exception is None or isinstance(result.exception, SystemExit) - assert result.output.strip() != "" + assert "exceedsthe100-byteworkflowsizelimit" in "".join(result.output.split()) dest_dir = project_dir / ".specify" / "workflows" / "align-wf" assert not dest_dir.exists() assert not WorkflowRegistry(project_dir).is_installed("align-wf")