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
|
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'
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue
Block a user