Decrypt-proxied-gradio-payloads

app.py

@@ -40,6 +40,7 @@ import numpy as np
 import py7zr
 import spaces
 import torch
+from cryptography.exceptions import InvalidTag
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
 from fastapi import Body, Depends, HTTPException, Query
@@ -100,6 +101,10 @@ SALT_SIZE = 16
 NONCE_SIZE = 12
 SCRYPT_N = 2**14
 FILE_MAGIC = b"EPNG1"
+PROXY_ACCESS_VALUE = b"haip"
+PROXY_ENCRYPTION = b"aes-256-gcm"
+PROXY_ENCRYPTION_HEADER = b"x-gradio-comfy-encryption"
+PROXY_MAGIC = b"GCV1"
 IMAGE_SUFFIXES = (".epng",)
 PREVIEW_CACHE_SIZE = 256
 PREVIEW_QUALITY = 70
@@ -1874,6 +1879,42 @@ def image_key(salt):
     ).derive(PASSWORD.encode())
 
 
+def proxy_key(salt):
+    return Scrypt(
+        salt=salt,
+        length=32,
+        n=SCRYPT_N,
+        r=8,
+        p=1,
+    ).derive(PROXY_ACCESS_VALUE)
+
+
+def encrypt_proxy_payload(data):
+    salt = os.urandom(SALT_SIZE)
+    nonce = os.urandom(NONCE_SIZE)
+    return (
+        PROXY_MAGIC
+        + salt
+        + nonce
+        + AESGCM(proxy_key(salt)).encrypt(nonce, data, PROXY_MAGIC)
+    )
+
+
+def decrypt_proxy_payload(data):
+    if len(data) < len(PROXY_MAGIC) + SALT_SIZE + NONCE_SIZE + 16:
+        raise ValueError("invalid encrypted payload")
+    if not data.startswith(PROXY_MAGIC):
+        raise ValueError("invalid encrypted payload")
+    salt_start = len(PROXY_MAGIC)
+    nonce_start = salt_start + SALT_SIZE
+    data_start = nonce_start + NONCE_SIZE
+    return AESGCM(proxy_key(data[salt_start:nonce_start])).decrypt(
+        data[nonce_start:data_start],
+        data[data_start:],
+        PROXY_MAGIC,
+    )
+
+
 def save_image(image, from_api):
     path = IMAGE_DIR / datetime.now(TIMEZONE).date().isoformat()
     path.mkdir(parents=True, exist_ok=True)
@@ -2759,7 +2800,107 @@ def run_job(job_id, workflow):
         }
 
 
+def replace_asgi_headers(headers, replacements):
+    names = {name for name, _ in replacements}
+    return [item for item in headers if item[0].lower() not in names] + replacements
+
+
+class GradioEncryptionMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+        headers = dict(scope.get("headers", []))
+        encrypted = (
+            headers.get(PROXY_ENCRYPTION_HEADER) == PROXY_ENCRYPTION
+            and "/gradio_api/call/health" not in scope.get("path", "")
+        )
+        if not encrypted:
+            await self.app(scope, receive, send)
+            return
+
+        decrypted_receive = receive
+        if scope.get("method") not in {"GET", "HEAD"}:
+            chunks = []
+            while True:
+                message = await receive()
+                if message["type"] == "http.disconnect":
+                    return
+                chunks.append(message.get("body", b""))
+                if not message.get("more_body", False):
+                    break
+            try:
+                body = decrypt_proxy_payload(b"".join(chunks))
+            except (InvalidTag, ValueError):
+                content = b'{"detail":"Invalid encrypted payload"}'
+                await send({
+                    "type": "http.response.start",
+                    "status": 400,
+                    "headers": [
+                        (b"content-type", b"application/json"),
+                        (b"content-length", str(len(content)).encode()),
+                    ],
+                })
+                await send({"type": "http.response.body", "body": content})
+                return
+
+            scope = dict(scope)
+            scope["headers"] = replace_asgi_headers(
+                scope.get("headers", []),
+                [
+                    (b"content-type", b"application/json"),
+                    (b"content-length", str(len(body)).encode()),
+                ],
+            )
+            delivered = False
+
+            async def decrypted_receive():
+                nonlocal delivered
+                if delivered:
+                    return {"type": "http.request", "body": b"", "more_body": False}
+                delivered = True
+                return {"type": "http.request", "body": body, "more_body": False}
+
+        start = None
+        response_chunks = []
+
+        async def encrypted_send(message):
+            nonlocal start
+            if message["type"] == "http.response.start":
+                start = message
+                return
+            if message["type"] == "http.response.pathsend":
+                await send(start)
+                await send(message)
+                return
+            if message["type"] != "http.response.body":
+                await send(message)
+                return
+            response_chunks.append(message.get("body", b""))
+            if message.get("more_body", False):
+                return
+            content = b"".join(response_chunks)
+            if not content.startswith(FILE_MAGIC):
+                content = encrypt_proxy_payload(content)
+                start = dict(start)
+                start["headers"] = replace_asgi_headers(
+                    start.get("headers", []),
+                    [
+                        (PROXY_ENCRYPTION_HEADER, PROXY_ENCRYPTION),
+                        (b"content-length", str(len(content)).encode()),
+                    ],
+                )
+            await send(start)
+            await send({"type": "http.response.body", "body": content})
+
+        await self.app(scope, decrypted_receive, encrypted_send)
+
+
 api = App()
+api.add_middleware(GradioEncryptionMiddleware)
 api.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],