From 3fd028480282c3ffacf5f0226ba69ef202eb47ac Mon Sep 17 00:00:00 2001
From: Mattabyte <matthew.inger@intelligentpathways.com.au>
Date: Tue, 28 Jan 2025 14:48:56 +1000
Subject: [PATCH] URL and Environment Variable Validation

---
 nodes/secure_api_call.py | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/nodes/secure_api_call.py b/nodes/secure_api_call.py
index 424d3da..cbac8d6 100644
--- a/nodes/secure_api_call.py
+++ b/nodes/secure_api_call.py
@@ -2,10 +2,24 @@ import requests
 import json
 import server
 import os
+import urllib.parse
+from typing import Optional
 
 from .util import ComfyAnyType
 
 class SecureApiCall:
+    ALLOWED_SCHEMES = {'https', 'http'}
+    
+    @staticmethod
+    def validate_url(url: str) -> Optional[str]:
+        try:
+            parsed = urllib.parse.urlparse(url)
+            if parsed.scheme not in SecureApiCall.ALLOWED_SCHEMES:
+                return f"URL scheme must be one of: {', '.join(SecureApiCall.ALLOWED_SCHEMES)}"
+            return None
+        except Exception:
+            return "Invalid URL format"
+
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -26,24 +40,40 @@ class SecureApiCall:
     CATEGORY = "notifications"
 
     def hook(self, any, api_url, api_auth, data, full_comfyui_info, timeout, verify_ssl):
+        # Validate URL before proceeding
+        url_error = self.validate_url(api_url)
+        if url_error:
+            raise ValueError(f"Invalid API URL: {url_error}")
+
         # Get the payload from the data string
-        payload = json.loads(data)
+        try:
+            payload = json.loads(data)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON data provided")
+
         # Add ComfyUI execution information to the payload
         current_queue = server.PromptServer.instance.prompt_queue.get_current_queue()
 
         # The response is a massive object (containing all nodes and workflow info) -- we only want the prompt_id for tracking, but if full_comfyui_info is true, we want to include the entire object
         prompt_id = current_queue[0][0][1]
         payload.update({"comfyui_execution_info": {"prompt_id": prompt_id}})
-        if include_full_comfyui_info:
+        if full_comfyui_info:
             payload.update({"comfyui_info": {"full_comfyui_info": current_queue}})
 
 
         # Resolve any environment variables in the api_url, api_auth fields -- Only allow COMFYUI_SECUREAPICALL_ prefixed variables
         if api_url.startswith("$ENV."):
-            env_var = f"COMFYUI_SECUREAPICALL_{api_url.removeprefix('$ENV.')}"
+            # Only allow alphanumeric and underscore characters in env var names
+            env_name = api_url.removeprefix("$ENV.")
+            if not env_name.replace("_", "").isalnum():
+                raise ValueError("Environment variable names must only contain letters, numbers, and underscores")
+            env_var = f"COMFYUI_SECUREAPICALL_{env_name}"
             api_url = os.getenv(env_var) or 'NOT_SET'
         if api_auth.startswith("$ENV."):
-            env_var = f"COMFYUI_SECUREAPICALL_{api_auth.removeprefix('$ENV.')}"
+            env_name = api_auth.removeprefix("$ENV.")
+            if not env_name.replace("_", "").isalnum():
+                raise ValueError("Environment variable names must only contain letters, numbers, and underscores")
+            env_var = f"COMFYUI_SECUREAPICALL_{env_name}"
             api_auth = os.getenv(env_var) or 'NOT_SET'
 
         # Send the payload to the API