feat(WebSocket): implement WebSocket handling and upgrade support

This commit is contained in:
Jaime Idolpx 2026-06-08 14:08:37 -04:00
parent e2197c33fd
commit ff664fa9d3

View File

@ -43,7 +43,7 @@ from http.server import HTTPServer, BaseHTTPRequestHandler
from socketserver import ThreadingMixIn
import urllib.request, urllib.parse, urllib.error, urllib.parse
from time import timezone, strftime, localtime, gmtime
import os, sys, re, shutil, uuid, hashlib, mimetypes, base64, socket
import os, sys, re, shutil, struct, uuid, hashlib, mimetypes, base64, socket
# Debug message ( True / False )
sys_debug = False
@ -687,7 +687,63 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
self.end_headers()
w.flush()
def _ws_recv(self, sock, n):
data = b''
while len(data) < n:
chunk = sock.recv(n - len(data))
if not chunk:
raise ConnectionError('closed')
data += chunk
return data
def _handle_websocket(self):
key = self.headers.get('Sec-WebSocket-Key', '')
accept = base64.b64encode(
hashlib.sha1((key + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11').encode()).digest()
).decode()
self.send_response(101, 'Switching Protocols')
self.send_header('Upgrade', 'websocket')
self.send_header('Connection', 'Upgrade')
self.send_header('Sec-WebSocket-Accept', accept)
self.end_headers()
sock = self.connection
sock.settimeout(None)
try:
while True:
b0, b1 = self._ws_recv(sock, 2)
opcode = b0 & 0x0F
masked = bool(b1 & 0x80)
length = b1 & 0x7F
if length == 126:
length = struct.unpack('>H', self._ws_recv(sock, 2))[0]
elif length == 127:
length = struct.unpack('>Q', self._ws_recv(sock, 8))[0]
mask = self._ws_recv(sock, 4) if masked else b''
payload = self._ws_recv(sock, length)
if masked:
payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
if opcode == 0x8: # close
sock.send(b'\x88\x00')
break
elif opcode == 0x9: # ping → pong
frame = bytes([0x8A, len(payload)]) + payload
sock.send(frame)
elif opcode in (0x1, 0x2): # text or binary — echo back as text
if length < 126:
header = bytes([0x81, length])
elif length < 65536:
header = bytes([0x81, 126]) + struct.pack('>H', length)
else:
header = bytes([0x81, 127]) + struct.pack('>Q', length)
sock.send(header + payload)
except Exception:
pass
def do_GET(self, onlyhead=False):
if (self.path == '/ws'
and self.headers.get('Upgrade', '').lower() == 'websocket'):
self._handle_websocket()
return
if self.WebAuth():
return
path, elem = self.path_elem()