feat(WebSocket): enhance WebSocket message handling and connection logging
This commit is contained in:
parent
ff664fa9d3
commit
ab5d9bb486
35
webdav3.py
35
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'
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user