Feature(Go): Screen frame relay end-to-end with graceful client BYE (Phase 4)

server/go/web/ws.go

@@ -31,27 +31,46 @@ var upgrader = websocket.Upgrader{
 
 // ----- per-connection client state ----------------------------------------
 
+// wsMsg is one queued WebSocket frame. binary toggles between
+// websocket.TextMessage (JSON signaling) and websocket.BinaryMessage
+// (screen frames).
+type wsMsg struct {
+	binary bool
+	data   []byte
+}
+
 type wsClient struct {
 	conn   *websocket.Conn
-	send   chan []byte
+	send   chan wsMsg
 	closed chan struct{}
 	once   sync.Once
 
 	// Mutated under wsHub.mu (or only by the read loop owning this client).
-	nonce string // outstanding challenge — cleared after a successful login
-	token string // set once authenticated
-	role  string // mirrors session role after login
-	addr  string // client address for logs
+	nonce    string // outstanding challenge — cleared after a successful login
+	token    string // set once authenticated
+	role     string // mirrors session role after login
+	addr     string // client address for logs
+	watching string // device ID this browser is currently streaming, "" when on the list
 }
 
-// queue writes a payload onto the send buffer. Drops silently if the buffer
-// is full so a stuck reader can't back-pressure the broadcast path.
+// queue writes a JSON text frame onto the send buffer. Drops silently if the
+// buffer is full so a stuck reader can't back-pressure the broadcast path.
 func (c *wsClient) queue(payload []byte) {
+	c.enqueue(wsMsg{binary: false, data: payload})
+}
+
+// queueBinary writes a binary WS frame. Used for screen-stream packets.
+func (c *wsClient) queueBinary(payload []byte) {
+	c.enqueue(wsMsg{binary: true, data: payload})
+}
+
+func (c *wsClient) enqueue(m wsMsg) {
 	select {
-	case c.send <- payload:
+	case c.send <- m:
 	case <-c.closed:
 	default:
-		// queue full — caller is responsible for noticing if it matters.
+		// queue full — drop (acceptable for video; signaling clients are
+		// typically not behind enough for the small text buffer to fill).
 	}
 }
 
@@ -105,6 +124,58 @@ func (h *wsHub) OnDeviceOffline(_ string) {
 	h.broadcastAuthenticated(`{"cmd":"devices_changed"}`)
 }
 
+// OnCursorChange relays the remote cursor index to every viewer of this
+// device. The browser maps the index to a CSS cursor (desktop) or overlay
+// SVG variant (touch). Hub already de-duplicates so we always have a real
+// transition here.
+func (h *wsHub) OnCursorChange(deviceID string, index byte) {
+	msg := mustJSON(map[string]any{
+		"cmd":   "cursor",
+		"index": index,
+	})
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	for c := range h.clients {
+		if c.watching == deviceID && c.token != "" {
+			c.queue(msg)
+		}
+	}
+}
+
+// OnResolutionChange notifies viewers so the browser-side WebCodecs decoder
+// can be (re)initialized with the right frame size. Without this, incoming
+// binary frames after connect_result are decoded by an uninitialized
+// VideoDecoder and the page stays on "Waiting for video...".
+func (h *wsHub) OnResolutionChange(deviceID string, width, height int) {
+	msg := mustJSON(map[string]any{
+		"cmd":    "resolution_changed",
+		"id":     deviceID,
+		"width":  width,
+		"height": height,
+	})
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	for c := range h.clients {
+		if c.watching == deviceID && c.token != "" {
+			c.queue(msg)
+		}
+	}
+}
+
+// OnScreenFrame ships a screen packet to every browser currently watching
+// this device. We hold the read lock for the whole iteration, but each
+// queueBinary is non-blocking (drops on backpressure) so a slow viewer
+// cannot stall the fast ones.
+func (h *wsHub) OnScreenFrame(deviceID string, packet []byte, _ bool) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	for c := range h.clients {
+		if c.watching == deviceID && c.token != "" {
+			c.queueBinary(packet)
+		}
+	}
+}
+
 // OnDeviceUpdate forwards heartbeat-derived liveness data so the device-list
 // rows can refresh RTT and active-window labels without re-fetching.
 func (h *wsHub) OnDeviceUpdate(id string, rtt int, activeWindow string) {
@@ -144,6 +215,12 @@ func (h *wsHub) unregister(c *wsClient) {
 	h.mu.Lock()
 	delete(h.clients, c)
 	h.mu.Unlock()
+	// If this client was the last viewer of a device, tear down the screen
+	// session so the device stops encoding. Done OUTSIDE the lock so the
+	// hub's mutators can take their own locks without risk of recursion.
+	if c.watching != "" && h.countWatchers(c.watching) == 0 {
+		h.devices.CloseScreen(c.watching)
+	}
 	// Do NOT revoke the token: tokens are session-scoped, not WS-scoped.
 	// Frontend may close+reopen the WS at any time (visibilitychange handler,
 	// brief network blip, reload) and must be able to resume with the same
@@ -170,7 +247,7 @@ func (h *wsHub) serve(w http.ResponseWriter, r *http.Request) {
 
 	client := &wsClient{
 		conn:   conn,
-		send:   make(chan []byte, wsSendBuffer),
+		send:   make(chan wsMsg, wsSendBuffer),
 		closed: make(chan struct{}),
 		nonce:  nonce,
 		addr:   r.RemoteAddr,
@@ -192,8 +269,12 @@ func (h *wsHub) writeLoop(c *wsClient) {
 	for {
 		select {
 		case msg := <-c.send:
+			msgType := websocket.TextMessage
+			if msg.binary {
+				msgType = websocket.BinaryMessage
+			}
 			_ = c.conn.SetWriteDeadline(time.Now().Add(wsWriteWait))
-			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
+			if err := c.conn.WriteMessage(msgType, msg.data); err != nil {
 				c.close()
 				return
 			}