diff --git a/server.py b/server.py index 27b14825e..71d9e93a2 100644 --- a/server.py +++ b/server.py @@ -555,7 +555,7 @@ class PromptServer(): buffer.seek(0) return web.Response(body=buffer.read(), content_type=f'image/{image_format}', - headers={"Content-Disposition": f"filename=\"{filename}\""}) + headers={"Content-Disposition": f"attachment; filename=\"{filename}\""}) if 'channel' not in request.rel_url.query: channel = 'rgba' @@ -575,7 +575,7 @@ class PromptServer(): buffer.seek(0) return web.Response(body=buffer.read(), content_type='image/png', - headers={"Content-Disposition": f"filename=\"{filename}\""}) + headers={"Content-Disposition": f"attachment; filename=\"{filename}\""}) elif channel == 'a': with Image.open(file) as img: @@ -592,7 +592,7 @@ class PromptServer(): alpha_buffer.seek(0) return web.Response(body=alpha_buffer.read(), content_type='image/png', - headers={"Content-Disposition": f"filename=\"{filename}\""}) + headers={"Content-Disposition": f"attachment; filename=\"{filename}\""}) else: # Use the content type from asset resolution if available, # otherwise guess from the filename. @@ -609,7 +609,7 @@ class PromptServer(): return web.FileResponse( file, headers={ - "Content-Disposition": f"filename=\"{filename}\"", + "Content-Disposition": f"attachment; filename=\"{filename}\"", "Content-Type": content_type } ) diff --git a/tests/test_view_content_disposition.py b/tests/test_view_content_disposition.py new file mode 100644 index 000000000..4fa12db63 --- /dev/null +++ b/tests/test_view_content_disposition.py @@ -0,0 +1,156 @@ +""" +Tests for the /view endpoint Content-Disposition header fix. + +Verifies that the Content-Disposition header complies with RFC 2183 +by using `attachment; filename="..."` format instead of just `filename="..."`. + +See: https://github.com/Comfy-Org/ComfyUI/issues/8914 +""" + +import pytest +import subprocess +import time +import os +import tempfile +import struct +import pytest + +# Simple 1x1 PNG image (minimal valid PNG) +MINIMAL_PNG = bytes([ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, # PNG signature + 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, # IHDR chunk length + type + 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, # 1x1 dimensions + 0x08, 0x02, 0x00, 0x00, 0x00, # bit depth, color type, etc + 0x90, 0x77, 0x53, 0xDE, # IHDR CRC + 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54, # IDAT chunk + 0x08, 0xD7, 0x63, 0xF8, 0xFF, 0xFF, 0xFF, 0x00, # image data + 0x05, 0xFE, 0x02, 0xFE, # IDAT CRC + 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, # IEND chunk + 0xAE, 0x42, 0x60, 0x82, # IEND CRC +]) + + +@pytest.mark.execution +class TestViewContentDisposition: + """Test suite for Content-Disposition header in /view endpoint.""" + + @pytest.fixture(scope="class", autouse=True) + def output_dir(self, args_pytest): + """Use the test output directory.""" + return args_pytest["output_dir"] + + @pytest.fixture(scope="class", autouse=True) + def _server(self, args_pytest): + """Start ComfyUI server for testing.""" + pargs = [ + 'python', 'main.py', + '--output-directory', args_pytest["output_dir"], + '--listen', args_pytest["listen"], + '--port', str(args_pytest["port"]), + '--cpu', + ] + p = subprocess.Popen(pargs, cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + yield + p.kill() + p.wait() + + @pytest.fixture(scope="class") + def server_url(self, args_pytest): + """Get the server base URL.""" + return f"http://{args_pytest['listen']}:{args_pytest['port']}" + + @pytest.fixture(scope="class") + def client(self, server_url, _server): + """Create HTTP client with retry.""" + import requests + session = requests.Session() + n_tries = 10 + for i in range(n_tries): + time.sleep(2) + try: + resp = session.get(f"{server_url}/system_stats", timeout=5) + if resp.status_code == 200: + break + except requests.exceptions.ConnectionError: + if i == n_tries - 1: + raise + return session + + @pytest.fixture + def test_image_file(self, output_dir): + """Create a temporary test image file and clean it up after the test.""" + filename = "test_content_disposition.png" + filepath = os.path.join(output_dir, filename) + with open(filepath, 'wb') as f: + f.write(MINIMAL_PNG) + yield filename, filepath + # Cleanup + if os.path.exists(filepath): + os.remove(filepath) + + def test_view_content_disposition_is_attachment_format(self, client, server_url, test_image_file): + """Test that /view endpoint returns Content-Disposition with RFC 2183 compliant format. + + The header should be: Content-Disposition: attachment; filename="test_content_disposition.png" + NOT: Content-Disposition: filename="test_content_disposition.png" + + This test validates the fix for: https://github.com/Comfy-Org/ComfyUI/issues/8914 + """ + filename, _ = test_image_file + response = client.get( + f"{server_url}/view", + params={"filename": filename}, + timeout=10 + ) + + assert response.status_code == 200, f"Expected 200, got {response.status_code}" + + content_disposition = response.headers.get("Content-Disposition", "") + assert content_disposition, "Content-Disposition header should be present" + + # RFC 2183 compliant format: "attachment; filename=\"name.ext\"" + expected_prefix = f"attachment; filename=\"{filename}\"" + assert content_disposition == expected_prefix, ( + f"Content-Disposition header should be RFC 2183 compliant.\n" + f"Expected: {expected_prefix}\n" + f"Got: {content_disposition}" + ) + + def test_view_with_subfolder_content_disposition(self, client, server_url, output_dir): + """Test Content-Disposition with subfolder parameter.""" + import requests + + # Create a subfolder and test image + subfolder = "test_subfolder" + subfolder_path = os.path.join(output_dir, subfolder) + os.makedirs(subfolder_path, exist_ok=True) + + filename = "test_subfolder.png" + filepath = os.path.join(subfolder_path, filename) + with open(filepath, 'wb') as f: + f.write(MINIMAL_PNG) + + try: + response = client.get( + f"{server_url}/view", + params={"filename": filename, "subfolder": subfolder}, + timeout=10 + ) + + assert response.status_code == 200, f"Expected 200, got {response.status_code}" + + content_disposition = response.headers.get("Content-Disposition", "") + assert content_disposition, "Content-Disposition header should be present" + + expected_prefix = f"attachment; filename=\"{filename}\"" + assert content_disposition == expected_prefix, ( + f"Content-Disposition header should be RFC 2183 compliant.\n" + f"Expected: {expected_prefix}\n" + f"Got: {content_disposition}" + ) + finally: + # Cleanup + if os.path.exists(filepath): + os.remove(filepath) + if os.path.exists(subfolder_path): + os.rmdir(subfolder_path)