Files
caissonfiv-comfyui-https-br…/__init__.py
T

312 lines
12 KiB
Python

"""
ComfyUI HTTPS Bridge
====================
在 ComfyUI 启动时自动开启一个本地 HTTPS 代理服务器,
让部署在公网 HTTPS 页面的套壳应用可以安全连接到本地 ComfyUI。
安装:将本文件夹放入 ComfyUI/custom_nodes/ 目录,重启 ComfyUI 即可。
地址检测优先级:
1. 环境变量 COMFY_HOST / COMFY_PORT(手动覆盖,优先级最高)
2. ComfyUI 启动参数 --listen / --port(自动读取)
3. 默认值 127.0.0.1:8188
环境变量(可选):
COMFY_BRIDGE_PORT 桥接端口,默认 9999
COMFY_HOST 覆盖 ComfyUI 地址
COMFY_PORT 覆盖 ComfyUI 端口
"""
import asyncio
import ssl
import os
import socket
import threading
import ipaddress
import datetime
import pathlib
import traceback
# ── 配置 ──────────────────────────────────────────────────────────────────────
BRIDGE_PORT = int(os.environ.get('COMFY_BRIDGE_PORT', 9999))
def _get_local_ip():
"""获取本机对外的实际 IP 地址(非 127.0.0.1)。"""
try:
# 连接一个外部地址(不会真正发包),借此确定本机出口 IP
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.connect(('8.8.8.8', 80))
ip = s.getsockname()[0]
s.close()
return ip
except Exception:
return '127.0.0.1'
def _detect_comfy_address():
"""
按优先级检测 ComfyUI 实际监听地址:
1. 环境变量 COMFY_HOST / COMFY_PORT
2. ComfyUI 启动参数 --listen / --port
3. 默认 127.0.0.1:8188
"""
host, port = '127.0.0.1', 8188
# 从 ComfyUI 内部读取启动参数
try:
from comfy.cli_args import args
if args.port:
port = int(args.port)
if args.listen:
listen = args.listen.strip()
# 0.0.0.0 / :: 表示监听所有接口,取本机实际出口 IP
host = _get_local_ip() if listen in ('0.0.0.0', '::') else listen
except Exception:
pass # 不在 ComfyUI 环境中也能降级运行
# 环境变量最终覆盖(优先级最高)
host = os.environ.get('COMFY_HOST', host)
port = int(os.environ.get('COMFY_PORT', port))
return host, port
COMFY_HOST, COMFY_PORT = _detect_comfy_address()
CERT_DIR = pathlib.Path(__file__).parent / '.certs'
CORS_HEADERS = {
'Access-Control-Allow-Origin': '*',
'Access-Control-Allow-Methods': '*',
'Access-Control-Allow-Headers': '*',
}
# ── 证书生成 ──────────────────────────────────────────────────────────────────
def generate_cert():
"""使用 cryptography 库生成自签名证书,无需任何外部工具。"""
try:
from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
except ImportError:
print('[HTTPS Bridge] ❌ 缺少 cryptography 库,请在 ComfyUI 的 Python 环境中运行:')
print(' pip install cryptography')
return None, None
cert_file = CERT_DIR / 'cert.pem'
key_file = CERT_DIR / 'key.pem'
if cert_file.exists() and key_file.exists():
return str(cert_file), str(key_file)
CERT_DIR.mkdir(exist_ok=True)
print('[HTTPS Bridge] 🔐 首次运行,生成自签名证书(有效期 10 年)...')
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
subject = issuer = x509.Name([
x509.NameAttribute(NameOID.COMMON_NAME, 'localhost'),
])
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(issuer)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(datetime.datetime.utcnow())
.not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
.add_extension(
x509.SubjectAlternativeName([
x509.DNSName('localhost'),
x509.IPAddress(ipaddress.IPv4Address('127.0.0.1')),
]),
critical=False,
)
.sign(key, hashes.SHA256())
)
key_file.write_bytes(key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.TraditionalOpenSSL,
serialization.NoEncryption(),
))
cert_file.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
print(f'[HTTPS Bridge] ✅ 证书已生成: {CERT_DIR}')
return str(cert_file), str(key_file)
# ── 请求处理 ──────────────────────────────────────────────────────────────────
def _is_websocket(request):
return (
request.headers.get('Upgrade', '').lower() == 'websocket'
and 'upgrade' in request.headers.get('Connection', '').lower()
)
def _proxy_headers(headers, skip_keys):
return {k: v for k, v in headers.items() if k.lower() not in skip_keys}
async def handle_websocket(request):
"""双向代理 WebSocket 连接。"""
import aiohttp
from aiohttp import web
ws_server = web.WebSocketResponse()
await ws_server.prepare(request)
target = f'ws://{COMFY_HOST}:{COMFY_PORT}{request.path_qs}'
skip = {'host', 'upgrade', 'connection',
'sec-websocket-key', 'sec-websocket-version', 'sec-websocket-extensions'}
headers = _proxy_headers(request.headers, skip)
try:
async with aiohttp.ClientSession() as session:
async with session.ws_connect(target, headers=headers) as ws_client:
async def forward(src, dst, is_text_fn, send_str, send_bytes):
async for msg in src:
if dst.closed:
break
if msg.type == aiohttp.WSMsgType.TEXT:
await send_str(msg.data)
elif msg.type == aiohttp.WSMsgType.BINARY:
await send_bytes(msg.data)
elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
break
await asyncio.gather(
forward(ws_client, ws_server, None,
ws_server.send_str, ws_server.send_bytes),
forward(ws_server, ws_client, None,
ws_client.send_str, ws_client.send_bytes),
return_exceptions=True,
)
except Exception as e:
print(f'[HTTPS Bridge] WS 错误: {e}')
return ws_server
async def handle_http(request):
"""代理普通 HTTP 请求,支持流式响应(大文件/图片)。"""
import aiohttp
from aiohttp import web
if request.method == 'OPTIONS':
return web.Response(status=204, headers=CORS_HEADERS)
target = f'http://{COMFY_HOST}:{COMFY_PORT}{request.path_qs}'
skip = {'host', 'content-length', 'transfer-encoding'}
headers = _proxy_headers(request.headers, skip)
try:
async with aiohttp.ClientSession() as session:
async with session.request(
request.method, target,
headers=headers,
data=await request.read(),
allow_redirects=False,
) as resp:
resp_headers = {
k: v for k, v in resp.headers.items()
if k.lower() not in {'transfer-encoding'}
}
resp_headers.update(CORS_HEADERS)
response = web.StreamResponse(status=resp.status, headers=resp_headers)
await response.prepare(request)
async for chunk in resp.content.iter_chunked(65536):
await response.write(chunk)
await response.write_eof()
return response
except aiohttp.ClientConnectorError:
return web.Response(
status=502,
text=f'无法连接到 ComfyUI ({COMFY_HOST}:{COMFY_PORT}),请确认 ComfyUI 正在运行。',
headers=CORS_HEADERS,
)
except Exception as e:
return web.Response(status=500, text=f'代理错误: {e}', headers=CORS_HEADERS)
async def handle_request(request):
if _is_websocket(request):
return await handle_websocket(request)
return await handle_http(request)
# ── 服务器启动 ────────────────────────────────────────────────────────────────
async def _run_server(cert_file, key_file):
from aiohttp import web
ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ssl_ctx.load_cert_chain(cert_file, key_file)
app = web.Application()
app.router.add_route('*', '/', handle_request)
app.router.add_route('*', '/{path_info:.*}', handle_request)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '0.0.0.0', BRIDGE_PORT, ssl_context=ssl_ctx)
await site.start()
print(f'\n✅ [HTTPS Bridge] 已启动')
print(f' 桥接地址: https://localhost:{BRIDGE_PORT}')
print(f' 代理目标: http://{COMFY_HOST}:{COMFY_PORT}')
print(f'\n⚠️ 首次使用需在浏览器授权证书(只需一次):')
print(f' 1. 打开 https://localhost:{BRIDGE_PORT}')
print(f' 2. 点击「高级」→「继续访问 localhost」')
print(f' 3. 看到空白页即表示成功\n')
await asyncio.Event().wait() # 永久运行
def _start_in_thread(cert_file, key_file):
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(_run_server(cert_file, key_file))
except OSError as e:
if e.errno == 98 or 'address already in use' in str(e).lower():
print(f'[HTTPS Bridge] ⚠️ 端口 {BRIDGE_PORT} 已被占用,桥接未启动。')
print(f' 可通过环境变量更换端口: COMFY_BRIDGE_PORT=9998')
else:
print(f'[HTTPS Bridge] 启动失败: {e}')
traceback.print_exc()
thread = threading.Thread(target=_run, daemon=True, name='https-bridge')
thread.start()
# ── ComfyUI Node 注册 ─────────────────────────────────────────────────────────
class HttpsBridgeNode:
"""显示桥接状态的信息节点(可选使用,桥接服务在后台自动运行)。"""
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('STRING',)
FUNCTION = 'get_status'
CATEGORY = '🔒 HTTPS Bridge'
OUTPUT_NODE = True
def get_status(self):
return (f'HTTPS Bridge: https://localhost:{BRIDGE_PORT}',)
NODE_CLASS_MAPPINGS = {'HttpsBridgeNode': HttpsBridgeNode}
NODE_DISPLAY_NAME_MAPPINGS = {'HttpsBridgeNode': '🔒 HTTPS Bridge Info'}
# ── 入口:ComfyUI 加载时自动执行 ─────────────────────────────────────────────
cert_file, key_file = generate_cert()
if cert_file:
_start_in_thread(cert_file, key_file)