feat(WebSocket): enhance WebSocket message handling and connection logging

This commit is contained in:
Jaime Idolpx 2026-06-08 14:19:17 -04:00
parent ff664fa9d3
commit ab5d9bb486

View File

@ -696,6 +696,20 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
data += chunk data += chunk
return data return data
def ws_process_message(self, message: str) -> str:
"""Override this to handle incoming WebSocket messages. Return the response string."""
return message # echo by default
def _ws_send(self, sock, payload: bytes):
length = len(payload)
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)
def _handle_websocket(self): def _handle_websocket(self):
key = self.headers.get('Sec-WebSocket-Key', '') key = self.headers.get('Sec-WebSocket-Key', '')
accept = base64.b64encode( accept = base64.b64encode(
@ -706,6 +720,8 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
self.send_header('Connection', 'Upgrade') self.send_header('Connection', 'Upgrade')
self.send_header('Sec-WebSocket-Accept', accept) self.send_header('Sec-WebSocket-Accept', accept)
self.end_headers() self.end_headers()
client = f'{self.client_address[0]}:{self.client_address[1]}'
print(f'[WS] {client} connected')
sock = self.connection sock = self.connection
sock.settimeout(None) sock.settimeout(None)
try: try:
@ -723,21 +739,20 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
if masked: if masked:
payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload)) payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
if opcode == 0x8: # close if opcode == 0x8: # close
print(f'[WS] {client} closed')
sock.send(b'\x88\x00') sock.send(b'\x88\x00')
break break
elif opcode == 0x9: # ping → pong elif opcode == 0x9: # ping → pong
frame = bytes([0x8A, len(payload)]) + payload sock.send(bytes([0x8A, len(payload)]) + payload)
sock.send(frame) elif opcode in (0x1, 0x2): # text or binary
elif opcode in (0x1, 0x2): # text or binary — echo back as text message = payload.decode('utf-8', errors='replace')
if length < 126: print(f'[WS] {client} recv ({length}b): {message}')
header = bytes([0x81, length]) response = self.ws_process_message(message)
elif length < 65536: if response is not None:
header = bytes([0x81, 126]) + struct.pack('>H', length) self._ws_send(sock, response.encode('utf-8'))
else:
header = bytes([0x81, 127]) + struct.pack('>Q', length)
sock.send(header + payload)
except Exception: except Exception:
pass pass
print(f'[WS] {client} disconnected')
def do_GET(self, onlyhead=False): def do_GET(self, onlyhead=False):
if (self.path == '/ws' if (self.path == '/ws'