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
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):
key = self.headers.get('Sec-WebSocket-Key', '')
accept = base64.b64encode(
@ -706,6 +720,8 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
self.send_header('Connection', 'Upgrade')
self.send_header('Sec-WebSocket-Accept', accept)
self.end_headers()
client = f'{self.client_address[0]}:{self.client_address[1]}'
print(f'[WS] {client} connected')
sock = self.connection
sock.settimeout(None)
try:
@ -723,21 +739,20 @@ class DAVRequestHandler(BaseHTTPRequestHandler):
if masked:
payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
if opcode == 0x8: # close
print(f'[WS] {client} closed')
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)
sock.send(bytes([0x8A, len(payload)]) + payload)
elif opcode in (0x1, 0x2): # text or binary
message = payload.decode('utf-8', errors='replace')
print(f'[WS] {client} recv ({length}b): {message}')
response = self.ws_process_message(message)
if response is not None:
self._ws_send(sock, response.encode('utf-8'))
except Exception:
pass
print(f'[WS] {client} disconnected')
def do_GET(self, onlyhead=False):
if (self.path == '/ws'