Summary
ComfyUI v0.23.0 contains an unsafe deserialization flaw in the LoadTrainingDataset node that allows crafted pickle files to be deserialized using torch.load without safety constraints.
By uploading a malicious .pkl shard via unauthenticated endpoints and referencing it in a queued workflow, an attacker can cause arbitrary Python code execution in the ComfyUI process context.
Affected context
This vulnerability appears in ComfyUI deployments that expose the image upload and workflow queuing endpoints over a network, such as in remote or shared inference environments.
It is most relevant when:
- The
/upload/imageendpoint accepts pickle-backed dataset shards - The
LoadTrainingDatasetnode is available in user workflows torch.loadis used to deserialize uploaded files within the ComfyUI runtime
Risk increases where the ComfyUI process runs with elevated permissions or has access to sensitive data or internal services.
Exploit preconditions
- Attack surface: HTTP API endpoints, specifically unauthenticated POST requests to
/upload/imageand/prompt - Required access: Network reachability to the ComfyUI instance; no authentication needed under default exposure
- Trigger point: A workflow graph that uses
LoadTrainingDatasetand references an attacker-uploadedshard_*.pklfile, causingtorch.loadto deserialize it - Why exploitation is possible: Pickle deserialization via
torch.loadimplicitly trusts file contents and invokes special methods like__reduce__, enabling arbitrary object construction and code execution when the payload is attacker-controlled
Why this happens
The core issue is the use of a general-purpose, code-capable serialization format (Python pickle) for data loading in a context where inputs can be supplied by unauthenticated users.
The LoadTrainingDataset node invokes torch.load on user-uploaded .pkl shards, which internally uses pickle deserialization and executes constructors and reducer functions embedded in the payload.
This aligns with CWE-502: Unsafe Deserialization, where untrusted data is deserialized into live objects without validation, sandboxing, or format restrictions, turning the deserialization step into an execution vector rather than a pure data transformation.
Potential impact
- Remote execution of arbitrary Python code within the ComfyUI process
- Reading, modifying, or deleting data accessible to the ComfyUI runtime
- Lateral movement to other services or hosts reachable from the ComfyUI server
- Tampering with model weights, training data, or inference output
- Use of the compromised host for further attacks, such as launching malware or pivoting inside the network
General mitigation guidance
Effective mitigation focuses on avoiding insecure deserialization patterns and reducing trust in user-provided artifacts.
Recommended general approaches include:
- Replacing pickle-based workflows with data formats that cannot encode executable objects (e.g., JSON, protocol buffers, or domain-specific binary formats)
- Introducing strict validation and whitelisting for files consumed by training or dataset-loading nodes, ensuring they match expected schemas and types
- Segregating upload and processing paths so that user-supplied files are never directly fed into deserialization APIs capable of executing code
- Running ComfyUI with least privilege, minimizing filesystem and network access to limit impact if deserialization is abused
- Enforcing authentication and authorization on all upload and workflow queuing endpoints to restrict who can introduce new artifacts into the system
How teams detect issues like this
Detection typically centers on identifying unsafe deserialization usage and monitoring runtime behavior.
Useful strategies include:
- Static analysis rules that flag calls to
pickle.load,pickle.loads, ortorch.loadon data originating from HTTP requests or other untrusted sources - Code review practices that require explicit justification and documentation when using executable serialization formats like pickle
- Runtime logging and anomaly detection around file uploads and dataset loading operations, focusing on unusual
.pklfile patterns or unexpected workflow graphs - Host-level monitoring for suspicious child processes, file changes, or network connections initiated by the ComfyUI process
- Incorporating automated security scanning tools into CI pipelines to enforce policies on deserialization APIs and high-risk data flows
