Style fix. (#8390)

* Update fix for potential XSS on /view

This commit uses mimetypes to add more restricted filetypes to prevent from being served, since mimetypes are what browsers use to determine how to serve files.

* Fix typo

Fixed a typo that prevented the program from running

comfy_extras/nodes_images.py

@@ -16,8 +16,7 @@ from inspect import cleandoc
 import torch
 import comfy.utils
 
-from comfy.comfy_types import FileLocator, IO
+from comfy.comfy_types import FileLocator
-from server import PromptServer
 
 MAX_RESOLUTION = nodes.MAX_RESOLUTION
 
@@ -492,36 +491,6 @@ class SaveSVGNode:
             counter += 1
         return { "ui": { "images": results } }
 
-class GetImageSize:
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": (IO.IMAGE,),
-            },
-            "hidden": {
-                "unique_id": "UNIQUE_ID",
-            }
-        }
-
-    RETURN_TYPES = (IO.INT, IO.INT)
-    RETURN_NAMES = ("width", "height")
-    FUNCTION = "get_size"
-
-    CATEGORY = "image"
-    DESCRIPTION = """Returns width and height of the image, and passes it through unchanged."""
-
-    def get_size(self, image, unique_id=None) -> tuple[int, int]:
-        height = image.shape[1]
-        width = image.shape[2]
-
-        # Send progress text to display size on the node
-        if unique_id:
-            PromptServer.instance.send_progress_text(f"width: {width}, height: {height}", unique_id)
-
-        return width, height
-
 NODE_CLASS_MAPPINGS = {
     "ImageCrop": ImageCrop,
     "RepeatImageBatch": RepeatImageBatch,
@@ -531,5 +500,4 @@ NODE_CLASS_MAPPINGS = {
     "SaveAnimatedPNG": SaveAnimatedPNG,
     "SaveSVGNode": SaveSVGNode,
     "ImageStitch": ImageStitch,
-    "GetImageSize": GetImageSize,
 }

nodes.py

@@ -2067,7 +2067,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImageQuantize": "Image Quantize",
     "ImageSharpen": "Image Sharpen",
     "ImageScaleToTotalPixels": "Scale Image to Total Pixels",
-    "GetImageSize": "Get Image Size",
     # _for_testing
     "VAEDecodeTiled": "VAE Decode (Tiled)",
     "VAEEncodeTiled": "VAE Encode (Tiled)",

requirements.txt

@@ -1,4 +1,4 @@
-comfyui-frontend-package==1.21.5
+comfyui-frontend-package==1.21.6
 comfyui-workflow-templates==0.1.25
 comfyui-embedded-docs==0.2.0
 torch

server.py

@@ -390,7 +390,7 @@ class PromptServer():
         async def view_image(request):
             if "filename" in request.rel_url.query:
                 filename = request.rel_url.query["filename"]
-                filename,output_dir = folder_paths.annotated_filepath(filename)
+                filename, output_dir = folder_paths.annotated_filepath(filename)
 
                 if not filename:
                     return web.Response(status=400)
@@ -476,9 +476,8 @@ class PromptServer():
                         # Get content type from mimetype, defaulting to 'application/octet-stream'
                         content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
 
-                        # For security, force certain extensions to download instead of display
+                        # For security, force certain mimetypes to download instead of display
-                        file_extension = os.path.splitext(filename)[1].lower()
+                        if content_type in {'text/html', 'text/html-sandboxed', 'application/xhtml+xml', 'text/javascript', 'text/css'}:
-                        if file_extension in {'.html', '.htm', '.js', '.css'}:
                             content_type = 'application/octet-stream'  # Forces download
 
                         return web.FileResponse(

tests-unit/comfy_extras_test/image_stitch_test.py

@@ -5,10 +5,7 @@ from unittest.mock import patch, MagicMock
 mock_nodes = MagicMock()
 mock_nodes.MAX_RESOLUTION = 16384
 
-# Mock server module for PromptServer
+with patch.dict('sys.modules', {'nodes': mock_nodes}):
-mock_server = MagicMock()
-
-with patch.dict('sys.modules', {'nodes': mock_nodes, 'server': mock_server}):
     from comfy_extras.nodes_images import ImageStitch