Address CodeRabbit feedback on VNC server agent matching and session lifecycle

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
 }

client/vnc/server/input_uinput_linux.go

@@ -1,4 +1,4 @@
-//go:build unix && !darwin && !ios && !android
+//go:build linux
 
 package server
 

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)
 	}
 }
 

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

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.

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()