From 98d533c8e848a1da03de198c9249770be3297bc1 Mon Sep 17 00:00:00 2001 From: Viktor Liu Date: Thu, 21 May 2026 12:01:25 +0200 Subject: [PATCH] Address CodeRabbit feedback on VNC server agent matching and session lifecycle --- client/vnc/server/agent_darwin.go | 34 ++++++------- ...t_uinput_unix.go => input_uinput_linux.go} | 2 +- client/vnc/server/input_windows.go | 6 +++ client/vnc/server/scancodes.go | 2 +- client/vnc/server/server.go | 49 ++++++++++--------- client/vnc/server/virtual_x11.go | 15 +++++- 6 files changed, 64 insertions(+), 44 deletions(-) rename client/vnc/server/{input_uinput_unix.go => input_uinput_linux.go} (99%) diff --git a/client/vnc/server/agent_darwin.go b/client/vnc/server/agent_darwin.go index 31adad5b9..da8083bcf 100644 --- a/client/vnc/server/agent_darwin.go +++ b/client/vnc/server/agent_darwin.go @@ -10,7 +10,6 @@ import ( "net" "os" "os/exec" - "path/filepath" "strconv" "sync" "syscall" @@ -253,15 +252,19 @@ func killAllVNCAgents() { } // vncAgentPIDs returns the pids of vnc-agent subprocesses spawned from -// this binary. Matches on (argv[0] basename == our own basename) AND -// argv contains the "vnc-agent" subcommand. Skips pid 0 and 1 defensively. +// this binary. Matches exactly on argv[0] == our own executable path +// AND argv[1] == "vnc-agent" so unrelated processes that happen to have +// the same name elsewhere in argv are not targeted. Skips pid 0 and 1 +// defensively. func vncAgentPIDs() ([]int, error) { procs, err := unix.SysctlKinfoProcSlice("kern.proc.all") if err != nil { return nil, fmt.Errorf("sysctl kern.proc.all: %w", err) } - ownExe, _ := os.Executable() - ownBase := filepath.Base(ownExe) + ownExe, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("resolve own executable: %w", err) + } var out []int for i := range procs { pid := int(procs[i].Proc.P_pid) @@ -269,7 +272,7 @@ func vncAgentPIDs() ([]int, error) { continue } argv, err := procArgv(pid) - if err != nil || !argvIsVNCAgent(argv, ownBase) { + if err != nil || !argvIsVNCAgent(argv, ownExe) { continue } out = append(out, pid) @@ -313,19 +316,12 @@ func procArgv(pid int) ([]string, error) { } // argvIsVNCAgent reports whether argv belongs to a vnc-agent subprocess -// spawned from our binary. Requires argv[0]'s basename to match ownBase -// and the "vnc-agent" subcommand to appear among the positional args. -func argvIsVNCAgent(argv []string, ownBase string) bool { - if len(argv) < 2 || ownBase == "" { +// spawned from our binary. Requires argv[0] to match ownExe exactly and +// argv[1] to be the vnc-agent subcommand. Matches the spawn shape in +// spawnAgentForUser and rejects anything else. +func argvIsVNCAgent(argv []string, ownExe string) bool { + if len(argv) < 2 || ownExe == "" { return false } - if filepath.Base(argv[0]) != ownBase { - return false - } - for _, a := range argv[1:] { - if a == vncAgentSubcommand { - return true - } - } - return false + return argv[0] == ownExe && argv[1] == vncAgentSubcommand } diff --git a/client/vnc/server/input_uinput_unix.go b/client/vnc/server/input_uinput_linux.go similarity index 99% rename from client/vnc/server/input_uinput_unix.go rename to client/vnc/server/input_uinput_linux.go index 86f6f0148..505fbc880 100644 --- a/client/vnc/server/input_uinput_unix.go +++ b/client/vnc/server/input_uinput_linux.go @@ -1,4 +1,4 @@ -//go:build unix && !darwin && !ios && !android +//go:build linux package server diff --git a/client/vnc/server/input_windows.go b/client/vnc/server/input_windows.go index cf18d96a3..cf3e2505e 100644 --- a/client/vnc/server/input_windows.go +++ b/client/vnc/server/input_windows.go @@ -482,6 +482,7 @@ var ( procGlobalAlloc = kernel32.NewProc("GlobalAlloc") procGlobalLock = kernel32.NewProc("GlobalLock") procGlobalUnlock = kernel32.NewProc("GlobalUnlock") + procGlobalFree = kernel32.NewProc("GlobalFree") ) const ( @@ -522,6 +523,7 @@ func (w *WindowsInputInjector) doSetClipboard(text string) { ptr, _, _ := procGlobalLock.Call(hMem) if ptr == 0 { log.Tracef("GlobalLock for clipboard: lock returned nil") + _, _, _ = procGlobalFree.Call(hMem) return } copy(unsafe.Slice((*uint16)(unsafe.Pointer(ptr)), len(utf16)), utf16) @@ -530,6 +532,7 @@ func (w *WindowsInputInjector) doSetClipboard(text string) { r, _, lerr := procOpenClipboard.Call(0) if r == 0 { log.Tracef("OpenClipboard: %v", lerr) + _, _, _ = procGlobalFree.Call(hMem) return } defer logCleanupCall("CloseClipboard", procCloseClipboard) @@ -538,6 +541,9 @@ func (w *WindowsInputInjector) doSetClipboard(text string) { r, _, lerr = procSetClipboardData.Call(cfUnicodeText, hMem) if r == 0 { log.Tracef("SetClipboardData: %v", lerr) + // Ownership only transfers to the OS on success; on failure we + // still own hMem and must free it. + _, _, _ = procGlobalFree.Call(hMem) } } diff --git a/client/vnc/server/scancodes.go b/client/vnc/server/scancodes.go index 54db42a63..9ae204805 100644 --- a/client/vnc/server/scancodes.go +++ b/client/vnc/server/scancodes.go @@ -23,7 +23,7 @@ package server // // Linux KEY_* codes. Only the ones we reference, since the full // linux/input-event-codes.h list isn't useful here. Naming mirrors the -// existing constants in input_uinput_unix.go (mixed case, no underscores). +// existing constants in input_uinput_linux.go (mixed case, no underscores). const ( keyEsc = 1 key1 = 2 diff --git a/client/vnc/server/server.go b/client/vnc/server/server.go index 3be5044cd..7f3b98816 100644 --- a/client/vnc/server/server.go +++ b/client/vnc/server/server.go @@ -418,11 +418,15 @@ func (s *Server) Stop() error { s.cancel = nil } - // Close active client connections before tearing down capturers and - // listeners. The per-session serve goroutines unblock from their Read - // loop with an error and run their deferred conn.Close, which surfaces - // a clean disconnect on the client side instead of leaving the - // connection hanging until the OS reclaims it on process exit. + // Close the listener first so the accept loop exits and cannot + // register any further connections in acceptedConns. Then close every + // already-accepted connection so per-session serve goroutines unblock + // and run their deferred conn.Close. + var listenerErr error + if s.listener != nil { + listenerErr = s.listener.Close() + s.listener = nil + } s.closeActiveSessions() if s.vmgr != nil { @@ -437,12 +441,8 @@ func (s *Server) Stop() error { c.Close() } - if s.listener != nil { - err := s.listener.Close() - s.listener = nil - if err != nil { - return fmt.Errorf("close VNC listener: %w", err) - } + if listenerErr != nil { + return fmt.Errorf("close VNC listener: %w", listenerErr) } s.log.Info("stopped") @@ -894,7 +894,8 @@ func (s *Server) acquireSessionResources(conn net.Conn, header *connectionHeader case ModeSession: return s.acquireVirtualSession(conn, header, connLog) default: - return s.acquireAttachSession(), s.injector, attachSessionCleanup, true + capturer, cleanup := s.acquireAttachSession() + return capturer, s.injector, cleanup, true } } @@ -921,17 +922,21 @@ func (s *Server) acquireVirtualSession(conn net.Conn, header *connectionHeader, return vs.Capturer(), vs.Injector(), vs.ClientDisconnect, true } -func (s *Server) acquireAttachSession() ScreenCapturer { - if cc, ok := s.capturer.(interface{ ClientConnect() }); ok { - cc.ClientConnect() +// acquireAttachSession bumps the shared capturer's per-session refcount +// (if it implements the optional ClientConnect/ClientDisconnect pair) and +// returns a cleanup func that releases it. X11Poller and the Windows +// capturer rely on the disconnect path to drop SHM/DXGI resources when no +// client is active. +func (s *Server) acquireAttachSession() (ScreenCapturer, func()) { + type connectDisconnect interface { + ClientConnect() + ClientDisconnect() } - return s.capturer -} - -// attachSessionCleanup is the no-op cleanup used by attach mode. Returned as a -// named func rather than an inline closure so the empty body is unambiguous. -func attachSessionCleanup() { - // Attach mode keeps the shared capturer; nothing to release per session. + if cc, ok := s.capturer.(connectDisconnect); ok { + cc.ClientConnect() + return s.capturer, cc.ClientDisconnect + } + return s.capturer, func() {} } // modeString returns a human-readable session mode name. diff --git a/client/vnc/server/virtual_x11.go b/client/vnc/server/virtual_x11.go index 0a764fca5..9f8886bc2 100644 --- a/client/vnc/server/virtual_x11.go +++ b/client/vnc/server/virtual_x11.go @@ -176,10 +176,19 @@ func (vs *VirtualSession) ClientDisconnect() { // idleExpired is called by the idle timer. It stops the session and // notifies the session manager via onIdle so it removes us from the map. +// Bails out early if a client reconnected before the timer callback won +// the race (Stop() doesn't cancel an already-firing AfterFunc, so the +// state check has to happen here under vs.mu). func (vs *VirtualSession) idleExpired() { + vs.mu.Lock() + if vs.stopped || vs.clients > 0 { + vs.mu.Unlock() + return + } + vs.mu.Unlock() + vs.log.Info("idle timeout reached, destroying virtual session") vs.Stop() - // onIdle acquires sessionManager.mu; safe because Stop() has released vs.mu. if vs.onIdle != nil { vs.onIdle() } @@ -231,6 +240,10 @@ func (vs *VirtualSession) Stop() { if vs.injector != nil { vs.injector.Close() } + if vs.poller != nil { + vs.poller.Close() + vs.poller = nil + } vs.stopDesktop() vs.stopXvfb()