62 lines
3.3 KiB
Python
62 lines
3.3 KiB
Python
"""Direct authenticated transport to Deck's own ComfyUI, including WebSocket."""
|
|
import http.client
|
|
import socket
|
|
import select
|
|
from pathlib import Path
|
|
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
|
|
import os
|
|
MAX_BODY=256*1024*1024
|
|
REQUEST_HEADERS={'content-type','accept','range','if-range','if-none-match','if-modified-since'}
|
|
HOP_HEADERS={'connection','keep-alive','proxy-authenticate','proxy-authorization','te','trailer','transfer-encoding','upgrade'}
|
|
|
|
def relay(handler,port):
|
|
if handler.headers.get('Transfer-Encoding'):raise ValueError('Content-Length erforderlich.')
|
|
length=int(handler.headers.get('Content-Length','0'))
|
|
if not 0<=length<=MAX_BODY:raise ValueError('Maximal 256 MiB pro Upload.')
|
|
if handler.headers.get('Upgrade','').lower()=='websocket':
|
|
with socket.create_connection(('127.0.0.1',port),timeout=10) as upstream:
|
|
lines=[f'{handler.command} {handler.path} HTTP/1.1',f'Host: 127.0.0.1:{port}','Authorization: Bearer '+Path('/run/secrets/comfy-token').read_text().strip()]
|
|
for key in ('Upgrade','Connection','Sec-WebSocket-Key','Sec-WebSocket-Version','Sec-WebSocket-Protocol'):
|
|
if handler.headers.get(key):lines.append(key+': '+handler.headers[key])
|
|
upstream.sendall(('\r\n'.join(lines)+'\r\n\r\n').encode());head=b''
|
|
while b'\r\n\r\n' not in head:
|
|
chunk=upstream.recv(4096)
|
|
if not chunk:raise ValueError('WebSocket-Verbindung abgebrochen.')
|
|
head+=chunk
|
|
if len(head)>65536:raise ValueError('Ungültige WebSocket-Antwort.')
|
|
handler.sent=True;handler.connection.sendall(head);upstream.settimeout(None)
|
|
while True:
|
|
readable,_,_=select.select([upstream,handler.connection],[],[],60)
|
|
if not readable:continue
|
|
for source in readable:
|
|
data=source.recv(65536)
|
|
if not data:return
|
|
(handler.connection if source is upstream else upstream).sendall(data)
|
|
else:
|
|
upstream=http.client.HTTPConnection('127.0.0.1',port,timeout=3600)
|
|
try:
|
|
headers={k:v for k,v in handler.headers.items() if k.lower() in REQUEST_HEADERS};headers['Content-Length']=str(length);headers['Authorization']='Bearer '+Path('/run/secrets/comfy-token').read_text().strip()
|
|
upstream.putrequest(handler.command,handler.path)
|
|
for key,value in headers.items():upstream.putheader(key,value)
|
|
upstream.endheaders();remaining=length
|
|
while remaining:
|
|
chunk=handler.rfile.read(min(1024*1024,remaining))
|
|
if not chunk:raise ValueError('Unvollständiger Upload.')
|
|
upstream.send(chunk);remaining-=len(chunk)
|
|
response=upstream.getresponse();handler.sent=True;handler.send_response(response.status)
|
|
for key,value in response.getheaders():
|
|
if key.lower() not in HOP_HEADERS|{'server','date'}:handler.send_header(key,value)
|
|
handler.send_header('Connection','close');handler.end_headers();handler.close_connection=True
|
|
if handler.command!='HEAD':
|
|
while chunk:=response.read(1024*1024):handler.wfile.write(chunk)
|
|
finally:upstream.close()
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self,*args):pass
|
|
def route(self):
|
|
try:relay(self,int(os.environ.get('SWARM_DECK_PORT','8120')))
|
|
except (OSError,ValueError):
|
|
if not getattr(self,'sent',False):self.send_error(503,'Deck video backend unavailable')
|
|
do_GET=do_POST=do_PUT=do_PATCH=do_DELETE=do_HEAD=route
|
|
if __name__=='__main__':ThreadingHTTPServer(('127.0.0.1',int(os.environ.get('SWARM_PROXY_PORT','8126'))),Handler).serve_forever()
|