Add WebSocket transport (ws:// + wss://) for Remote Desktop

A new MessageChannel abstraction lets the host and viewer speak the
existing typed-message protocol over either raw TCP framing or
WebSocket BINARY frames. Each WS frame carries one full encoded typed
message (magic + type + length + payload), so decode_frame_header /
encode_frame are reused unchanged and only the wire layer changes.

ws_protocol.py is a small RFC 6455 implementation (no extra deps):
server / client handshake helpers, single-frame BINARY send, recv that
transparently handles PING / PONG / CLOSE control frames, and explicit
rejection of fragmented data frames so messages always fit in one
~16 MiB frame. Clients mask outgoing payloads as required; servers do
not.

WebSocketDesktopHost and WebSocketDesktopViewer are thin subclasses
that override the channel-creation hook to perform the upgrade
handshake before falling back to the shared auth + receive loop. The
existing ssl_context plumbing stays in place — passing a context to
WebSocketDesktopHost/Viewer transparently upgrades the connection to
wss://, so no separate TLS-WS class is needed.

Tests cover ws_protocol round trips (handshake, masked + unmasked
binary frames, extended payload length, bad-request rejection) and
end-to-end host<->viewer scenarios (auth, frame stream, input
dispatch, host_id announce, mixed-transport rejection in both
directions, path validation).

Author: JE-Chen

je_auto_control/utils/remote_desktop/__init__.py

@@ -23,9 +23,14 @@
 )
 from je_auto_control.utils.remote_desktop.registry import registry
 from je_auto_control.utils.remote_desktop.viewer import RemoteDesktopViewer
+from je_auto_control.utils.remote_desktop.ws_host import WebSocketDesktopHost
+from je_auto_control.utils.remote_desktop.ws_viewer import (
+    WebSocketDesktopViewer,
+)
 
 __all__ = [
     "RemoteDesktopHost", "RemoteDesktopViewer",
+    "WebSocketDesktopHost", "WebSocketDesktopViewer",
     "InputDispatchError", "AuthenticationError", "ProtocolError",
     "MessageType", "encode_frame", "decode_frame_header",
     "dispatch_input", "registry",

je_auto_control/utils/remote_desktop/host.py

@@ -19,7 +19,9 @@
 )
 from je_auto_control.utils.remote_desktop.protocol import (
     AuthenticationError, MessageType, ProtocolError,
-    encode_frame, read_message,
+)
+from je_auto_control.utils.remote_desktop.transport import (
+    MessageChannel, TcpMessageChannel,
 )
 
 FrameProvider = Callable[[], bytes]
@@ -51,12 +53,11 @@ def provide() -> bytes:
 class _ClientHandler:
     """Per-connection auth + input-receive + frame-send state."""
 
-    def __init__(self, host: "RemoteDesktopHost", sock: socket.socket,
-                 address) -> None:
+    def __init__(self, host: "RemoteDesktopHost",
+                 channel: MessageChannel, address) -> None:
         self._host = host
-        self._sock = sock
+        self._channel = channel
         self._address = address
-        self._send_lock = threading.Lock()
         self._shutdown = threading.Event()
         self._sender_thread: Optional[threading.Thread] = None
         self._receiver_thread: Optional[threading.Thread] = None
@@ -95,27 +96,23 @@ def stop(self) -> None:
 
     def _authenticate(self) -> None:
         nonce = make_nonce()
-        self._sock.settimeout(_AUTH_TIMEOUT_S)
-        self._send(MessageType.AUTH_CHALLENGE, nonce)
-        msg_type, payload = read_message(self._sock)
+        self._channel.settimeout(_AUTH_TIMEOUT_S)
+        self._channel.send_typed(MessageType.AUTH_CHALLENGE, nonce)
+        msg_type, payload = self._channel.read_typed()
         if msg_type is not MessageType.AUTH_RESPONSE:
-            self._send(MessageType.AUTH_FAIL, b"expected AUTH_RESPONSE")
+            self._channel.send_typed(MessageType.AUTH_FAIL,
+                                     b"expected AUTH_RESPONSE")
             raise AuthenticationError(
                 f"expected AUTH_RESPONSE, got {msg_type.name}"
             )
         if not verify_response(self._host._token, nonce, payload):
-            self._send(MessageType.AUTH_FAIL, b"bad token")
+            self._channel.send_typed(MessageType.AUTH_FAIL, b"bad token")
             raise AuthenticationError("bad token")
         ok_payload = json.dumps(
             {"host_id": self._host.host_id}, ensure_ascii=False,
         ).encode("utf-8")
-        self._send(MessageType.AUTH_OK, ok_payload)
-        self._sock.settimeout(None)
-
-    def _send(self, message_type: MessageType, payload: bytes) -> None:
-        data = encode_frame(message_type, payload)
-        with self._send_lock:
-            self._sock.sendall(data)
+        self._channel.send_typed(MessageType.AUTH_OK, ok_payload)
+        self._channel.settimeout(None)
 
     def _send_loop(self) -> None:
         last_sent = 0
@@ -131,7 +128,7 @@ def _send_loop(self) -> None:
             if frame is None:
                 continue
             try:
-                self._send(MessageType.FRAME, frame)
+                self._channel.send_typed(MessageType.FRAME, frame)
             except (OSError, ConnectionError) as error:
                 autocontrol_logger.info(
                     "remote_desktop send to %s failed: %r",
@@ -144,7 +141,7 @@ def _send_loop(self) -> None:
     def _recv_loop(self) -> None:
         while not self._shutdown.is_set():
             try:
-                msg_type, payload = read_message(self._sock)
+                msg_type, payload = self._channel.read_typed()
             except (OSError, ConnectionError, ProtocolError) as error:
                 if not self._shutdown.is_set():
                     autocontrol_logger.info(
@@ -186,14 +183,7 @@ def _handle_input_payload(self, payload: bytes) -> None:
             )
 
     def _close(self) -> None:
-        try:
-            self._sock.shutdown(socket.SHUT_RDWR)
-        except OSError:
-            pass
-        try:
-            self._sock.close()
-        except OSError:
-            pass
+        self._channel.close()
 
 
 class RemoteDesktopHost:
@@ -338,7 +328,19 @@ def _accept_loop(self) -> None:
             wrapped = self._maybe_wrap_tls(client_sock, address)
             if wrapped is None:
                 continue
-            handler = _ClientHandler(self, wrapped, address)
+            try:
+                channel = self._build_channel(wrapped, address)
+            except (OSError, RuntimeError) as error:
+                autocontrol_logger.info(
+                    "remote_desktop channel handshake from %s failed: %r",
+                    address, error,
+                )
+                try:
+                    wrapped.close()
+                except OSError:
+                    pass
+                continue
+            handler = _ClientHandler(self, channel, address)
             with self._clients_lock:
                 if len(self._clients) >= self._max_clients:
                     autocontrol_logger.info(
@@ -351,6 +353,12 @@ def _accept_loop(self) -> None:
             handler.start()
             self._reap_dead_clients()
 
+    def _build_channel(self, sock: socket.socket,
+                       address) -> MessageChannel:
+        """Hook for transports: TCP wraps directly, WS overrides this."""
+        del address
+        return TcpMessageChannel(sock)
+
     def _maybe_wrap_tls(self, client_sock: socket.socket,
                         address) -> Optional[socket.socket]:
         """Return a TLS-wrapped socket when an ssl_context is configured."""

je_auto_control/utils/remote_desktop/transport.py

@@ -0,0 +1,123 @@
+"""Pluggable typed-message transport for the remote-desktop protocol.
+
+The host and viewer always exchange the same typed messages
+(``MessageType`` from :mod:`protocol`), but the wire layer can be either
+the original raw-TCP framing or WebSocket binary frames. ``MessageChannel``
+hides that distinction so the rest of the codebase deals with
+``send_typed`` / ``read_typed`` only.
+"""
+import socket
+import threading
+from typing import Tuple
+
+from je_auto_control.utils.remote_desktop.protocol import (
+    HEADER_SIZE, MessageType, ProtocolError,
+    decode_frame_header, encode_frame, read_message,
+)
+from je_auto_control.utils.remote_desktop.ws_protocol import (
+    recv_message as ws_recv_message,
+    send_binary as ws_send_binary,
+    send_close as ws_send_close,
+)
+
+
+class MessageChannel:
+    """Abstract bidirectional typed-message endpoint."""
+
+    def send_typed(self, message_type: MessageType, payload: bytes) -> None:
+        raise NotImplementedError
+
+    def read_typed(self) -> Tuple[MessageType, bytes]:
+        raise NotImplementedError
+
+    def settimeout(self, timeout) -> None:
+        raise NotImplementedError
+
+    def close(self) -> None:
+        raise NotImplementedError
+
+
+class TcpMessageChannel(MessageChannel):
+    """Original transport: each typed message is one length-prefixed frame."""
+
+    def __init__(self, sock: socket.socket) -> None:
+        self._sock = sock
+        self._send_lock = threading.Lock()
+
+    def send_typed(self, message_type: MessageType, payload: bytes) -> None:
+        data = encode_frame(message_type, payload)
+        with self._send_lock:
+            self._sock.sendall(data)
+
+    def read_typed(self) -> Tuple[MessageType, bytes]:
+        return read_message(self._sock)
+
+    def settimeout(self, timeout) -> None:
+        self._sock.settimeout(timeout)
+
+    def close(self) -> None:
+        try:
+            self._sock.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+        try:
+            self._sock.close()
+        except OSError:
+            pass
+
+    @property
+    def sock(self) -> socket.socket:
+        return self._sock
+
+
+class WsMessageChannel(MessageChannel):
+    """WebSocket transport: each WS BINARY frame carries one typed message.
+
+    The WS payload is the existing typed-frame encoding (magic + type +
+    length + body), so :func:`decode_frame_header` and :func:`encode_frame`
+    are reused unchanged. ``mask_outgoing`` follows RFC 6455: clients must
+    mask, servers must not.
+    """
+
+    def __init__(self, sock: socket.socket, mask_outgoing: bool) -> None:
+        self._sock = sock
+        self._mask = bool(mask_outgoing)
+        self._send_lock = threading.Lock()
+
+    def send_typed(self, message_type: MessageType, payload: bytes) -> None:
+        data = encode_frame(message_type, payload)
+        with self._send_lock:
+            ws_send_binary(self._sock, data, mask=self._mask)
+
+    def read_typed(self) -> Tuple[MessageType, bytes]:
+        ws_payload = ws_recv_message(self._sock)
+        if len(ws_payload) < HEADER_SIZE:
+            raise ProtocolError("WS payload too short to contain typed header")
+        msg_type, length = decode_frame_header(ws_payload[:HEADER_SIZE])
+        body = ws_payload[HEADER_SIZE:HEADER_SIZE + length]
+        if len(body) != length:
+            raise ProtocolError(
+                f"declared length {length} but ws payload had {len(body)}"
+            )
+        return msg_type, body
+
+    def settimeout(self, timeout) -> None:
+        self._sock.settimeout(timeout)
+
+    def close(self) -> None:
+        try:
+            ws_send_close(self._sock, mask=self._mask)
+        except OSError:
+            pass
+        try:
+            self._sock.shutdown(socket.SHUT_RDWR)
+        except OSError:
+            pass
+        try:
+            self._sock.close()
+        except OSError:
+            pass
+
+    @property
+    def sock(self) -> socket.socket:
+        return self._sock