From 698bcf681e141a92ee7fd3e23abef9febec2e224 Mon Sep 17 00:00:00 2001 From: KingKDot Date: Sun, 19 Jul 2026 02:33:30 -0500 Subject: [PATCH] keyframe fixes and manually allow users to refresh keyframes if shit gets all fucked up --- Overlord-Client/cmd/agent/capture/capture.go | 1 + .../cmd/agent/capture/capture_test.go | 11 +++ .../cmd/agent/capture/h264_direct_windows.go | 19 +++++- .../cmd/agent/capture/win_duplication_test.go | 25 +++++++ .../cmd/agent/handlers/screenshot.go | 15 ++--- .../agent/handlers/screenshot_windows_test.go | 66 ++++++++++-------- .../public/assets/remotedesktop.js | 52 ++++++++++++-- Overlord-Server/public/remotedesktop.html | 9 +++ .../routes/websocket-lifecycle-routes.ts | 6 +- .../server/ws-console-rd-backstage.test.ts | 67 +++++++++++++++++++ .../src/server/ws-console-rd-backstage.ts | 26 +++++++ Overlord-Server/src/wsHandlers.test.ts | 40 ++++++++++- Overlord-Server/src/wsHandlers.ts | 53 ++++++++++----- .../test/e2e/remotedesktop-input.e2e.ts | 34 ++++++++++ 14 files changed, 360 insertions(+), 64 deletions(-) diff --git a/Overlord-Client/cmd/agent/capture/capture.go b/Overlord-Client/cmd/agent/capture/capture.go index 63495a7..5d482a8 100644 --- a/Overlord-Client/cmd/agent/capture/capture.go +++ b/Overlord-Client/cmd/agent/capture/capture.go @@ -645,6 +645,7 @@ func ResetPrevbackstage() { func RequestDesktopFullFrame() { requestFullFrames(2) + webrtcpub.RequestKeyframe() RequestDesktopH264Keyframe() } diff --git a/Overlord-Client/cmd/agent/capture/capture_test.go b/Overlord-Client/cmd/agent/capture/capture_test.go index 8f6676d..316b54f 100644 --- a/Overlord-Client/cmd/agent/capture/capture_test.go +++ b/Overlord-Client/cmd/agent/capture/capture_test.go @@ -11,6 +11,7 @@ import ( "overlord-client/cmd/agent/config" rt "overlord-client/cmd/agent/runtime" + "overlord-client/cmd/agent/webrtcpub" "overlord-client/cmd/agent/wire" "github.com/vmihailenco/msgpack/v5" @@ -476,3 +477,13 @@ func TestSetFrameFlowTargetFPSEnvOverride(t *testing.T) { t.Fatalf("expected env slot limit 12, got %d", got) } } + +func TestRequestDesktopFullFrameRequestsWebRTCKeyframe(t *testing.T) { + _ = webrtcpub.ConsumeKeyframeRequest() + + RequestDesktopFullFrame() + + if !webrtcpub.ConsumeKeyframeRequest() { + t.Fatal("expected desktop full-frame request to force the next WebRTC video frame to be a keyframe") + } +} diff --git a/Overlord-Client/cmd/agent/capture/h264_direct_windows.go b/Overlord-Client/cmd/agent/capture/h264_direct_windows.go index aac0931..85ce65a 100644 --- a/Overlord-Client/cmd/agent/capture/h264_direct_windows.go +++ b/Overlord-Client/cmd/agent/capture/h264_direct_windows.go @@ -9,12 +9,24 @@ import ( "overlord-client/cmd/agent/wire" ) +func directVideoKeyframeDue(requested bool, nowNs, lastKeyframeNs int64) bool { + if requested || lastKeyframeNs == 0 { + return true + } + return time.Duration(nowNs-lastKeyframeNs) >= keyframeEvery +} + func tryBuildDirectH264Frame(display int) (wire.Frame, time.Duration, time.Duration, bool, error) { codec := blockCodec() if (codec != "h264" && codec != "hevc") || (codec == "h264" && useDesktopSoftwareH264()) || !useDesktopDuplication() { return wire.Frame{}, 0, 0, false, nil } - forceKeyframe := webrtcpub.ConsumeKeyframeRequest() + now := time.Now() + forceKeyframe := directVideoKeyframeDue( + webrtcpub.ConsumeKeyframeRequest(), + now.UnixNano(), + lastKeyframe.Load(), + ) var data []byte var width, height int var captureDur, encodeDur time.Duration @@ -31,8 +43,9 @@ func tryBuildDirectH264Frame(display int) (wire.Frame, time.Duration, time.Durat if len(data) == 0 { return wire.Frame{}, captureDur, encodeDur, true, nil } - now := time.Now() - lastKeyframe.Store(now.UnixNano()) + if forceKeyframe { + lastKeyframe.Store(now.UnixNano()) + } statFullFrames.Add(1) return wire.Frame{Type: "frame", Header: wire.FrameHeader{Monitor: display, FPS: 0, Format: codec, Width: width, Height: height}, Data: data}, captureDur, encodeDur, true, nil } diff --git a/Overlord-Client/cmd/agent/capture/win_duplication_test.go b/Overlord-Client/cmd/agent/capture/win_duplication_test.go index ee396a3..6e96073 100644 --- a/Overlord-Client/cmd/agent/capture/win_duplication_test.go +++ b/Overlord-Client/cmd/agent/capture/win_duplication_test.go @@ -90,3 +90,28 @@ func TestCaptureDisplayDXGI(t *testing.T) { t.Fatalf("invalid capture bounds: %dx%d", bounds.Dx(), bounds.Dy()) } } + +func TestDirectVideoKeyframeDue(t *testing.T) { + now := (10 * time.Second).Nanoseconds() + last := now - keyframeEvery.Nanoseconds() + + tests := []struct { + name string + requested bool + last int64 + want bool + }{ + {name: "explicit request", requested: true, last: now, want: true}, + {name: "first frame", last: 0, want: true}, + {name: "before periodic interval", last: last + 1, want: false}, + {name: "at periodic interval", last: last, want: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := directVideoKeyframeDue(test.requested, now, test.last); got != test.want { + t.Fatalf("directVideoKeyframeDue() = %v, want %v", got, test.want) + } + }) + } +} diff --git a/Overlord-Client/cmd/agent/handlers/screenshot.go b/Overlord-Client/cmd/agent/handlers/screenshot.go index 9f050d5..c0a40fb 100644 --- a/Overlord-Client/cmd/agent/handlers/screenshot.go +++ b/Overlord-Client/cmd/agent/handlers/screenshot.go @@ -18,15 +18,13 @@ import ( ) var ( - monitorCountFn = capture.MonitorCount - displayBoundsFn = capture.DisplayBounds - captureDisplayRGBABitBltFn = capture.CaptureDisplayRGBABitBlt + monitorCountFn = capture.MonitorCount + displayBoundsFn = capture.DisplayBounds + captureRectFn = screenshot.CaptureRect ) func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays bool) error { //garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2 - restore := capture.BypassResolutionCap() - defer restore() if allDisplays { log.Printf("screenshot: capturing all displays") @@ -86,8 +84,6 @@ func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplay return err } - capture.RequestDesktopH264Keyframe() - return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{ Type: "command_result", CommandID: cmdID, @@ -157,7 +153,8 @@ func captureScreenshotImageWindows(allDisplays bool) (*image.RGBA, int, image.Re } if !allDisplays { - img, err := captureDisplayRGBABitBltFn(0) + bounds := displayBoundsFn(0) + img, err := captureRectFn(bounds) if err != nil || img == nil { return nil, 0, image.Rectangle{}, err } @@ -174,7 +171,7 @@ func captureScreenshotImageWindows(allDisplays bool) (*image.RGBA, int, image.Re for i := 0; i < monitorCount; i++ { bounds := displayBoundsFn(i) - img, err := captureDisplayRGBABitBltFn(i) + img, err := captureRectFn(bounds) if err != nil || img == nil { return nil, 0, image.Rectangle{}, err } diff --git a/Overlord-Client/cmd/agent/handlers/screenshot_windows_test.go b/Overlord-Client/cmd/agent/handlers/screenshot_windows_test.go index b76d4f1..235bafb 100644 --- a/Overlord-Client/cmd/agent/handlers/screenshot_windows_test.go +++ b/Overlord-Client/cmd/agent/handlers/screenshot_windows_test.go @@ -16,14 +16,17 @@ import ( func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) { origMonitorCountFn := monitorCountFn - origCaptureFn := captureDisplayRGBABitBltFn + origBoundsFn := displayBoundsFn + origCaptureFn := captureRectFn t.Cleanup(func() { monitorCountFn = origMonitorCountFn - captureDisplayRGBABitBltFn = origCaptureFn + displayBoundsFn = origBoundsFn + captureRectFn = origCaptureFn }) monitorCountFn = func() int { return 2 } - captureDisplayRGBABitBltFn = func(display int) (*image.RGBA, error) { + displayBoundsFn = func(int) image.Rectangle { return image.Rect(0, 0, 3, 2) } + captureRectFn = func(image.Rectangle) (*image.RGBA, error) { img := image.NewRGBA(image.Rect(0, 0, 3, 2)) img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255}) return img, nil @@ -47,11 +50,11 @@ func TestCaptureScreenshotImageWindows_PrimaryOnly(t *testing.T) { func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) { origMonitorCountFn := monitorCountFn origBoundsFn := displayBoundsFn - origCaptureFn := captureDisplayRGBABitBltFn + origCaptureFn := captureRectFn t.Cleanup(func() { monitorCountFn = origMonitorCountFn displayBoundsFn = origBoundsFn - captureDisplayRGBABitBltFn = origCaptureFn + captureRectFn = origCaptureFn }) monitorCountFn = func() int { return 2 } @@ -61,9 +64,9 @@ func TestCaptureScreenshotImageWindows_AllDisplaysStitched(t *testing.T) { } return image.Rect(2, 0, 4, 2) } - captureDisplayRGBABitBltFn = func(display int) (*image.RGBA, error) { + captureRectFn = func(bounds image.Rectangle) (*image.RGBA, error) { img := image.NewRGBA(image.Rect(0, 0, 2, 2)) - if display == 0 { + if bounds.Min.X == 0 { img.SetRGBA(0, 0, color.RGBA{R: 255, A: 255}) } else { img.SetRGBA(0, 0, color.RGBA{G: 255, A: 255}) @@ -112,16 +115,19 @@ func TestCaptureScreenshotImageWindows_NoDisplays(t *testing.T) { } } -func TestHandleScreenshot_DoesNotEmitLiveFrame(t *testing.T) { +func TestHandleScreenshot_RepeatedRequestsDoNotEmitLiveFrames(t *testing.T) { origMonitorCountFn := monitorCountFn - origCaptureFn := captureDisplayRGBABitBltFn + origBoundsFn := displayBoundsFn + origCaptureFn := captureRectFn t.Cleanup(func() { monitorCountFn = origMonitorCountFn - captureDisplayRGBABitBltFn = origCaptureFn + displayBoundsFn = origBoundsFn + captureRectFn = origCaptureFn }) monitorCountFn = func() int { return 1 } - captureDisplayRGBABitBltFn = func(display int) (*image.RGBA, error) { + displayBoundsFn = func(int) image.Rectangle { return image.Rect(0, 0, 2, 2) } + captureRectFn = func(image.Rectangle) (*image.RGBA, error) { img := image.NewRGBA(image.Rect(0, 0, 2, 2)) img.SetRGBA(0, 0, color.RGBA{B: 255, A: 255}) return img, nil @@ -129,26 +135,30 @@ func TestHandleScreenshot_DoesNotEmitLiveFrame(t *testing.T) { ctx := context.Background() writer := &testWriter{} - if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, "cmd-1", false); err != nil { - t.Fatalf("HandleScreenshot failed: %v", err) + for i, commandID := range []string{"cmd-1", "cmd-2"} { + if err := HandleScreenshot(ctx, &rt.Env{Conn: writer}, commandID, false); err != nil { + t.Fatalf("HandleScreenshot request %d failed: %v", i+1, err) + } } - if len(writer.msgs) != 2 { - t.Fatalf("expected 2 messages without live frame emission, got %d", len(writer.msgs)) + if len(writer.msgs) != 4 { + t.Fatalf("expected 4 result messages without live frame emission, got %d", len(writer.msgs)) } - var first wire.ScreenshotResult - if err := msgpack.Unmarshal(writer.msgs[0], &first); err != nil { - t.Fatalf("unmarshal screenshot_result: %v", err) - } - if first.Type != "screenshot_result" { - t.Fatalf("expected screenshot_result, got %q", first.Type) - } + for i := range 2 { + var screenshotResult wire.ScreenshotResult + if err := msgpack.Unmarshal(writer.msgs[i*2], &screenshotResult); err != nil { + t.Fatalf("unmarshal screenshot_result %d: %v", i+1, err) + } + if screenshotResult.Type != "screenshot_result" { + t.Fatalf("expected screenshot_result for request %d, got %q", i+1, screenshotResult.Type) + } - var second wire.CommandResult - if err := msgpack.Unmarshal(writer.msgs[1], &second); err != nil { - t.Fatalf("unmarshal command_result: %v", err) - } - if second.Type != "command_result" || !second.OK { - t.Fatalf("expected successful command_result, got %+v", second) + var commandResult wire.CommandResult + if err := msgpack.Unmarshal(writer.msgs[i*2+1], &commandResult); err != nil { + t.Fatalf("unmarshal command_result %d: %v", i+1, err) + } + if commandResult.Type != "command_result" || !commandResult.OK { + t.Fatalf("expected successful command_result for request %d, got %+v", i+1, commandResult) + } } } diff --git a/Overlord-Server/public/assets/remotedesktop.js b/Overlord-Server/public/assets/remotedesktop.js index 91b4e07..7a1aa94 100644 --- a/Overlord-Server/public/assets/remotedesktop.js +++ b/Overlord-Server/public/assets/remotedesktop.js @@ -68,6 +68,7 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s const streamProfileSelect = document.getElementById("streamProfileSelect"); const streamProfileDetail = document.getElementById("streamProfileDetail"); const bitrateSelect = document.getElementById("bitrateSelect"); + const requestKeyframeBtn = document.getElementById("requestKeyframeBtn"); const smoothingSlider = document.getElementById("smoothingSlider"); const smoothingValue = document.getElementById("smoothingValue"); const qualitySlider = document.getElementById("qualitySlider"); @@ -143,6 +144,7 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s let h264KeyframeErrorStreak = 0; let h264RecoveryAttempts = 0; let h264LastDecodeWarnAt = 0; + let lastAutomaticKeyframeAt = 0; const H264_LOW_FPS_THRESHOLD = 6; const H264_FALLBACK_WARMUP_MS = 10000; const H264_MIN_FRAMES_BEFORE_FALLBACK = 120; @@ -150,6 +152,8 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s const H264_KEYFRAME_ERROR_RESTART_THRESHOLD = 24; const H264_MAX_RECOVERY_ATTEMPTS = 1; const H264_DECODE_WARN_THROTTLE_MS = 2000; + const VIDEO_FRAME_GAP_KEYFRAME_MS = 500; + const AUTOMATIC_KEYFRAME_COOLDOWN_MS = 2000; const inputBackpressureBytes = 256 * 1024; const diagnostics = { agent: null, @@ -621,6 +625,9 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s if (stopBtn) { stopBtn.disabled = !wsOpen || (!isStarting && !isStreaming && !isStopping && !isStalled); } + if (requestKeyframeBtn) { + requestKeyframeBtn.disabled = !wsOpen || (!isStreaming && !isStalled); + } if (recordBtn) { recordBtn.disabled = !recordingActive && @@ -1733,6 +1740,13 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s sharedSettingsSaver.scheduleSave(); }); } + if (requestKeyframeBtn) { + requestKeyframeBtn.addEventListener("click", function () { + if (streamState !== "streaming" && streamState !== "stalled") return; + sendCmd("desktop_request_keyframe", { reason: "manual_viewer" }); + }); + } + function needsWebrtcFirewallWarning(mode) { if (firewallWarningAcked) return false; @@ -2012,8 +2026,32 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s return buf.length >= 8 && buf[0] === 0x46 && buf[1] === 0x52 && buf[2] === 0x4d; } + function requestAutomaticKeyframe(reason, details = {}) { + if (!desiredStreaming || getWebrtcMode() !== "off") return false; + if (streamState !== "streaming" && streamState !== "stalled" && streamState !== "starting") { + return false; + } + const codec = String(diagnostics.codec || negotiatedCodec || "").toLowerCase(); + if (codec !== "h264" && codec !== "hevc") return false; + if (!ws || ws.readyState !== WebSocket.OPEN) return false; + + const now = performance.now(); + if (lastAutomaticKeyframeAt && now - lastAutomaticKeyframeAt < AUTOMATIC_KEYFRAME_COOLDOWN_MS) { + return false; + } + lastAutomaticKeyframeAt = now; + rdDebug("automatic keyframe request", { reason, codec, ...details }); + sendCmd("desktop_request_keyframe", { reason, codec }); + return true; + } + function markFrameReceived() { - lastFrameAt = performance.now(); + const now = performance.now(); + const frameGapMs = lastFrameAt ? now - lastFrameAt : 0; + lastFrameAt = now; + if (frameGapMs >= VIDEO_FRAME_GAP_KEYFRAME_MS) { + requestAutomaticKeyframe("viewer_frame_gap", { frameGapMs: Math.round(frameGapMs) }); + } clearOfflineTimer(); if (!firstFrameLogged) { firstFrameLogged = true; @@ -2512,17 +2550,23 @@ import { createSharedUiSettingsSaver, loadSharedUiSettings } from "./shared-ui-s h264KeyframeErrorStreak = 0; } catch (err) { h264PendingTimings.pop(); - if (codecName === "h264" && isKeyframeRequiredError(err)) { + if (isKeyframeRequiredError(err)) { h264KeyframeErrorStreak += 1; + requestAutomaticKeyframe(`${codecName}_decoder_keyframe_required`, { + streak: h264KeyframeErrorStreak, + }); const now = Date.now(); if (now - h264LastDecodeWarnAt >= H264_DECODE_WARN_THROTTLE_MS) { h264LastDecodeWarnAt = now; - console.warn("rd: h264 decode waiting for keyframe", { + console.warn(`rd: ${codecName} decode waiting for keyframe`, { streak: h264KeyframeErrorStreak, recoveries: h264RecoveryAttempts, }); } - if (h264KeyframeErrorStreak >= H264_KEYFRAME_ERROR_RESTART_THRESHOLD) { + if ( + codecName === "h264" && + h264KeyframeErrorStreak >= H264_KEYFRAME_ERROR_RESTART_THRESHOLD + ) { const restarted = tryRecoverH264Stream("h264_keyframe_required"); if (!restarted) { fallbackToJpegCodec("h264_keyframe_required_loop"); diff --git a/Overlord-Server/public/remotedesktop.html b/Overlord-Server/public/remotedesktop.html index 7279d09..1291cea 100644 --- a/Overlord-Server/public/remotedesktop.html +++ b/Overlord-Server/public/remotedesktop.html @@ -219,6 +219,15 @@ +
diff --git a/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts b/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts index b202bf8..a70aa3a 100644 --- a/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts +++ b/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts @@ -17,13 +17,14 @@ import { decodeMessage, encodeMessage, type WireMessage, type Hello, type Ping } import * as sessionManager from "../../sessions/sessionManager"; import type { SocketData } from "../../sessions/types"; import type { ClientInfo } from "../../types"; -import { clearClientSyncState, handleFrame, handleHello, handlePing, handlePong, handleScreenshotThumbnailResult } from "../../wsHandlers"; +import { clearClientSyncState, handleFrame, handleHello, handlePing, handlePong, handleScreenshotThumbnailResult, shouldRelayFrameToViewers } from "../../wsHandlers"; import { queueClientDbUpdate, scheduleQueuedClientDbFlush } from "../../client-db-sync"; import { getMaxPayloadLimit, getMessageByteLength, isAllowedClientMessageType } from "../../wsValidation"; import { stopAllProxiesForClient } from "../socks5-proxy-manager"; import { verifyBuildToken, isBuildBanned } from "../build-signing"; import { clearThumbnail } from "../../thumbnails"; import { stopRemoteDesktopRecording } from "../rd-recording"; +import { requestRemoteDesktopKeyframeAfterScreenshot } from "../ws-console-rd-backstage"; import { isAuthenticatedViewerRole, registerViewerSocket, @@ -795,13 +796,14 @@ export async function handleWebSocketMessage( } } } - if (handleFrame(client, payload)) { + if (handleFrame(client, payload, shouldRelayFrameToViewers(payload))) { try { ws.send(encodeMessage({ type: "frame_ack" })); } catch {} } break; case "screenshot_result": handleScreenshotThumbnailResult(client, payload); deps.handleNotificationScreenshotResult(client.id, payload); + requestRemoteDesktopKeyframeAfterScreenshot(client.id); break; case "console_output": deps.handleConsoleOutput(client.id, payload); diff --git a/Overlord-Server/src/server/ws-console-rd-backstage.test.ts b/Overlord-Server/src/server/ws-console-rd-backstage.test.ts index 5572d03..e274869 100644 --- a/Overlord-Server/src/server/ws-console-rd-backstage.test.ts +++ b/Overlord-Server/src/server/ws-console-rd-backstage.test.ts @@ -14,6 +14,7 @@ import { handleDesktopEncoderCapabilities, backstageStreamingState, rdStreamingState, + requestRemoteDesktopKeyframeAfterScreenshot, } from "./ws-console-rd-backstage"; type MockWs = { @@ -118,6 +119,72 @@ describe("remote desktop viewer control", () => { expect(rdStreamingState.get(clientId)?.isStreaming).toBe(false); }); + test("forwards keyframe requests only while the desktop stream is active", () => { + const clientId = `rd-keyframe-${Date.now().toString(36)}`; + const { agentWs } = createClient(clientId); + const viewer = createMockWs({ clientId }); + const viewerSocket = viewer as unknown as Parameters[0]; + + handleRemoteDesktopViewerOpen(viewerSocket); + handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ + type: "desktop_request_keyframe", + reason: "manual_viewer", + })); + + expect(agentCommands(agentWs).filter((msg) => msg.commandType === "desktop_request_keyframe")).toHaveLength(0); + + handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ type: "desktop_start" })); + handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ + type: "desktop_request_keyframe", + reason: "manual_viewer", + })); + + const keyframeCommands = agentCommands(agentWs) + .filter((msg) => msg.commandType === "desktop_request_keyframe"); + expect(keyframeCommands).toHaveLength(1); + expect(keyframeCommands[0]?.payload).toEqual({ reason: "manual_viewer" }); + + handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ + type: "desktop_request_keyframe", + reason: "viewer_frame_gap", + })); + const automaticCommand = agentCommands(agentWs) + .filter((msg) => msg.commandType === "desktop_request_keyframe") + .at(-1); + expect(automaticCommand?.payload).toEqual({ reason: "viewer_frame_gap" }); + }); + + test("requests HEVC recovery after a screenshot only while streaming", () => { + const clientId = `rd-screenshot-keyframe-${Date.now().toString(36)}`; + const { agentWs } = createClient(clientId); + const viewer = createMockWs({ clientId }); + const viewerSocket = viewer as unknown as Parameters[0]; + + handleRemoteDesktopViewerOpen(viewerSocket); + expect(requestRemoteDesktopKeyframeAfterScreenshot(clientId)).toBe(false); + + handleRemoteDesktopViewerMessage(viewerSocket, JSON.stringify({ type: "desktop_start" })); + const state = rdStreamingState.get(clientId); + expect(state).toBeDefined(); + if (!state) return; + state.codec = "hevc"; + agentWs.sent.length = 0; + + expect(requestRemoteDesktopKeyframeAfterScreenshot(clientId)).toBe(true); + expect(agentCommands(agentWs) + .filter((msg) => msg.commandType === "desktop_request_keyframe")) + .toEqual([ + expect.objectContaining({ + payload: { reason: "post_screenshot_hevc_recovery" }, + }), + ]); + + state.codec = "h264"; + agentWs.sent.length = 0; + expect(requestRemoteDesktopKeyframeAfterScreenshot(clientId)).toBe(false); + expect(agentWs.sent).toHaveLength(0); + }); + test("does not forward desktop_start when a macOS client is missing required permissions", () => { const clientId = `rd-mac-perms-${Date.now().toString(36)}`; const { info, agentWs } = createClient(clientId); diff --git a/Overlord-Server/src/server/ws-console-rd-backstage.ts b/Overlord-Server/src/server/ws-console-rd-backstage.ts index fe726e2..0c5bb71 100644 --- a/Overlord-Server/src/server/ws-console-rd-backstage.ts +++ b/Overlord-Server/src/server/ws-console-rd-backstage.ts @@ -178,6 +178,11 @@ function broadcastFrameToViewers( const rdSendStats = { lastLog: 0, frames: 0, sendMs: 0, bytes: 0 }; const rdDebugFrameLogAt = new Map(); const rdCanvasFrameAckPending = new Map(); +const AUTOMATIC_DESKTOP_KEYFRAME_REASONS: Record = { + viewer_frame_gap: true, + h264_decoder_keyframe_required: true, + hevc_decoder_keyframe_required: true, +}; export const rdStreamingState = new Map, raw: string | ArrayBuffer | Uint8Array) { const payload = decodeViewerPayload(raw); if (!payload) return; @@ -489,6 +506,15 @@ export function handleRemoteDesktopViewerMessage(ws: ServerWebSocket safeSendViewer(ws, { type: "status", status: "stopped" }); break; } + case "desktop_request_keyframe": { + if (state.isStreaming) { + const requestedReason = String(payload.reason || ""); + sendDesktopCommand(target, "desktop_request_keyframe", { + reason: AUTOMATIC_DESKTOP_KEYFRAME_REASONS[requestedReason] ? requestedReason : "manual_viewer", + }); + } + break; + } case "desktop_record_start": { if (!state.isStreaming) { safeSendViewer(ws, { diff --git a/Overlord-Server/src/wsHandlers.test.ts b/Overlord-Server/src/wsHandlers.test.ts index 8912223..2201aa6 100644 --- a/Overlord-Server/src/wsHandlers.test.ts +++ b/Overlord-Server/src/wsHandlers.test.ts @@ -1,7 +1,8 @@ import { describe, expect, test } from "bun:test"; +import type { ClientInfo } from "./types"; import { decodeMessage } from "./protocol"; import { metrics } from "./metrics"; -import { handlePing, handlePong, sendPingRequest } from "./wsHandlers"; +import { handleFrame, handlePing, handlePong, sendPingRequest, shouldRelayFrameToViewers } from "./wsHandlers"; type MockWs = { sent: Uint8Array[]; send: (msg: Uint8Array) => void }; @@ -178,3 +179,40 @@ describe("wsHandlers ping/pong", () => { expect(snapshot.ping.count).toBe(0); }); }); + +describe("wsHandlers frame routing", () => { + test("keeps screenshot frames out of remote desktop viewers", () => { + const broadcasts: unknown[] = []; + const runtime = globalThis as unknown as { + __rdBroadcast?: (clientId: string, bytes: Uint8Array, header: unknown) => boolean; + }; + runtime.__rdBroadcast = (clientId, bytes, header) => { + broadcasts.push({ clientId, bytes, header }); + return true; + }; + const client = { + id: "screenshot-client", + lastSeen: Date.now(), + online: true, + isAdmin: false, + } as unknown as ClientInfo; + const screenshotFrame = { + type: "frame", + header: { fps: 0, format: "h264", width: 1920, height: 1080 }, + data: new Uint8Array([1, 2, 3]), + }; + + try { + expect(shouldRelayFrameToViewers(screenshotFrame)).toBe(false); + handleFrame(client, screenshotFrame, shouldRelayFrameToViewers(screenshotFrame)); + expect(broadcasts).toHaveLength(0); + + const liveFrame = { ...screenshotFrame, header: { ...screenshotFrame.header, fps: 60 } }; + expect(shouldRelayFrameToViewers(liveFrame)).toBe(true); + handleFrame(client, liveFrame, shouldRelayFrameToViewers(liveFrame)); + expect(broadcasts).toHaveLength(1); + } finally { + delete runtime.__rdBroadcast; + } + }); +}); diff --git a/Overlord-Server/src/wsHandlers.ts b/Overlord-Server/src/wsHandlers.ts index eff8f97..bf99387 100644 --- a/Overlord-Server/src/wsHandlers.ts +++ b/Overlord-Server/src/wsHandlers.ts @@ -273,30 +273,49 @@ export function handleScreenshotThumbnailResult(info: ClientInfo, payload: any): return true; } -export function handleFrame(info: ClientInfo, payload: any): boolean { - const bytes = payload.data as unknown as Uint8Array; - const header = (payload as any).header; +export function shouldRelayFrameToViewers(payload: unknown): boolean { + if (!payload || typeof payload !== "object" || !("header" in payload)) return false; + const header = payload.header; + if (!header || typeof header !== "object" || !("fps" in header)) return false; + return Number(header.fps) > 0; +} + +export function handleFrame(info: ClientInfo, payload: unknown, relayToViewers = true): boolean { + const frame = payload && typeof payload === "object" + ? payload as Record + : {}; + const rawData = frame.data; + const bytes = rawData instanceof Uint8Array + ? rawData + : rawData instanceof ArrayBuffer + ? new Uint8Array(rawData) + : new Uint8Array(); + const header = frame.header && typeof frame.header === "object" + ? frame.header as Record + : {}; const allowedFormats = ["jpeg", "jpg", "webp"]; - const fmt = String(header?.format || "").toLowerCase(); + const fmt = String(header.format || "").toLowerCase(); const safeFormat = allowedFormats.includes(fmt) ? fmt : ""; metrics.recordBytesReceived(bytes.length); let handledByViewerRelay = false; - try { - const globalAny: any = globalThis as any; - if (header?.webcam) { - if (globalAny.__webcamBroadcast) { - handledByViewerRelay = globalAny.__webcamBroadcast(info.id, bytes, header); + if (relayToViewers) { + try { + const runtime = globalThis as unknown as { + __webcamBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record) => boolean; + __backstageBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record) => boolean; + __rdBroadcast?: (clientId: string, data: Uint8Array, frameHeader: Record) => boolean; + }; + if (header.webcam) { + handledByViewerRelay = runtime.__webcamBroadcast?.(info.id, bytes, header) ?? false; + } else if (header.backstage) { + handledByViewerRelay = runtime.__backstageBroadcast?.(info.id, bytes, header) ?? false; + } else { + handledByViewerRelay = runtime.__rdBroadcast?.(info.id, bytes, header) ?? false; } - } else if (header?.backstage) { - if (globalAny.__backstageBroadcast) { - handledByViewerRelay = globalAny.__backstageBroadcast(info.id, bytes, header); - } - } else if (globalAny.__rdBroadcast) { - handledByViewerRelay = globalAny.__rdBroadcast(info.id, bytes, header); - } - } catch {} + } catch {} + } if (safeFormat) { const now = Date.now(); diff --git a/Overlord-Server/test/e2e/remotedesktop-input.e2e.ts b/Overlord-Server/test/e2e/remotedesktop-input.e2e.ts index 3ac79f7..6618fcf 100644 --- a/Overlord-Server/test/e2e/remotedesktop-input.e2e.ts +++ b/Overlord-Server/test/e2e/remotedesktop-input.e2e.ts @@ -64,6 +64,7 @@ test("WebRTC input stays inactive until start and then forwards mouse and keyboa }); await expect(page.locator("#streamProfileSelect option[value='1440:60']")).toHaveCount(1); await expect(page.locator("#streamProfileSelect option[value='2160:60']")).toHaveCount(1); + await expect(page.locator("#requestKeyframeBtn")).toBeDisabled(); await page.evaluate(() => { const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent; @@ -198,4 +199,37 @@ test("WebRTC input stays inactive until start and then forwards mouse and keyboa expect.objectContaining({ type: "key_down", key: "Enter", code: "Enter" }), expect.objectContaining({ type: "key_up", key: "Enter", code: "Enter" }), ])); + + await page.evaluate(() => { + const socket = (window as typeof window & { __rdSocket: EventTarget }).__rdSocket; + const framePacket = Uint8Array.from([0x46, 0x52, 0x4d, 0, 0, 0, 0, 0]).buffer; + socket.dispatchEvent(new MessageEvent("message", { data: framePacket })); + }); + await expect(page.locator("#requestKeyframeBtn")).toBeEnabled(); + await page.waitForTimeout(550); + await page.evaluate(() => { + const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent; + sent.length = 0; + const socket = (window as typeof window & { __rdSocket: EventTarget }).__rdSocket; + const framePacket = Uint8Array.from([0x46, 0x52, 0x4d, 0, 0, 0, 0, 0]).buffer; + socket.dispatchEvent(new MessageEvent("message", { data: framePacket })); + }); + await expect.poll(async () => page.evaluate(async () => { + const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent; + const { decodeMsgpack } = await import("/assets/msgpack-helpers.js"); + return sent.map((message) => decodeMsgpack(message)) + .filter((message) => message.type === "desktop_request_keyframe"); + })).toEqual([ + expect.objectContaining({ type: "desktop_request_keyframe", reason: "viewer_frame_gap" }), + ]); + await page.locator("#rdSettingsBtn").click(); + await page.locator("#requestKeyframeBtn").click(); + await expect.poll(async () => page.evaluate(async () => { + const sent = (window as typeof window & { __rdSent: ArrayBuffer[] }).__rdSent; + const { decodeMsgpack } = await import("/assets/msgpack-helpers.js"); + return sent.map((message) => decodeMsgpack(message)) + .filter((message) => message.type === "desktop_request_keyframe" && message.reason === "manual_viewer"); + })).toEqual([ + expect.objectContaining({ type: "desktop_request_keyframe", reason: "manual_viewer" }), + ]); });