From f08b9aa9eb50685835433d54c4e7ae73e50e47b2 Mon Sep 17 00:00:00 2001
From: bymyself <cbyrne@comfy.org>
Date: Fri, 4 Jul 2025 14:11:09 -0700
Subject: [PATCH] [fix] Add validation and improve device detection in registry
 search

- Add sortField validation to prevent malformed API requests
- Improve GPU device type detection with additional vendor patterns (nvidia, amd, apple)
- Prevent potential injection via sort parameters
---
 browser_tests/assets/save_animated_png.json   | 190 ++++++++++++++++++
 .../providers/registrySearchProvider.ts       |  15 +-
 2 files changed, 201 insertions(+), 4 deletions(-)
 create mode 100644 browser_tests/assets/save_animated_png.json

diff --git a/browser_tests/assets/save_animated_png.json b/browser_tests/assets/save_animated_png.json
new file mode 100644
index 000000000..6cfaeed1a
--- /dev/null
+++ b/browser_tests/assets/save_animated_png.json
@@ -0,0 +1,190 @@
+{
+  "id": "cffcce2d-a13c-4a5f-929b-82f274bacc36",
+  "revision": 0,
+  "last_node_id": 14,
+  "last_link_id": 14,
+  "nodes": [
+    {
+      "id": 10,
+      "type": "LoadImage",
+      "pos": [
+        -361.02374267578125,
+        -40.05255126953125
+      ],
+      "size": [
+        274.080078125,
+        314
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            11
+          ]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "ComfyUI_00137_.png",
+        "image"
+      ]
+    },
+    {
+      "id": 12,
+      "type": "ImageBatch",
+      "pos": [
+        146.92184448242188,
+        104.8472671508789
+      ],
+      "size": [
+        140,
+        46
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image1",
+          "type": "IMAGE",
+          "link": 11
+        },
+        {
+          "name": "image2",
+          "type": "IMAGE",
+          "link": 12
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            14
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "ImageBatch"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 14,
+      "type": "SaveAnimatedPNG",
+      "pos": [
+        457.4212646484375,
+        39.56276321411133
+      ],
+      "size": [
+        270,
+        368
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 14
+        }
+      ],
+      "outputs": [],
+      "properties": {},
+      "widgets_values": [
+        "ComfyUI",
+        6,
+        4
+      ]
+    },
+    {
+      "id": 11,
+      "type": "LoadImage",
+      "pos": [
+        -360.4931640625,
+        326.1943664550781
+      ],
+      "size": [
+        274.080078125,
+        314
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            12
+          ]
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "ComfyUI_00153_.png",
+        "image"
+      ]
+    }
+  ],
+  "links": [
+    [
+      11,
+      10,
+      0,
+      12,
+      0,
+      "IMAGE"
+    ],
+    [
+      12,
+      11,
+      0,
+      12,
+      1,
+      "IMAGE"
+    ],
+    [
+      14,
+      12,
+      0,
+      14,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.129559245649766,
+      "offset": [
+        768.6140137916129,
+        203.6152852376302
+      ]
+    },
+    "frontendVersion": "1.22.2"
+  },
+  "version": 0.4
+}
\ No newline at end of file
diff --git a/src/services/providers/registrySearchProvider.ts b/src/services/providers/registrySearchProvider.ts
index f3ee58c4e..86a1d986a 100644
--- a/src/services/providers/registrySearchProvider.ts
+++ b/src/services/providers/registrySearchProvider.ts
@@ -86,6 +86,10 @@ export const useComfyRegistrySearchProvider = (): NodePackSearchProvider => {
 
       // Apply sort if provided (only supported by list endpoint)
       if (sortField) {
+        // Validate sort field to prevent malformed API requests
+        if (!/^[a-zA-Z_]+$/.test(sortField)) {
+          throw new Error(`Invalid sort field: ${sortField}`)
+        }
         const sortParam =
           sortDirection === 'desc' ? `${sortField};desc` : sortField
         listParams.sort = [sortParam]
@@ -161,12 +165,15 @@ export const useComfyRegistrySearchProvider = (): NodePackSearchProvider => {
     const stats = systemStatsStore.systemStats
     if (!stats?.devices || stats.devices.length === 0) return undefined
 
-    // Look for the first GPU device
+    // Look for the first GPU device - check for additional patterns
     for (const device of stats.devices) {
       const deviceType = device.type.toLowerCase()
-      if (deviceType.includes('cuda')) return 'cuda'
-      if (deviceType.includes('mps')) return 'mps'
-      if (deviceType.includes('rocm')) return 'rocm'
+      if (deviceType.includes('nvidia') || deviceType.includes('cuda'))
+        return 'cuda'
+      if (deviceType.includes('apple') || deviceType.includes('mps'))
+        return 'mps'
+      if (deviceType.includes('amd') || deviceType.includes('rocm'))
+        return 'rocm'
       if (deviceType.includes('directml')) return 'directml'
     }
     return undefined