upload Image cryptio and key exchange

2025-10-11 18:46:44 +08:00
parent 268d5559b4
commit a5ba357ed7
4 changed files with 911 additions and 2 deletions
@@ -0,0 +1,226 @@
+# CryptIO 加密图片上传系统
+
+## 系统架构
+
+本系统实现了端到端的加密图片上传和预览功能，确保图片在客户端加密，在服务器以加密形式存储，只有在使用时才在服务端解密。
+
+## 密钥管理方案
+
+### 密钥分布
+- **客户端持有**：
+  - 客户端公钥/私钥对
+  - 服务端公钥/私钥对（通过加密传输获得）
+
+- **服务端持有**：
+  - 服务端公钥/私钥对
+  - 客户端公钥
+
+### 密钥交换流程
+
+1. **客户端初始化**：
+   - 检查 `localStorage`，如果没有客户端密钥对则生成新的 RSA-2048 密钥对
+   - 密钥对保存在 `localStorage` 中（键：`cryptio_client_keypair`）
+
+2. **密钥交换**（`POST /cryptio/exchange_keys`）：
+   ```
+   客户端 -> 服务器: 发送客户端公钥
+   服务器: 1. 保存客户端公钥到 keys.json
+          2. 生成随机 AES-256 密钥
+          3. 用 AES 加密服务端密钥对（JSON格式）
+          4. 用客户端公钥加密 AES 密钥
+          5. 返回: {encrypted_aes_key, iv, encrypted_data}
+   客户端: 1. 用客户端私钥解密 AES 密钥
+          2. 用 AES 密钥解密服务端密钥对
+          3. 保存到 localStorage（键：`cryptio_server_keys`）
+   ```
+
+3. **密钥缓存**：
+   - 客户端密钥对和服务端密钥都缓存在 `localStorage`
+   - 下次访问时直接使用缓存，无需重新交换
+
+## 图片上传流程
+
+### 1. 客户端加密上传
+
+```typescript
+文件选择 -> 读取文件 -> 混合加密 -> 上传加密数据 -> 服务器存储
+```
+
+**混合加密方案**（AES + RSA）：
+
+1. 生成随机 AES-256-GCM 密钥
+2. 生成随机 IV（12字节）
+3. 用 AES-GCM 加密图片数据
+4. 用**服务端公钥**加密 AES 密钥
+5. 组合数据格式：
+   ```
+   [4字节: AES密钥长度] +
+   [加密的AES密钥] +
+   [4字节: IV长度] +
+   [IV] +
+   [加密的图片数据]
+   ```
+
+### 2. 服务器存储
+
+- 端点：`POST /cryptio/upload_encrypted`
+- 接收加密数据，直接保存到 `input/` 目录
+- 文件名格式：`cryptio_{uuid}.{ext}.encrypted`
+- **文件以完全加密的形式存储在磁盘上**
+
+### 3. 服务端解密（节点执行时）
+
+当 workflow 执行到 `UploadImageCryptIO` 节点时：
+
+1. 从 `input/` 目录读取 `.encrypted` 文件
+2. 解析组合数据（AES密钥长度、AES密钥、IV长度、IV、加密数据）
+3. 用**服务端私钥**解密 AES 密钥
+4. 用 AES-GCM 解密图片数据
+5. 转换为 PIL Image 对象
+6. 处理并返回 IMAGE tensor 和 MASK tensor
+
+代码位置：`src/cryptio/upload_image.py:48`
+
+## 图片预览流程
+
+### 客户端预览加密图片
+
+```
+选择加密图片 -> 请求加密数据 -> 客户端解密 -> 显示预览
+```
+
+1. **获取加密数据**：
+   - 端点：`GET /cryptio/view_encrypted?filename=xxx.encrypted`
+   - 服务器返回 base64 编码的加密文件内容
+
+2. **客户端解密**：
+   - 用**服务端私钥**解密 AES 密钥
+   - 用 AES 密钥解密图片数据
+   - 创建 Blob URL 显示图片
+
+代码位置：`ts/cryptioUploadImage.ts:350`
+
+## 文件结构
+
+### 服务端
+
+```
+src/cryptio/
+├── keys.py              # 密钥管理
+├── api.py               # API端点
+│   ├── POST /cryptio/exchange_keys        # 密钥交换
+│   ├── POST /cryptio/upload_encrypted     # 上传加密图片
+│   ├── GET  /cryptio/view_encrypted       # 获取加密图片
+│   └── GET  /cryptio/public_key           # 兼容旧接口
+├── upload_image.py      # UploadImageCryptIO节点
+└── nodes.py             # 节点注册
+```
+
+### 客户端
+
+```
+ts/
+├── cryptioUploadImage.ts   # 完整的加密上传实现
+│   ├── 密钥生成和管理
+│   ├── 密钥交换
+│   ├── 文件加密/解密
+│   ├── 自定义上传widget
+│   └── 预览功能
+└── clientEncrypt.ts        # 文本加密节点（旧）
+```
+
+## 安全特性
+
+1. **端到端加密**：
+   - 图片在客户端加密后才上传
+   - 服务器存储加密文件
+   - 只在节点执行时临时解密
+
+2. **混合加密**：
+   - 结合 RSA（非对称）和 AES（对称）的优点
+   - RSA 用于密钥交换，AES 用于数据加密
+   - 每次上传使用不同的随机 AES 密钥和 IV
+
+3. **密钥安全传输**：
+   - 服务端密钥通过客户端公钥加密传输
+   - 客户端密钥永不离开客户端
+
+4. **文件级加密**：
+   - 加密文件直接存储在磁盘
+   - 服务器重启不影响加密数据
+   - 无明文缓存
+
+## 使用方法
+
+### 1. 基础使用
+
+1. 在 ComfyUI 中添加 `UploadImageCryptIO` 节点
+2. 点击 "choose file to upload" 按钮
+3. 选择图片文件
+4. 图片自动加密、上传，并显示预览
+5. 将节点的 `encrypted` 参数设为 `True`
+6. 连接到其他图片处理节点
+
+### 2. 密钥管理
+
+- **重新生成客户端密钥**：清空 localStorage 中的 `cryptio_client_keypair`
+- **重新获取服务端密钥**：清空 localStorage 中的 `cryptio_server_keys`
+- **服务端重新生成密钥**：删除 `keys/keys.json` 文件，重启 ComfyUI
+
+### 3. 兼容性
+
+- 支持所有标准图片格式（JPEG, PNG, WebP, etc.）
+- 支持透明通道和多帧图片
+- 输出与标准 `LoadImage` 节点完全兼容
+
+## 技术栈
+
+- **加密算法**：
+  - RSA-2048 (OAEP, SHA-256)
+  - AES-256-GCM (图片加密)
+  - AES-256-CBC (密钥传输)
+
+- **前端**：
+  - TypeScript
+  - Web Crypto API
+  - ComfyUI Extension API
+
+- **后端**：
+  - Python 3
+  - cryptography library
+  - aiohttp
+
+## 性能考虑
+
+- **客户端加密**：现代浏览器的 Web Crypto API 性能优秀
+- **文件大小**：加密后文件略大（约增加 300-400 字节的元数据）
+- **内存使用**：服务端解密时需要完整加载文件到内存
+- **缓存策略**：密钥缓存在 localStorage，避免重复交换
+
+## 故障排查
+
+### 问题：上传失败
+
+- 检查浏览器控制台错误信息
+- 确认 `/cryptio/exchange_keys` 端点可访问
+- 验证客户端密钥对已生成
+
+### 问题：预览不显示
+
+- 确认文件名以 `.encrypted` 结尾
+- 检查服务端私钥是否正确获取
+- 查看控制台解密错误信息
+
+### 问题：节点执行失败
+
+- 确认 `encrypted` 参数设为 `True`
+- 验证加密文件存在于 `input/` 目录
+- 检查服务端私钥是否可用
+
+## 未来改进
+
+- [ ] 支持批量加密上传
+- [ ] 添加加密进度显示
+- [ ] 支持流式加密（大文件）
+- [ ] 密钥轮换机制
+- [ ] 加密文件完整性校验
@@ -1,13 +1,101 @@
 import os
 import uuid
+import json
+import base64
 from aiohttp import web
 from server import PromptServer
-from .keys import _get_keys
+from .keys import _get_keys, update_client_public_key
 import folder_paths
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.backends import default_backend
+
+
+@PromptServer.instance.routes.post("/cryptio/exchange_keys")
+async def exchange_keys(request):
+    """
+    密钥交换端点：
+    1. 接收客户端公钥
+    2. 保存客户端公钥到服务器
+    3. 使用客户端公钥加密服务端密钥对
+    4. 返回加密后的服务端密钥
+    """
+    try:
+        data = await request.json()
+        client_public_key_pem = data.get("client_public_key")
+
+        if not client_public_key_pem:
+            return web.json_response({"error": "client_public_key is required"}, status=400)
+
+        # 保存客户端公钥
+        client_public_key_bytes = client_public_key_pem.encode("utf-8")
+        update_client_public_key(client_public_key_bytes)
+
+        # 获取服务端密钥
+        keys = _get_keys()
+        server_public_key = keys.get("server_public_key")
+        server_private_key = keys.get("server_private_key")
+
+        if not server_public_key or not server_private_key:
+            return web.json_response({"error": "Server keys not found"}, status=500)
+
+        # 使用客户端公钥加密服务端密钥
+        client_public_key = serialization.load_pem_public_key(
+            client_public_key_bytes, backend=default_backend()
+        )
+
+        # 准备要加密的数据（服务端公钥和私钥）
+        keys_data = json.dumps({
+            "server_public_key": server_public_key.decode("utf-8"),
+            "server_private_key": server_private_key.decode("utf-8"),
+        })
+
+        # 由于RSA加密有大小限制，使用混合加密方案
+        # 1. 生成随机AES密钥
+        from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+        from cryptography.hazmat.primitives import padding as sym_padding
+        import secrets
+
+        aes_key = secrets.token_bytes(32)  # 256-bit AES key
+        iv = secrets.token_bytes(16)  # 128-bit IV
+
+        # 2. 使用AES加密数据
+        cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=default_backend())
+        encryptor = cipher.encryptor()
+
+        # PKCS7 padding
+        padder = sym_padding.PKCS7(128).padder()
+        padded_data = padder.update(keys_data.encode("utf-8")) + padder.finalize()
+
+        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
+
+        # 3. 使用RSA加密AES密钥
+        encrypted_aes_key = client_public_key.encrypt(
+            aes_key,
+            padding.OAEP(
+                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None,
+            ),
+        )
+
+        # 4. 返回加密后的数据
+        return web.json_response({
+            "encrypted_aes_key": base64.b64encode(encrypted_aes_key).decode("ascii"),
+            "iv": base64.b64encode(iv).decode("ascii"),
+            "encrypted_data": base64.b64encode(encrypted_data).decode("ascii"),
+        })
+
+    except Exception as e:
+        print(f"Error in key exchange: {e}")
+        import traceback
+        traceback.print_exc()
+        return web.json_response({"error": str(e)}, status=500)


@PromptServer.instance.routes.get("/cryptio/public_key")
 async def get_public_key(request):
+    """兼容旧的公钥获取接口（不加密传输）"""
    keys = _get_keys()
    pub = keys.get("server_public_key")
    if pub:
@@ -55,4 +143,33 @@ async def upload_encrypted_image(request):

    except Exception as e:
        print(f"Error uploading encrypted image: {e}")
+        return web.json_response({"error": str(e)}, status=500)
+
+
+@PromptServer.instance.routes.get("/cryptio/view_encrypted")
+async def view_encrypted_image(request):
+    """获取加密图片用于前端预览（返回加密数据，由前端解密）"""
+    try:
+        filename = request.query.get("filename")
+        if not filename:
+            return web.json_response({"error": "filename parameter is required"}, status=400)
+
+        # 获取文件路径
+        image_path = folder_paths.get_annotated_filepath(filename)
+
+        if not os.path.exists(image_path):
+            return web.json_response({"error": "File not found"}, status=404)
+
+        # 读取加密文件
+        with open(image_path, 'rb') as f:
+            encrypted_data = f.read()
+
+        # 返回base64编码的加密数据
+        return web.json_response({
+            "encrypted_data": base64.b64encode(encrypted_data).decode("ascii"),
+            "filename": filename
+        })
+
+    except Exception as e:
+        print(f"Error viewing encrypted image: {e}")
        return web.json_response({"error": str(e)}, status=500)
@@ -111,4 +111,24 @@ def _get_keys() -> dict:
    else:
        _KEY_CACHE = {}
        _KEY_CACHE_MTIME = None
-        return _KEY_CACHE
+        return _KEY_CACHE
+
+
+def update_client_public_key(client_public_pem: bytes):
+    """更新客户端公钥到 keys.json"""
+    global _KEY_CACHE, _KEY_CACHE_MTIME
+
+    # 确保服务端密钥存在
+    _ensure_keys_exist()
+    keys = _load_keys_json()
+
+    # 更新客户端公钥
+    _save_keys_json(
+        keys["server_public_key"],
+        keys["server_private_key"],
+        client_public_pem
+    )
+
+    # 清空缓存，强制重新加载
+    _KEY_CACHE = None
+    _KEY_CACHE_MTIME = None
@@ -0,0 +1,546 @@
+//@ts-ignore
+import { app as rawApp } from "../../scripts/app.js";
+//@ts-ignore
+import { api as rawApi } from "../../scripts/api.js";
+import type {
+    ComfyApp,
+    ComfyApi,
+    ComfyNodeDef,
+} from "@comfyorg/comfyui-frontend-types";
+
+const app: ComfyApp = rawApp;
+const api: ComfyApi = rawApi;
+
+// 密钥存储键名
+const CLIENT_KEYPAIR_STORAGE_KEY = "cryptio_client_keypair";
+const SERVER_KEYS_STORAGE_KEY = "cryptio_server_keys";
+
+// 密钥接口
+interface KeyPair {
+    publicKey: string;
+    privateKey: string;
+}
+
+interface ServerKeys {
+    publicKey: string;
+    privateKey: string;
+}
+
+// 生成客户端密钥对
+async function generateClientKeyPair(): Promise<KeyPair> {
+    // 生成RSA密钥对
+    const keyPair = await window.crypto.subtle.generateKey(
+        {
+            name: "RSA-OAEP",
+            modulusLength: 2048,
+            publicExponent: new Uint8Array([1, 0, 1]),
+            hash: "SHA-256",
+        },
+        true,
+        ["encrypt", "decrypt"]
+    );
+
+    // 导出公钥为PEM格式
+    const publicKeyBuffer = await window.crypto.subtle.exportKey(
+        "spki",
+        keyPair.publicKey
+    );
+    const publicKeyBase64 = btoa(
+        String.fromCharCode(...new Uint8Array(publicKeyBuffer))
+    );
+    const publicKeyPem =
+        "-----BEGIN PUBLIC KEY-----\n" +
+        publicKeyBase64.match(/.{1,64}/g)!.join("\n") +
+        "\n-----END PUBLIC KEY-----";
+
+    // 导出私钥为PEM格式
+    const privateKeyBuffer = await window.crypto.subtle.exportKey(
+        "pkcs8",
+        keyPair.privateKey
+    );
+    const privateKeyBase64 = btoa(
+        String.fromCharCode(...new Uint8Array(privateKeyBuffer))
+    );
+    const privateKeyPem =
+        "-----BEGIN PRIVATE KEY-----\n" +
+        privateKeyBase64.match(/.{1,64}/g)!.join("\n") +
+        "\n-----END PRIVATE KEY-----";
+
+    return {
+        publicKey: publicKeyPem,
+        privateKey: privateKeyPem,
+    };
+}
+
+// 获取或生成客户端密钥对
+async function getClientKeyPair(): Promise<KeyPair> {
+    // 从localStorage获取
+    const stored = localStorage.getItem(CLIENT_KEYPAIR_STORAGE_KEY);
+    if (stored) {
+        try {
+            return JSON.parse(stored);
+        } catch (e) {
+            console.error("Failed to parse stored client keypair:", e);
+        }
+    }
+
+    // 生成新的密钥对
+    console.log("Generating new client keypair...");
+    const keyPair = await generateClientKeyPair();
+
+    // 保存到localStorage
+    localStorage.setItem(CLIENT_KEYPAIR_STORAGE_KEY, JSON.stringify(keyPair));
+
+    return keyPair;
+}
+
+// 从PEM导入公钥
+async function importPublicKeyFromPem(pem: string): Promise<CryptoKey> {
+    const pemHeader = "-----BEGIN PUBLIC KEY-----";
+    const pemFooter = "-----END PUBLIC KEY-----";
+    const pemContents = pem
+        .substring(
+            pem.indexOf(pemHeader) + pemHeader.length,
+            pem.indexOf(pemFooter)
+        )
+        .replace(/\s/g, "");
+
+    const binaryDer = window.atob(pemContents);
+    const binaryDerArray = new Uint8Array(binaryDer.length);
+    for (let i = 0; i < binaryDer.length; i++) {
+        binaryDerArray[i] = binaryDer.charCodeAt(i);
+    }
+
+    return await window.crypto.subtle.importKey(
+        "spki",
+        binaryDerArray,
+        {
+            name: "RSA-OAEP",
+            hash: "SHA-256",
+        },
+        false,
+        ["encrypt"]
+    );
+}
+
+// 从PEM导入私钥
+async function importPrivateKeyFromPem(pem: string): Promise<CryptoKey> {
+    const pemHeader = "-----BEGIN PRIVATE KEY-----";
+    const pemFooter = "-----END PRIVATE KEY-----";
+    const pemContents = pem
+        .substring(
+            pem.indexOf(pemHeader) + pemHeader.length,
+            pem.indexOf(pemFooter)
+        )
+        .replace(/\s/g, "");
+
+    const binaryDer = window.atob(pemContents);
+    const binaryDerArray = new Uint8Array(binaryDer.length);
+    for (let i = 0; i < binaryDer.length; i++) {
+        binaryDerArray[i] = binaryDer.charCodeAt(i);
+    }
+
+    return await window.crypto.subtle.importKey(
+        "pkcs8",
+        binaryDerArray,
+        {
+            name: "RSA-OAEP",
+            hash: "SHA-256",
+        },
+        false,
+        ["decrypt"]
+    );
+}
+
+// 与服务器交换密钥
+async function exchangeKeys(): Promise<ServerKeys> {
+    // 检查缓存
+    const cached = localStorage.getItem(SERVER_KEYS_STORAGE_KEY);
+    if (cached) {
+        try {
+            return JSON.parse(cached);
+        } catch (e) {
+            console.error("Failed to parse stored server keys:", e);
+        }
+    }
+
+    // 获取客户端密钥对
+    const clientKeyPair = await getClientKeyPair();
+
+    // 发送客户端公钥到服务器
+    const response = await fetch("/cryptio/exchange_keys", {
+        method: "POST",
+        headers: {
+            "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+            client_public_key: clientKeyPair.publicKey,
+        }),
+    });
+
+    if (!response.ok) {
+        throw new Error(`Key exchange failed: ${response.statusText}`);
+    }
+
+    const data = await response.json();
+
+    // 使用客户端私钥解密服务端密钥
+    const clientPrivateKey = await importPrivateKeyFromPem(
+        clientKeyPair.privateKey
+    );
+
+    // 1. 解密AES密钥
+    const encryptedAesKey = Uint8Array.from(
+        atob(data.encrypted_aes_key),
+        (c) => c.charCodeAt(0)
+    );
+    const aesKeyBuffer = await window.crypto.subtle.decrypt(
+        {
+            name: "RSA-OAEP",
+        },
+        clientPrivateKey,
+        encryptedAesKey
+    );
+
+    // 2. 导入AES密钥
+    const aesKey = await window.crypto.subtle.importKey(
+        "raw",
+        aesKeyBuffer,
+        {
+            name: "AES-CBC",
+            length: 256,
+        },
+        false,
+        ["decrypt"]
+    );
+
+    // 3. 解密服务端密钥数据
+    const iv = Uint8Array.from(atob(data.iv), (c) => c.charCodeAt(0));
+    const encryptedData = Uint8Array.from(atob(data.encrypted_data), (c) =>
+        c.charCodeAt(0)
+    );
+
+    const decryptedBuffer = await window.crypto.subtle.decrypt(
+        {
+            name: "AES-CBC",
+            iv: iv,
+        },
+        aesKey,
+        encryptedData
+    );
+
+    // 4. 去除PKCS7 padding
+    const decryptedArray = new Uint8Array(decryptedBuffer);
+    const paddingLength = decryptedArray[decryptedArray.length - 1];
+    const unpaddedData = decryptedArray.slice(0, -paddingLength);
+
+    // 5. 解析JSON
+    const decryptedText = new TextDecoder().decode(unpaddedData);
+    const serverKeys = JSON.parse(decryptedText);
+
+    // 保存到localStorage
+    localStorage.setItem(SERVER_KEYS_STORAGE_KEY, JSON.stringify(serverKeys));
+
+    return serverKeys;
+}
+
+// 获取服务端密钥
+async function getServerKeys(): Promise<ServerKeys> {
+    return await exchangeKeys();
+}
+
+// 使用混合加密方案加密文件
+async function encryptFile(file: File): Promise<Uint8Array> {
+    const serverKeys = await getServerKeys();
+
+    // 读取文件内容
+    const arrayBuffer = await file.arrayBuffer();
+    const data = new Uint8Array(arrayBuffer);
+
+    // 1. 生成随机 AES 密钥
+    const aesKey = await window.crypto.subtle.generateKey(
+        {
+            name: "AES-GCM",
+            length: 256,
+        },
+        true,
+        ["encrypt"]
+    );
+
+    // 2. 生成随机 IV
+    const iv = window.crypto.getRandomValues(new Uint8Array(12));
+
+    // 3. 使用 AES 加密文件数据
+    const encryptedData = await window.crypto.subtle.encrypt(
+        {
+            name: "AES-GCM",
+            iv: iv,
+        },
+        aesKey,
+        data
+    );
+
+    // 4. 导出 AES 密钥
+    const rawAesKey = await window.crypto.subtle.exportKey("raw", aesKey);
+
+    // 5. 使用服务端公钥加密 AES 密钥
+    const serverPublicKey = await importPublicKeyFromPem(
+        serverKeys.publicKey
+    );
+    const encryptedAesKey = await window.crypto.subtle.encrypt(
+        {
+            name: "RSA-OAEP",
+        },
+        serverPublicKey,
+        rawAesKey
+    );
+
+    // 6. 组合数据
+    const encryptedAesKeyArray = new Uint8Array(encryptedAesKey);
+    const encryptedDataArray = new Uint8Array(encryptedData);
+    const combined = new Uint8Array(
+        4 + encryptedAesKeyArray.length + 4 + iv.length + encryptedDataArray.length
+    );
+
+    const view = new DataView(combined.buffer);
+    view.setUint32(0, encryptedAesKeyArray.length, false);
+    combined.set(encryptedAesKeyArray, 4);
+    view.setUint32(4 + encryptedAesKeyArray.length, iv.length, false);
+    combined.set(iv, 4 + encryptedAesKeyArray.length + 4);
+    combined.set(encryptedDataArray, 4 + encryptedAesKeyArray.length + 4 + iv.length);
+
+    return combined;
+}
+
+// 解密文件（用于预览）
+async function decryptFile(encryptedData: Uint8Array): Promise<Uint8Array> {
+    const serverKeys = await getServerKeys();
+
+    // 解析组合数据
+    let offset = 0;
+
+    // 1. 读取加密的AES密钥长度
+    const view = new DataView(encryptedData.buffer);
+    const encryptedAesKeyLength = view.getUint32(offset, false);
+    offset += 4;
+
+    // 2. 读取加密的AES密钥
+    const encryptedAesKey = encryptedData.slice(
+        offset,
+        offset + encryptedAesKeyLength
+    );
+    offset += encryptedAesKeyLength;
+
+    // 3. 读取IV长度
+    const ivLength = view.getUint32(offset, false);
+    offset += 4;
+
+    // 4. 读取IV
+    const iv = encryptedData.slice(offset, offset + ivLength);
+    offset += ivLength;
+
+    // 5. 读取加密的数据
+    const encrypted = encryptedData.slice(offset);
+
+    // 6. 使用服务端私钥解密AES密钥
+    const serverPrivateKey = await importPrivateKeyFromPem(
+        serverKeys.privateKey
+    );
+    const aesKeyBuffer = await window.crypto.subtle.decrypt(
+        {
+            name: "RSA-OAEP",
+        },
+        serverPrivateKey,
+        encryptedAesKey
+    );
+
+    // 7. 导入AES密钥
+    const aesKey = await window.crypto.subtle.importKey(
+        "raw",
+        aesKeyBuffer,
+        {
+            name: "AES-GCM",
+            length: 256,
+        },
+        false,
+        ["decrypt"]
+    );
+
+    // 8. 使用AES密钥解密数据
+    const decryptedBuffer = await window.crypto.subtle.decrypt(
+        {
+            name: "AES-GCM",
+            iv: iv,
+        },
+        aesKey,
+        encrypted
+    );
+
+    return new Uint8Array(decryptedBuffer);
+}
+
+// 自定义上传函数
+async function uploadEncryptedImage(file: File): Promise<any> {
+    console.log("Encrypting and uploading image:", file.name);
+
+    // 加密文件
+    const encryptedData = await encryptFile(file);
+
+    // 创建FormData
+    const formData = new FormData();
+    const encryptedBlob = new Blob([encryptedData], {
+        type: "application/octet-stream",
+    });
+    formData.append("image", encryptedBlob, file.name);
+
+    // 上传到服务器
+    const response = await fetch("/cryptio/upload_encrypted", {
+        method: "POST",
+        body: formData,
+    });
+
+    if (!response.ok) {
+        throw new Error(`Upload failed: ${response.statusText}`);
+    }
+
+    return await response.json();
+}
+
+// 加载并解密图片用于预览
+async function loadEncryptedImageForPreview(filename: string): Promise<string> {
+    // 从服务器获取加密数据
+    const response = await fetch(
+        `/cryptio/view_encrypted?filename=${encodeURIComponent(filename)}`
+    );
+
+    if (!response.ok) {
+        throw new Error(`Failed to load encrypted image: ${response.statusText}`);
+    }
+
+    const data = await response.json();
+
+    // 解密数据
+    const encryptedData = Uint8Array.from(atob(data.encrypted_data), (c) =>
+        c.charCodeAt(0)
+    );
+    const decryptedData = await decryptFile(encryptedData);
+
+    // 转换为Blob URL
+    const blob = new Blob([decryptedData]);
+    return URL.createObjectURL(blob);
+}
+
+// 注册扩展
+app.registerExtension({
+    name: "cryptio.UploadImageCryptIO",
+    async setup() {
+        // 初始化时交换密钥
+        try {
+            await exchangeKeys();
+            console.log("CryptIO: Keys exchanged successfully");
+        } catch (error) {
+            console.error("CryptIO: Failed to exchange keys:", error);
+        }
+    },
+    async beforeRegisterNodeDef(
+        nodeType: any,
+        nodeData: ComfyNodeDef,
+        app: ComfyApp
+    ) {
+        if (nodeType.comfyClass === "UploadImageCryptIO") {
+            // 添加自定义的upload widget
+            const onNodeCreated = nodeType.prototype.onNodeCreated;
+            nodeType.prototype.onNodeCreated = function (this: any) {
+                const r = onNodeCreated?.apply(this, arguments);
+
+                // 添加上传button widget
+                const uploadWidget = this.addWidget(
+                    "button",
+                    "choose file to upload (cryptIO)",
+                    "image",
+                    () => {
+                        const fileInput = document.createElement("input");
+                        fileInput.type = "file";
+                        fileInput.accept = "image/*";
+                        fileInput.style.display = "none";
+                        document.body.appendChild(fileInput);
+
+                        fileInput.onchange = async () => {
+                            if (fileInput.files && fileInput.files.length > 0) {
+                                const file = fileInput.files[0];
+                                try {
+                                    // 上传加密图片
+                                    const result = await uploadEncryptedImage(file);
+
+                                    // 更新image widget的值
+                                    const imageWidget = this.widgets.find(
+                                        (w: any) => w.name === "image"
+                                    );
+                                    if (imageWidget) {
+                                        imageWidget.value = result.name;
+                                    }
+
+                                    // 更新encrypted标志
+                                    const encryptedWidget = this.widgets.find(
+                                        (w: any) => w.name === "encrypted"
+                                    );
+                                    if (encryptedWidget) {
+                                        encryptedWidget.value = true;
+                                    }
+
+                                    // 更新预览
+                                    await this.updatePreview(result.name);
+
+                                    app.graph?.setDirtyCanvas(true, false);
+                                } catch (error) {
+                                    console.error("Upload error:", error);
+                                    alert("Failed to upload encrypted image: " + error);
+                                }
+                            }
+                            document.body.removeChild(fileInput);
+                        };
+
+                        fileInput.click();
+                    }
+                );
+
+                // 添加预览更新方法
+                this.updatePreview = async function (filename: string) {
+                    if (!filename || !filename.endsWith(".encrypted")) {
+                        return;
+                    }
+
+                    try {
+                        const imageUrl = await loadEncryptedImageForPreview(filename);
+
+                        // 更新节点的图片显示
+                        const img = new Image();
+                        img.onload = () => {
+                            this.imgs = [img];
+                            app.graph?.setDirtyCanvas(true, false);
+                        };
+                        img.src = imageUrl;
+                    } catch (error) {
+                        console.error("Preview error:", error);
+                    }
+                };
+
+                // 当image值改变时更新预览
+                const imageWidget = this.widgets.find((w: any) => w.name === "image");
+                if (imageWidget) {
+                    const originalCallback = imageWidget.callback;
+                    imageWidget.callback = async function (this: any, value: any) {
+                        if (originalCallback) {
+                            originalCallback.call(this, value);
+                        }
+                        if (value && value.endsWith(".encrypted")) {
+                            await this.node.updatePreview(value);
+                        }
+                    }.bind({ node: this });
+                }
+
+                return r;
+            };
+        }
+    },
+});