From ab5d9bb48601d887f1e8ea9870e56e3180911d6b Mon Sep 17 00:00:00 2001
From: Jaime Idolpx <jaime@idolpx.com>
Date: Mon, 8 Jun 2026 14:19:17 -0400
Subject: [PATCH] feat(WebSocket): enhance WebSocket message handling and
 connection logging

---
 webdav3.py | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/webdav3.py b/webdav3.py
index 80603dc..f0b390b 100644
--- a/webdav3.py
+++ b/webdav3.py
@@ -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'