diff --git a/server/go/cmd/main.go b/server/go/cmd/main.go index 84588de..d0d403c 100644 --- a/server/go/cmd/main.go +++ b/server/go/cmd/main.go @@ -138,11 +138,16 @@ func (h *MyHandler) OnReceive(ctx *connection.Context, data []byte) { case protocol.TokenClipboardText: h.handleClipboardText(ctx, data) case protocol.TokenKeyframe: - // Sent by the client only when frameID % m_GOP == 0; the client's - // DEFAULT_GOP is 0x7FFFFFFF (effectively infinite), so this token - // is essentially unused in practice. Treat as a no-op for now — - // IDRs always arrive in-band via TOKEN_NEXTSCREEN and we catch - // them via the H264 NAL scan in handleScreenFrame. + // The client switches to this token when it emits an IDR on request + // (COMMAND_SCREEN_KEYFRAME_REQ) or on a GOP boundary. The header + // layout is identical to TOKEN_NEXTSCREEN, but the payload is always + // an IDR, so relay it as frameType=1 without sniffing. Dropping it + // (the old behaviour, harmless while the client's DEFAULT_GOP of + // 0x7FFFFFFF meant this token never appeared) breaks keyframe + // recovery: the browser keeps decoding P-frames against a reference + // chain it never received, so the picture stays corrupt until the + // screen session restarts. + h.handleScreenFrame(ctx, data, true) case protocol.CmdCursorImage: // Custom cursor bitmaps — relayed in Phase 5+ when the web cursor // overlay learns to render arbitrary BGRA images. Drop silently for diff --git a/server/go/cmd/screen_keyframe_test.go b/server/go/cmd/screen_keyframe_test.go new file mode 100644 index 0000000..d5a7f44 --- /dev/null +++ b/server/go/cmd/screen_keyframe_test.go @@ -0,0 +1,72 @@ +package main + +import ( + "testing" + + "github.com/yuanyuanxiang/SimpleRemoter/server/go/connection" + "github.com/yuanyuanxiang/SimpleRemoter/server/go/hub" + "github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol" +) + +type frameRecorder struct { + frames []recordedFrame +} + +type recordedFrame struct { + deviceID string + packet []byte + isKey bool +} + +func (r *frameRecorder) OnScreenFrame(deviceID string, packet []byte, isKeyframe bool) { + r.frames = append(r.frames, recordedFrame{deviceID, packet, isKeyframe}) +} +func (r *frameRecorder) OnDeviceOnline(hub.DeviceInfo) {} +func (r *frameRecorder) OnDeviceOffline(string) {} +func (r *frameRecorder) OnDeviceUpdate(string, int, string) {} +func (r *frameRecorder) OnResolutionChange(string, int, int) {} +func (r *frameRecorder) OnCursorChange(string, byte) {} +func (r *frameRecorder) OnClipboardText(string, string) {} +func (r *frameRecorder) OnTerminalReady(string, bool) {} +func (r *frameRecorder) OnTerminalData(string, []byte) {} +func (r *frameRecorder) OnTerminalClosed(string, string) {} + +// TestOnReceiveRelaysKeyframeToken guards the forced-IDR recovery path: when +// the browser asks for a keyframe, the client replies with TOKEN_KEYFRAME. +// Dropping that token (the old no-op) leaves the browser decoding P-frames +// against a reference chain it never received, i.e. persistent corruption. +func TestOnReceiveRelaysKeyframeToken(t *testing.T) { + h := hub.New() + rec := &frameRecorder{} + h.Subscribe(rec) + + h.Register(&hub.Device{ID: "dev", PeerIP: "10.0.0.1"}, + &connection.Context{RemoteAddr: "10.0.0.1:1000"}) + screenCtx := &connection.Context{RemoteAddr: "10.0.0.1:1001"} + if !h.BindScreenConn("dev", screenCtx) { + t.Fatal("screen bind failed") + } + + // [TOKEN_KEYFRAME][algo=H264][cursorPos:8][cursorIdx:1][IDR NAL] + data := []byte{ + protocol.TokenKeyframe, 2, + 0, 0, 0, 0, 0, 0, 0, 0, + 0, + 0x00, 0x00, 0x00, 0x01, 0x65, 0x88, 0x84, + } + (&MyHandler{hub: h}).OnReceive(screenCtx, data) + + if len(rec.frames) != 1 { + t.Fatalf("frames relayed = %d, want 1 (TOKEN_KEYFRAME must not be dropped)", len(rec.frames)) + } + f := rec.frames[0] + if f.deviceID != "dev" { + t.Fatalf("deviceID = %q, want dev", f.deviceID) + } + if !f.isKey || f.packet[4] != 1 { + t.Fatalf("isKeyframe=%v frameType=%d, want true/1", f.isKey, f.packet[4]) + } + if want := 9 + len(data) - (1 + protocol.ScreenFrameHeaderLen); len(f.packet) != want { + t.Fatalf("packet len = %d, want %d", len(f.packet), want) + } +} diff --git a/server/go/hub/hub.go b/server/go/hub/hub.go index 52d62e3..65124b5 100644 --- a/server/go/hub/hub.go +++ b/server/go/hub/hub.go @@ -581,6 +581,13 @@ func (h *Hub) PublishResolution(deviceID string, width, height int) { if d, ok := h.devices[deviceID]; ok { d.screenWidth = width d.screenHeight = height + // The cached IDR belongs to the previous capture setup. The client + // re-sends TOKEN_BITMAPINFO on the same sub-conn after a window + // resize / capture restart, so this is not only the + // new-sub-conn path (BindScreenConn already clears there). Serving + // the old-resolution IDR alongside the new dimensions leaves the + // decoder misconfigured and garbles the picture until the next IDR. + d.lastKeyframe = nil } h.mu.Unlock() for _, s := range h.snapshotSubscribers() { diff --git a/server/go/protocol/commands.go b/server/go/protocol/commands.go index 5cc2829..d9129e6 100644 --- a/server/go/protocol/commands.go +++ b/server/go/protocol/commands.go @@ -112,6 +112,7 @@ const ( CommandShell byte = 40 // COMMAND_SHELL - ask device to open a shell sub-connection CommandTerminalRsize byte = 81 // CMD_TERMINAL_RESIZE - [cmd:1][cols:2 LE][rows:2 LE] CmdRestoreConsole byte = 82 // CMD_RESTORE_CONSOLE - RDP session "归位": switch back to the console session and restart capture + CommandScreenKeyframeReq byte = 169 // COMMAND_SCREEN_KEYFRAME_REQ - force the client to emit an IDR immediately CommandBye byte = 204 // COMMAND_BYE - disconnect CommandHeartbeat byte = 216 // CMD_HEARTBEAT_ACK diff --git a/server/go/web/watch_set_test.go b/server/go/web/watch_set_test.go new file mode 100644 index 0000000..934996d --- /dev/null +++ b/server/go/web/watch_set_test.go @@ -0,0 +1,96 @@ +package web + +import ( + "fmt" + "reflect" + "sort" + "testing" +) + +func sortedWatchIDs(c *wsClient) []string { + out := make([]string, 0, len(c.watchIDs)) + for id := range c.watchIDs { + out = append(out, id) + } + sort.Strings(out) + return out +} + +func TestApplyConnectWatchMulti(t *testing.T) { + c := &wsClient{watchIDs: make(map[string]struct{})} + + if stopped, tooMany := applyConnectWatch(c, "a", true); len(stopped) != 0 || tooMany { + t.Fatalf("first connect = (%v, %v), want ([], false)", stopped, tooMany) + } + if stopped, tooMany := applyConnectWatch(c, "b", true); len(stopped) != 0 || tooMany { + t.Fatalf("second connect = (%v, %v), want ([], false)", stopped, tooMany) + } + if got, want := sortedWatchIDs(c), []string{"a", "b"}; !reflect.DeepEqual(got, want) { + t.Fatalf("watch set = %v, want %v", got, want) + } + + // Repeated connect (focus switch) must be a no-op. + if stopped, tooMany := applyConnectWatch(c, "a", true); len(stopped) != 0 || tooMany { + t.Fatalf("repeat connect = (%v, %v), want ([], false)", stopped, tooMany) + } + if got, want := sortedWatchIDs(c), []string{"a", "b"}; !reflect.DeepEqual(got, want) { + t.Fatalf("watch set after repeat = %v, want %v", got, want) + } +} + +func TestApplyConnectWatchLegacyReplaces(t *testing.T) { + c := &wsClient{watchIDs: map[string]struct{}{"a": {}, "b": {}}} + + stopped, tooMany := applyConnectWatch(c, "c", false) + if tooMany { + t.Fatal("legacy connect hit the cap, want not") + } + sort.Strings(stopped) + if want := []string{"a", "b"}; !reflect.DeepEqual(stopped, want) { + t.Fatalf("stopped = %v, want %v", stopped, want) + } + if got, want := sortedWatchIDs(c), []string{"c"}; !reflect.DeepEqual(got, want) { + t.Fatalf("watch set = %v, want %v", got, want) + } +} + +func TestApplyConnectWatchCap(t *testing.T) { + c := &wsClient{watchIDs: make(map[string]struct{})} + for i := 0; i < maxWatchPerClient; i++ { + applyConnectWatch(c, fmt.Sprintf("d%02d", i), true) + } + if _, tooMany := applyConnectWatch(c, "overflow", true); !tooMany { + t.Fatal("expected the cap to reject a new device") + } + if len(c.watchIDs) != maxWatchPerClient { + t.Fatalf("watch set size = %d, want %d", len(c.watchIDs), maxWatchPerClient) + } + // An already watched device must still be accepted at the cap. + if _, tooMany := applyConnectWatch(c, "d00", true); tooMany { + t.Fatal("repeat connect at the cap was rejected") + } +} + +func TestResolveWatchDevice(t *testing.T) { + h := &wsHub{} + c := &wsClient{watchIDs: map[string]struct{}{"a": {}}} + + if got := h.resolveWatchDevice(c, []byte(`{"id":"a"}`)); got != "a" { + t.Fatalf("explicit watched id = %q, want a", got) + } + if got := h.resolveWatchDevice(c, []byte(`{"id":"nope"}`)); got != "" { + t.Fatalf("explicit unwatched id = %q, want empty", got) + } + // Old frontend: no id, exactly one watched device. + if got := h.resolveWatchDevice(c, []byte(`{}`)); got != "a" { + t.Fatalf("no id, single device = %q, want a", got) + } + + c.watchIDs["b"] = struct{}{} + if got := h.resolveWatchDevice(c, []byte(`{}`)); got != "" { + t.Fatalf("no id, two devices = %q, want empty", got) + } + if got := h.resolveWatchDevice(c, []byte(`{"id":"a"}`)); got != "a" { + t.Fatalf("explicit id, two devices = %q, want a", got) + } +} diff --git a/server/go/web/ws.go b/server/go/web/ws.go index 048730e..da99764 100644 --- a/server/go/web/ws.go +++ b/server/go/web/ws.go @@ -52,8 +52,16 @@ type wsClient struct { 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 termWatching string // device ID for an open web terminal session, "" otherwise + + // watchIDs is the set of devices this browser is streaming (the screen + // wall watches several at once). Empty while on the device list. + // Guarded by wsHub.mu. + watchIDs map[string]struct{} + + // keyframeAt throttles keyframe requests to one per device per second. + // Unix seconds, guarded by wsHub.mu. + keyframeAt map[string]int64 } // queue writes a JSON text frame onto the send buffer. Drops silently if the @@ -85,6 +93,24 @@ func (c *wsClient) close() { }) } +// challengeCaps advertises optional protocol features. It is appended to every +// challenge, including the rotation after a failed login, because the browser +// overwrites its capability state on each challenge. Old frontends ignore it. +const challengeCaps = `,"caps":{"multiWatch":true,"keyframeRequest":true}` + +// maxWatchPerClient bounds the number of devices one WebSocket may stream at +// once. The wall tops out at 25 tiles; the slack covers a focus switch racing +// a teardown. This keeps the per-frame broadcast loop cheap for a misbehaving +// or buggy client. +const maxWatchPerClient = 32 + +// isWatching reports whether this connection subscribed to deviceID. +// Callers must hold wsHub.mu (read or write). +func (c *wsClient) isWatching(deviceID string) bool { + _, ok := c.watchIDs[deviceID] + return ok +} + // ----- ws hub: registry of all connected browsers ------------------------- type wsHub struct { @@ -253,10 +279,19 @@ func (h *wsHub) OnDeviceOffline(id string) { "cmd": "device_offline", "id": id, }) - h.mu.RLock() - defer h.mu.RUnlock() + // Drop the id from every watch set as well: the device is gone, so this + // connection no longer holds a subscription for it. The browser does the + // same with its local refcount. Without this, a stale entry would keep + // countWatchers non-zero and later block the screen teardown. + h.mu.Lock() + defer h.mu.Unlock() for c := range h.clients { - if c.watching == id && c.token != "" { + if !c.isWatching(id) { + continue + } + delete(c.watchIDs, id) + delete(c.keyframeAt, id) + if c.token != "" { c.queue(msg) } } @@ -269,12 +304,13 @@ func (h *wsHub) OnDeviceOffline(id string) { func (h *wsHub) OnCursorChange(deviceID string, index byte) { msg := mustJSON(map[string]any{ "cmd": "cursor", + "id": deviceID, "index": index, }) h.mu.RLock() defer h.mu.RUnlock() for c := range h.clients { - if c.watching == deviceID && c.token != "" { + if c.isWatching(deviceID) && c.token != "" { c.queue(msg) } } @@ -294,7 +330,7 @@ func (h *wsHub) OnResolutionChange(deviceID string, width, height int) { h.mu.RLock() defer h.mu.RUnlock() for c := range h.clients { - if c.watching == deviceID && c.token != "" { + if c.isWatching(deviceID) && c.token != "" { c.queue(msg) } } @@ -311,7 +347,7 @@ func (h *wsHub) OnClipboardText(deviceID string, text string) { h.mu.RLock() defer h.mu.RUnlock() for c := range h.clients { - if c.watching == deviceID && c.token != "" { + if c.isWatching(deviceID) && c.token != "" { c.queue(msg) } } @@ -325,7 +361,7 @@ 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 != "" { + if c.isWatching(deviceID) && c.token != "" { c.queueBinary(packet) } } @@ -424,12 +460,22 @@ func (h *wsHub) register(c *wsClient) { func (h *wsHub) unregister(c *wsClient) { h.mu.Lock() delete(h.clients, c) + // Snapshot the watch set and clear it under the same lock. The teardown + // below must run OUTSIDE the lock: countWatchers takes h.mu.RLock and + // deferredCloseScreen takes h.mu.Lock. + watched := make([]string, 0, len(c.watchIDs)) + for id := range c.watchIDs { + watched = append(watched, id) + } + c.watchIDs = nil + c.keyframeAt = nil 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.deferredCloseScreen(c.watching) + // session so the device stops encoding. + for _, id := range watched { + if h.countWatchers(id) == 0 { + h.deferredCloseScreen(id) + } } // Terminal sessions are single-viewer by design, so any open session // belongs to this client. Tear it down so the next viewer doesn't @@ -492,11 +538,13 @@ func (h *wsHub) serve(w http.ResponseWriter, r *http.Request) { } client := &wsClient{ - conn: conn, - send: make(chan wsMsg, wsSendBuffer), - closed: make(chan struct{}), - nonce: nonce, - addr: clientIP(r, h.trustForwardedFor), + conn: conn, + send: make(chan wsMsg, wsSendBuffer), + closed: make(chan struct{}), + nonce: nonce, + addr: clientIP(r, h.trustForwardedFor), + watchIDs: make(map[string]struct{}), + keyframeAt: make(map[string]int64), } h.register(client) defer h.unregister(client) @@ -504,7 +552,7 @@ func (h *wsHub) serve(w http.ResponseWriter, r *http.Request) { go h.writeLoop(client) // Greet with a challenge nonce so the browser can compute the login response. - client.queue([]byte(`{"cmd":"challenge","nonce":"` + nonce + `"}`)) + client.queue([]byte(`{"cmd":"challenge","nonce":"` + nonce + `"` + challengeCaps + `}`)) h.readLoop(client) } diff --git a/server/go/web/ws_handlers.go b/server/go/web/ws_handlers.go index 435aca3..91813fe 100644 --- a/server/go/web/ws_handlers.go +++ b/server/go/web/ws_handlers.go @@ -24,7 +24,8 @@ func (h *wsHub) rotateChallenge(c *wsClient) { return } c.nonce = n - c.queue([]byte(`{"cmd":"challenge","nonce":"` + n + `"}`)) + // Must carry caps too: the browser replaces serverCaps on every challenge. + c.queue([]byte(`{"cmd":"challenge","nonce":"` + n + `"` + challengeCaps + `}`)) } // dispatch routes one inbound message to its handler. The `raw` payload is @@ -50,6 +51,8 @@ func (h *wsHub) dispatch(c *wsClient, cmd string, raw []byte) { case "connect": h.handleConnect(c, raw) + case "keyframe": + h.handleKeyframe(c, raw) case "rdp_reset": h.handleRdpReset(c, raw) case "audio_toggle": @@ -240,66 +243,196 @@ func (h *wsHub) resetLoginThrottle(c *wsClient, username string) { // screen sub-connection. Closing the TCP sub-conn is the signal the C++ // device firmware uses to stop screen capture, so this is how we ask the // device to free its encoder. -func (h *wsHub) handleDisconnect(c *wsClient, _ []byte) { - // Mirror handleConnect: take h.mu so event-handler readers - // (OnResolutionChange/OnScreenFrame) get a consistent view of c.watching. +func (h *wsHub) handleDisconnect(c *wsClient, raw []byte) { + var in struct { + ID string `json:"id"` + MultiWatch bool `json:"multiWatch"` + Req string `json:"req"` + } + _ = json.Unmarshal(raw, &in) + + // Take h.mu so event-handler readers (OnResolutionChange/OnScreenFrame) + // get a consistent view of the watch set. + var closed []string h.mu.Lock() - prev := c.watching - c.watching = "" + if in.MultiWatch && in.ID != "" { + // New frontend: drop only this device; the other wall tiles keep + // streaming on the same connection. + if c.isWatching(in.ID) { + delete(c.watchIDs, in.ID) + delete(c.keyframeAt, in.ID) + closed = append(closed, in.ID) + } + } else { + // Old frontend, or a disconnect without an id: it only ever watched + // one device, so clear everything. + for id := range c.watchIDs { + closed = append(closed, id) + } + c.watchIDs = make(map[string]struct{}) + c.keyframeAt = make(map[string]int64) + } h.mu.Unlock() - c.queue([]byte(`{"cmd":"disconnect_result","ok":true}`)) - if prev != "" && h.countWatchers(prev) == 0 { - h.deferredCloseScreen(prev) + + res := map[string]any{"cmd": "disconnect_result", "ok": true} + if in.ID != "" { + res["id"] = in.ID + } + if in.Req != "" { + res["req"] = in.Req + } + c.queue(mustJSON(res)) + + // deferredCloseScreen takes h.mu, so it must run outside the lock above. + for _, id := range closed { + if h.countWatchers(id) == 0 { + h.deferredCloseScreen(id) + } } } -// countWatchers returns how many authenticated clients still have their -// `watching` field pointing at deviceID. Called from disconnect paths. +// countWatchers returns how many clients still have deviceID in their watch +// set. Called from disconnect paths. func (h *wsHub) countWatchers(deviceID string) int { h.mu.RLock() defer h.mu.RUnlock() n := 0 for c := range h.clients { - if c.watching == deviceID { + if c.isWatching(deviceID) { n++ } } return n } +// applyConnectWatch mutates c.watchIDs for a connect request. The caller must +// hold h.mu. +// +// It returns the ids dropped by the legacy replace-all semantics; the caller +// must stop their screen sessions outside the lock (only if nobody else still +// watches them). tooMany reports that the per-connection cap blocked the id. +func applyConnectWatch(c *wsClient, id string, multiWatch bool) (stopped []string, tooMany bool) { + if c.watchIDs == nil { + c.watchIDs = make(map[string]struct{}) + } + if !multiWatch { + // Old frontend: one device per connection. Replace the whole set so + // switching devices repeatedly cannot grow it without bound. + for old := range c.watchIDs { + if old != id { + stopped = append(stopped, old) + } + } + c.watchIDs = make(map[string]struct{}) + } + if _, ok := c.watchIDs[id]; ok { + return stopped, false // repeated connect (focus switch) is a no-op + } + if len(c.watchIDs) >= maxWatchPerClient { + return stopped, true + } + c.watchIDs[id] = struct{}{} + return stopped, false +} + +// resolveWatchDevice maps an input command to the device it targets. The new +// frontend sends an explicit id, which must be in this connection's watch set +// so a stale or forged id can never inject input into a device the browser is +// not viewing. Old frontends omit id — accepted only while exactly one device +// is watched. With several devices and no id the command is dropped. +func (h *wsHub) resolveWatchDevice(c *wsClient, raw []byte) string { + var in struct { + ID string `json:"id"` + } + _ = json.Unmarshal(raw, &in) + h.mu.RLock() + defer h.mu.RUnlock() + if in.ID != "" { + if c.isWatching(in.ID) { + return in.ID + } + return "" + } + if len(c.watchIDs) == 1 { + for id := range c.watchIDs { + return id + } + } + return "" +} + func (h *wsHub) handleConnect(c *wsClient, raw []byte) { if !h.requireAuth(c, raw, "connect_result") { return } var in struct { - ID string `json:"id"` + ID string `json:"id"` + MultiWatch bool `json:"multiWatch"` + Req string `json:"req"` } if err := json.Unmarshal(raw, &in); err != nil || in.ID == "" { c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": false, "msg": "Bad request"})) return } - // If a screen session is already live for this device (another browser - // is already watching), reuse it: hand the new viewer the current - // resolution and the most recent IDR keyframe so its decoder can start - // rendering immediately, without waiting for the next IDR (~15 s). - cache := h.devices.ScreenState(in.ID) - if cache.Active { - c.queue(mustJSON(map[string]any{ - "cmd": "connect_result", "ok": true, - "width": cache.Width, "height": cache.Height, - })) - if len(cache.Keyframe) > 0 { - c.queueBinary(cache.Keyframe) + // Every response echoes id and req: one socket may now serve several + // devices (screen wall), so the browser routes by id (and req when it + // needs to pair a reply with a specific request). + reply := func(ok bool, msg string) map[string]any { + m := map[string]any{"cmd": "connect_result", "ok": ok, "id": in.ID} + if in.Req != "" { + m["req"] = in.Req } - h.mu.Lock() - // Cancel any pending deferred close so the relay stays alive. + if msg != "" { + m["msg"] = msg + } + return m + } + + var stopped []string + var tooMany bool + h.mu.Lock() + stopped, tooMany = applyConnectWatch(c, in.ID, in.MultiWatch) + // A pending deferred close means the previous session is still alive for + // a few seconds; cancel it so the cached resolution + IDR hot path wins. + // Only when the subscription was accepted: a rejected connect must not + // cancel the teardown of a device nobody is watching any more, or that + // session would keep streaming forever. + if !tooMany { if t, ok := h.screenCloseTimers[in.ID]; ok { t.Stop() delete(h.screenCloseTimers, in.ID) } - c.watching = in.ID - h.mu.Unlock() + } + h.mu.Unlock() + + if tooMany { + c.queue(mustJSON(reply(false, "Too many devices on one connection"))) + return + } + // Only the legacy replace-all semantics can produce stopped ids. Their + // screen sessions are torn down outside the lock; deferredCloseScreen is + // a no-op while another browser still watches the device. + for _, id := range stopped { + if h.countWatchers(id) == 0 { + h.deferredCloseScreen(id) + } + } + + // If a screen session is already live for this device (another browser + // is already watching, or this connection reconnected), reuse it: hand + // the viewer the cached resolution and the most recent IDR keyframe so + // its decoder can start rendering immediately, without waiting for the + // next IDR (~15 s). + cache := h.devices.ScreenState(in.ID) + if cache.Active { + m := reply(true, "") + m["width"] = cache.Width + m["height"] = cache.Height + c.queue(mustJSON(m)) + if len(cache.Keyframe) > 0 { + c.queueBinary(cache.Keyframe) + } return } @@ -312,41 +445,101 @@ func (h *wsHub) handleConnect(c *wsClient, raw []byte) { cmd[2] = protocol.AlgorithmH264 cmd[3] = 1 - // CRITICAL: bind c.watching BEFORE asking the device to start capturing. - // On fast reconnects the device's screen sub-conn handshake completes in - // <100 ms, so TOKEN_BITMAPINFO and even the first H264 frame can arrive - // before this handler finishes — and the resolution_changed / frame - // broadcasts in wsHub filter on c.watching. With the assignment after - // SendToDevice the new viewer silently misses the very first IDR and - // resolution_changed, leaving the page stuck on "Waiting for video". - // - // The write needs to share the lock event handlers use to read c.watching - // (they iterate h.clients under h.mu.RLock). Without that the write is a - // data race; on a fast reconnect the reader goroutine can keep observing - // the previous value ("") long enough to drop the first resolution_changed - // and the first IDR, which produces the exact "every other quick reconnect - // goes black" symptom — the C++ server avoids it because it does the same - // state mutation under std::mutex and reaps the memory-barrier as a bonus. - h.mu.Lock() - c.watching = in.ID - h.mu.Unlock() - + // CRITICAL: the watch entry is inserted above (applyConnectWatch) BEFORE + // asking the device to start capturing. On fast reconnects the device's + // screen sub-conn handshake completes in <100 ms, so TOKEN_BITMAPINFO and + // even the first H264 frame can arrive before this handler finishes — and + // the resolution_changed / frame broadcasts filter on the watch set. If + // the entry were added afterwards the new viewer would silently miss the + // very first IDR and resolution_changed, leaving the page stuck on + // "Waiting for video". if err := h.devices.SendToDevice(in.ID, cmd); err != nil { - // Roll back the watching flag if we never managed to kick capture. + // Roll back the subscription: no screen session was started. h.mu.Lock() - c.watching = "" + delete(c.watchIDs, in.ID) + delete(c.keyframeAt, in.ID) h.mu.Unlock() msg := "Device offline" if err != hub.ErrDeviceOffline { msg = "Failed to start screen capture" h.log.Error("SendToDevice(%s): %v", in.ID, err) } - c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": false, "msg": msg})) + c.queue(mustJSON(reply(false, msg))) return } h.log.Debug("[timing] COMMAND_SCREEN_SPY sent to device=%s (cold start)", in.ID) - c.queue(mustJSON(map[string]any{"cmd": "connect_result", "ok": true})) + c.queue(mustJSON(reply(true, ""))) +} + +// handleKeyframe asks the device to emit an IDR immediately. The browser calls +// it when it resumes decoding a wall tile, or right after connecting to reset +// the reference chain. The server first replays its cached IDR so the decoder +// has something to paint at once, then sends COMMAND_SCREEN_KEYFRAME_REQ to +// the client, whose fresh IDR replaces the cache (the cached one cannot +// guarantee the following P-frames reference it). +// +// Mirrors CWebService::HandleKeyframe. Requests are limited to one per device +// per second, so a retrying frontend cannot spin the encoder. +func (h *wsHub) handleKeyframe(c *wsClient, raw []byte) { + if !h.requireAuth(c, raw, "keyframe_result") { + return + } + var in struct { + ID string `json:"id"` + Req string `json:"req"` + } + _ = json.Unmarshal(raw, &in) + + reply := func(ok bool, msg string) map[string]any { + m := map[string]any{"cmd": "keyframe_result", "ok": ok} + if in.ID != "" { + m["id"] = in.ID + } + if in.Req != "" { + m["req"] = in.Req + } + if msg != "" { + m["msg"] = msg + } + return m + } + if in.ID == "" { + c.queue(mustJSON(reply(false, "Bad request"))) + return + } + + allowed := false + h.mu.Lock() + if c.isWatching(in.ID) { + now := time.Now().Unix() + if last, ok := c.keyframeAt[in.ID]; !ok || now > last { + if c.keyframeAt == nil { + c.keyframeAt = make(map[string]int64) + } + c.keyframeAt[in.ID] = now + allowed = true + } + } + h.mu.Unlock() + if !allowed { + c.queue(mustJSON(reply(false, "Not watching device or rate limited"))) + return + } + + cache := h.devices.ScreenState(in.ID) + if !cache.Active { + c.queue(mustJSON(reply(false, "No active screen session"))) + return + } + if len(cache.Keyframe) > 0 { + c.queueBinary(cache.Keyframe) + } + if err := h.devices.SendToScreen(in.ID, []byte{protocol.CommandScreenKeyframeReq}); err != nil { + c.queue(mustJSON(reply(false, "Failed to request keyframe"))) + return + } + c.queue(mustJSON(reply(true, ""))) } // handleRdpReset asks the device to switch its screen capture back to the @@ -363,7 +556,7 @@ func (h *wsHub) handleRdpReset(c *wsClient, raw []byte) { if !h.requireAuth(c, raw, "rdp_reset_result") { return } - deviceID := c.watching + deviceID := h.resolveWatchDevice(c, raw) if deviceID == "" { h.log.Warn("rdp_reset: no device watched (addr=%s role=%s)", c.addr, c.role) return @@ -440,7 +633,7 @@ func (h *wsHub) handleMouse(c *wsClient, raw []byte) { if err := json.Unmarshal(raw, &in); err != nil { return } - deviceID := c.watching + deviceID := h.resolveWatchDevice(c, raw) if deviceID == "" { return // browser hasn't picked a device yet } @@ -525,7 +718,7 @@ func (h *wsHub) handleKey(c *wsClient, raw []byte) { if in.KeyCode == protocol.VKLWin || in.KeyCode == protocol.VKRWin { return // Windows keys stay local; matches both C++ server and client } - deviceID := c.watching + deviceID := h.resolveWatchDevice(c, raw) if deviceID == "" { return } @@ -583,7 +776,7 @@ func (h *wsHub) handleClipboardSet(c *wsClient, raw []byte) { if in.Text == "" { return } - deviceID := c.watching + deviceID := h.resolveWatchDevice(c, raw) if deviceID == "" { return } @@ -607,7 +800,7 @@ func (h *wsHub) handleClipboardGet(c *wsClient, raw []byte) { if !h.requireAuth(c, raw, "clipboard_get_result") { return } - deviceID := c.watching + deviceID := h.resolveWatchDevice(c, raw) if deviceID == "" { return }