Persist virtual session processes for crash recovery and identify them by start time

client/internal/engine_vnc.go

@@ -146,6 +146,10 @@ func (e *Engine) startVNCServer() error {
 		// snapshot ourselves; otherwise the UI's session list goes stale until an
 		// unrelated peer change happens to fire one.
 		OnSessionsChanged: e.statusRecorder.NotifyStateChange,
+		// Persist the X server and desktop PIDs so a daemon that dies without
+		// running Stop still cleans them up on the next start; otherwise they
+		// keep a display and a full desktop session alive indefinitely.
+		OnVirtualProcesses: e.persistVNCProcesses,
 	})
 
 	listenAddr := netip.AddrPortFrom(netbirdIP, vnc.InternalPort)
@@ -327,3 +331,14 @@ func displayPeer(info vncserver.ApprovalInfo) string {
 	}
 	return "unknown peer"
 }
+
+// persistVNCProcesses records the live virtual-session processes in the state
+// file. Best effort: a failure here costs crash recovery, not the session.
+func (e *Engine) persistVNCProcesses(state *vncserver.ShutdownState) {
+	if e.stateManager == nil {
+		return
+	}
+	if err := e.stateManager.UpdateState(state); err != nil {
+		log.Debugf("update VNC session state: %v", err)
+	}
+}

client/server/state_linux.go

@@ -9,6 +9,7 @@ import (
 	"github.com/netbirdio/netbird/client/internal/routemanager/systemops"
 	"github.com/netbirdio/netbird/client/internal/statemanager"
 	"github.com/netbirdio/netbird/client/ssh/config"
+	vncserver "github.com/netbirdio/netbird/client/vnc/server"
 )
 
 // registerStates registers all states that need crash recovery cleanup.
@@ -18,4 +19,7 @@ func registerStates(mgr *statemanager.Manager) {
 	mgr.RegisterState(&nftables.ShutdownState{})
 	mgr.RegisterState(&iptables.ShutdownState{})
 	mgr.RegisterState(&config.ShutdownState{})
+	// Virtual VNC sessions leave an X server and a desktop behind if the daemon
+	// dies without stopping them.
+	mgr.RegisterState(&vncserver.ShutdownState{})
 }

client/vnc/server/server.go

@@ -183,6 +183,9 @@ type Server struct {
 	ctx            context.Context
 	cancel         context.CancelFunc
 	vmgr           virtualSessionManager
+	// onVirtualProcesses forwards live virtual-session process records to the
+	// daemon for crash recovery; nil when nothing is listening.
+	onVirtualProcesses func(*ShutdownState)
 	// serviceAgentMu guards the shared per-session agent manager below, which
 	// every service-mode accept loop resolves through; see Server.serviceAgent.
 	// Its own mutex rather than mu: Stop holds mu while tearing it down.
@@ -320,6 +323,12 @@ type Config struct {
 	// Approver brokers the per-connection prompt to the local user via the
 	// daemon→UI event channel. Nil disables the gate.
 	Approver Approver
+
+	// OnVirtualProcesses, when set, is called with the current virtual-session
+	// process records whenever one starts or stops. The daemon persists them
+	// through the state manager so a crash does not leave an orphaned X server
+	// and desktop running for the life of the host.
+	OnVirtualProcesses func(*ShutdownState)
 }
 
 // Approver decouples the VNC server from the approval broker. A non-nil
@@ -354,24 +363,25 @@ type ApprovalInfo struct {
 // auth. The protocol-level VNC password scheme is not supported.
 func New(cfg Config) *Server {
 	s := &Server{
-		capturer:          cfg.Capturer,
-		injector:          cfg.Injector,
-		identityKey:       cfg.IdentityKey,
-		serviceMode:       cfg.ServiceMode,
-		sessionRecorder:   cfg.SessionRecorder,
-		requireApproval:   cfg.RequireApproval,
-		approver:          cfg.Approver,
-		disableAuth:       cfg.DisableAuth,
-		netstackNet:       cfg.NetstackNet,
-		preListener:       cfg.Listener,
-		authorizer:        sshauth.NewAuthorizer(),
-		log:               log.WithField("component", "vnc-server"),
-		sessions:          make(map[uint64]ActiveSessionInfo),
-		sessionConns:      make(map[uint64]net.Conn),
-		onSessionsChanged: cfg.OnSessionsChanged,
-		acceptedConns:     make(map[net.Conn]struct{}),
-		connAuth:          make(map[net.Conn]connAuthInfo),
-		connSem:           make(chan struct{}, maxConcurrentVNCConns),
+		capturer:           cfg.Capturer,
+		injector:           cfg.Injector,
+		identityKey:        cfg.IdentityKey,
+		serviceMode:        cfg.ServiceMode,
+		sessionRecorder:    cfg.SessionRecorder,
+		requireApproval:    cfg.RequireApproval,
+		approver:           cfg.Approver,
+		disableAuth:        cfg.DisableAuth,
+		netstackNet:        cfg.NetstackNet,
+		preListener:        cfg.Listener,
+		authorizer:         sshauth.NewAuthorizer(),
+		log:                log.WithField("component", "vnc-server"),
+		sessions:           make(map[uint64]ActiveSessionInfo),
+		sessionConns:       make(map[uint64]net.Conn),
+		onSessionsChanged:  cfg.OnSessionsChanged,
+		onVirtualProcesses: cfg.OnVirtualProcesses,
+		acceptedConns:      make(map[net.Conn]struct{}),
+		connAuth:           make(map[net.Conn]connAuthInfo),
+		connSem:            make(chan struct{}, maxConcurrentVNCConns),
 	}
 	if len(cfg.IdentityKey) == 32 {
 		pub, err := curve25519.X25519(cfg.IdentityKey, curve25519.Basepoint)

client/vnc/server/server_x11.go

@@ -15,7 +15,7 @@ func (s *Server) serviceAcceptLoop(ln net.Listener) {
 }
 
 func (s *Server) platformSessionManager() virtualSessionManager {
-	return newSessionManager(s.log)
+	return newSessionManager(s.log, s.onVirtualProcesses)
 }
 
 func (s *Server) platformShutdown() {

client/vnc/server/shutdown_state.go

@@ -3,19 +3,36 @@
 package server
 
 import (
+	"bytes"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 	"syscall"
 
 	log "github.com/sirupsen/logrus"
 )
 
+// sessionProcess identifies one virtual-session process well enough to be
+// signalled safely after a crash.
+//
+// A PID on its own is not enough: by the time the daemon restarts, the kernel
+// may have handed that number to something else, and Cleanup signals the whole
+// process group. Start time is what makes the identity stable — it is fixed for
+// the life of a process and a reused PID always has a later one — and the UID
+// keeps us from signalling another user's processes even if both matched.
+type sessionProcess struct {
+	PID int `json:"pid"`
+	// StartTime is field 22 of /proc/<pid>/stat, in clock ticks since boot.
+	StartTime uint64 `json:"startTime,omitempty"`
+	UID       uint32 `json:"uid,omitempty"`
+}
+
 // ShutdownState tracks VNC virtual session processes for crash recovery.
 // Persisted by the state manager; on restart, residual processes are killed.
 type ShutdownState struct {
-	// Processes maps a description to its PID (e.g., "xvfb:50" -> 1234).
-	Processes map[string]int `json:"processes,omitempty"`
+	// Processes maps a description to the process it names (e.g. "xvfb:50").
+	Processes map[string]sessionProcess `json:"processes,omitempty"`
 }
 
 // Name returns the state name for the state manager.
@@ -29,20 +46,20 @@ func (s *ShutdownState) Cleanup() error {
 		return nil
 	}
 
-	for desc, pid := range s.Processes {
-		if pid <= 0 {
+	for desc, proc := range s.Processes {
+		if proc.PID <= 0 {
 			continue
 		}
-		if !isOurProcess(pid, desc) {
-			log.Debugf("cleanup:skipping PID %d (%s), not ours", pid, desc)
+		if !isOurProcess(proc, desc) {
+			log.Debugf("cleanup: skipping PID %d (%s), not ours", proc.PID, desc)
 			continue
 		}
-		log.Infof("cleanup:killing residual process %d (%s)", pid, desc)
+		log.Infof("cleanup: killing residual process %d (%s)", proc.PID, desc)
 		// Kill the process group (negative PID) to get children too.
-		if err := syscall.Kill(-pid, syscall.SIGTERM); err != nil {
+		if err := syscall.Kill(-proc.PID, syscall.SIGTERM); err != nil {
 			// Try individual process if group kill fails.
-			if killErr := syscall.Kill(pid, syscall.SIGKILL); killErr != nil {
-				log.Debugf("cleanup: kill pid %d (%s): group kill: %v, single kill: %v", pid, desc, err, killErr)
+			if killErr := syscall.Kill(proc.PID, syscall.SIGKILL); killErr != nil {
+				log.Debugf("cleanup: kill pid %d (%s): group kill: %v, single kill: %v", proc.PID, desc, err, killErr)
 			}
 		}
 	}
@@ -51,20 +68,59 @@ func (s *ShutdownState) Cleanup() error {
 	return nil
 }
 
-// isOurProcess verifies the PID still belongs to a VNC-related process by
-// matching desc against /proc/<pid>/cmdline. A PID that no longer exists, or
-// whose cmdline cannot be read, is treated as foreign and reported false, so
-// cleanup never signals a process it cannot identify.
-func isOurProcess(pid int, desc string) bool {
+// describeProcess captures the identity of a freshly started process so a later
+// Cleanup can tell it apart from whatever inherits its PID.
+func describeProcess(pid int) sessionProcess {
+	proc := sessionProcess{PID: pid}
+	if start, err := processStartTime(pid); err == nil {
+		proc.StartTime = start
+	} else {
+		log.Debugf("read start time for pid %d: %v", pid, err)
+	}
+	if uid, err := processUID(pid); err == nil {
+		proc.UID = uid
+	} else {
+		log.Debugf("read uid for pid %d: %v", pid, err)
+	}
+	return proc
+}
+
+// isOurProcess verifies the PID still belongs to the VNC-related process it was
+// recorded for, by matching desc against /proc/<pid>/cmdline and confirming the
+// process start time and owner are the ones recorded. Anything that cannot be
+// read, or does not match, is reported as foreign so cleanup never signals a
+// process it has not identified.
+func isOurProcess(proc sessionProcess, desc string) bool {
 	// Check if the process exists at all.
-	if err := syscall.Kill(pid, 0); err != nil {
+	if err := syscall.Kill(proc.PID, 0); err != nil {
 		return false
 	}
 
-	// On Linux, verify via /proc cmdline.
-	cmdline, err := os.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
+	// A recorded start time that no longer matches means the PID was reused.
+	// A record without one predates the check and cannot be trusted to be the
+	// same process, so it is refused as well.
+	if proc.StartTime == 0 {
+		log.Debugf("cleanup: pid %d (%s) has no recorded start time", proc.PID, desc)
+		return false
+	}
+	start, err := processStartTime(proc.PID)
 	if err != nil {
-		log.Debugf("cleanup: cannot read /proc/%d/cmdline: %v, treating PID as foreign", pid, err)
+		log.Debugf("cleanup: cannot read start time for pid %d: %v, treating PID as foreign", proc.PID, err)
+		return false
+	}
+	if start != proc.StartTime {
+		log.Debugf("cleanup: pid %d (%s) started at %d, recorded %d: PID was reused", proc.PID, desc, start, proc.StartTime)
+		return false
+	}
+
+	if uid, err := processUID(proc.PID); err != nil || uid != proc.UID {
+		log.Debugf("cleanup: pid %d (%s) owner mismatch (err=%v): treating PID as foreign", proc.PID, desc, err)
+		return false
+	}
+
+	cmdline, err := os.ReadFile(fmt.Sprintf("/proc/%d/cmdline", proc.PID))
+	if err != nil {
+		log.Debugf("cleanup: cannot read /proc/%d/cmdline: %v, treating PID as foreign", proc.PID, err)
 		return false
 	}
 
@@ -80,3 +136,34 @@ func isOurProcess(pid int, desc string) bool {
 	}
 	return false
 }
+
+// processStartTime reads field 22 of /proc/<pid>/stat, the process start time in
+// clock ticks since boot. Parsed from the last ')' so a comm containing spaces
+// or parentheses cannot shift the field offsets.
+func processStartTime(pid int) (uint64, error) {
+	raw, err := os.ReadFile(fmt.Sprintf("/proc/%d/stat", pid))
+	if err != nil {
+		return 0, err
+	}
+	commEnd := bytes.LastIndexByte(raw, ')')
+	if commEnd < 0 {
+		return 0, fmt.Errorf("malformed /proc/%d/stat", pid)
+	}
+	// Fields after comm: state is field 3, so start time (field 22) is the
+	// 20th entry of the remainder.
+	fields := strings.Fields(string(raw[commEnd+1:]))
+	const startTimeOffset = 19
+	if len(fields) <= startTimeOffset {
+		return 0, fmt.Errorf("/proc/%d/stat has %d fields after comm", pid, len(fields))
+	}
+	return strconv.ParseUint(fields[startTimeOffset], 10, 64)
+}
+
+// processUID reads the real UID that owns a process.
+func processUID(pid int) (uint32, error) {
+	var st syscall.Stat_t
+	if err := syscall.Stat(fmt.Sprintf("/proc/%d", pid), &st); err != nil {
+		return 0, err
+	}
+	return st.Uid, nil
+}

client/vnc/server/shutdown_state_test.go

@@ -0,0 +1,79 @@
+//go:build unix
+
+package server
+
+import (
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// The start time is what tells a recorded process apart from whatever later
+// inherits its PID, so the field offset has to be right and the value stable.
+func TestProcessStartTimeIsStable(t *testing.T) {
+	if _, err := os.Stat("/proc/self/stat"); err != nil {
+		t.Skip("no procfs")
+	}
+
+	pid := os.Getpid()
+	first, err := processStartTime(pid)
+	require.NoError(t, err)
+	assert.NotZero(t, first, "a running process has a non-zero start time")
+
+	second, err := processStartTime(pid)
+	require.NoError(t, err)
+	assert.Equal(t, first, second, "start time must not move for a live process")
+}
+
+// A comm containing spaces and parentheses must not shift the field offsets,
+// which is why parsing starts from the last ')' rather than splitting the line.
+func TestProcessStartTimeToleratesOddCommName(t *testing.T) {
+	if _, err := os.Stat("/proc/self/stat"); err != nil {
+		t.Skip("no procfs")
+	}
+	sh, err := exec.LookPath("sh")
+	if err != nil {
+		t.Skip("no shell")
+	}
+
+	// argv[0] becomes the comm, truncated to 15 chars by the kernel.
+	cmd := exec.Command(sh, "-c", "sleep 30")
+	cmd.Args[0] = "a (b) c"
+	require.NoError(t, cmd.Start())
+	t.Cleanup(func() {
+		_ = cmd.Process.Kill()
+		_, _ = cmd.Process.Wait()
+	})
+
+	got, err := processStartTime(cmd.Process.Pid)
+	require.NoError(t, err)
+	assert.NotZero(t, got)
+}
+
+// A PID that never existed must not be reported as ours, and neither must a
+// record that carries no start time to compare against.
+func TestIsOurProcessRefusesUnidentifiableRecords(t *testing.T) {
+	assert.False(t, isOurProcess(sessionProcess{PID: -1}, "xvfb:50"))
+
+	pid := os.Getpid()
+	assert.False(t, isOurProcess(sessionProcess{PID: pid}, "xvfb:50"),
+		"a record with no recorded start time cannot be matched and must be refused")
+}
+
+// describeProcess captures enough to match the process back to itself.
+func TestDescribeProcessRoundTrips(t *testing.T) {
+	if _, err := os.Stat("/proc/self/stat"); err != nil {
+		t.Skip("no procfs")
+	}
+
+	proc := describeProcess(os.Getpid())
+	assert.Equal(t, os.Getpid(), proc.PID)
+	assert.NotZero(t, proc.StartTime)
+
+	start, err := processStartTime(proc.PID)
+	require.NoError(t, err)
+	assert.Equal(t, start, proc.StartTime)
+}

client/vnc/server/shutdown_state_windows.go

@@ -0,0 +1,19 @@
+//go:build windows
+
+package server
+
+// ShutdownState exists on Windows only so the shared server Config can name it.
+// Virtual sessions are an X11 feature: the Windows path proxies to an agent the
+// service control manager owns, so there are no residual processes of ours to
+// reap after a crash.
+type ShutdownState struct{}
+
+// Name returns the state name for the state manager.
+func (s *ShutdownState) Name() string {
+	return "vnc_sessions_state"
+}
+
+// Cleanup has nothing to do on Windows.
+func (s *ShutdownState) Cleanup() error {
+	return nil
+}

client/vnc/server/virtual_x11.go

@@ -251,6 +251,26 @@ func (vs *VirtualSession) Injector() InputInjector {
 	return vs.injector
 }
 
+// processes returns the identities of this session's live X server and desktop
+// processes, keyed by a description Cleanup uses to sanity-check them.
+func (vs *VirtualSession) processes() map[string]sessionProcess {
+	vs.mu.Lock()
+	defer vs.mu.Unlock()
+
+	if vs.stopped {
+		return nil
+	}
+	out := make(map[string]sessionProcess, 2)
+	display := strings.TrimPrefix(vs.display, ":")
+	if vs.xvfb != nil && vs.xvfb.Process != nil {
+		out["xvfb:"+display] = describeProcess(vs.xvfb.Process.Pid)
+	}
+	if vs.desktop != nil && vs.desktop.Process != nil {
+		out["desktop:"+display] = describeProcess(vs.desktop.Process.Pid)
+	}
+	return out
+}
+
 // Display returns the X11 display string (e.g., ":99").
 func (vs *VirtualSession) Display() string {
 	return vs.display
@@ -738,17 +758,37 @@ type sessionManager struct {
 	mu       sync.Mutex
 	sessions map[string]*VirtualSession
 	log      *log.Entry
+	// onProcesses publishes the live X server and desktop processes so the
+	// daemon can persist them for crash recovery. Nil when nothing is
+	// listening.
+	onProcesses func(*ShutdownState)
 }
 
-func newSessionManager(logger *log.Entry) *sessionManager {
+func newSessionManager(logger *log.Entry, onProcesses func(*ShutdownState)) *sessionManager {
 	sm := &sessionManager{
-		sessions: make(map[string]*VirtualSession),
-		log:      logger,
+		sessions:    make(map[string]*VirtualSession),
+		log:         logger,
+		onProcesses: onProcesses,
 	}
 	sm.sweepStaleXAuth()
 	return sm
 }
 
+// publishProcesses hands the current set of session processes to the daemon.
+// Called with sm.mu held, after any change to sm.sessions.
+func (sm *sessionManager) publishProcessesLocked() {
+	if sm.onProcesses == nil {
+		return
+	}
+	state := &ShutdownState{Processes: make(map[string]sessionProcess)}
+	for _, vs := range sm.sessions {
+		for desc, proc := range vs.processes() {
+			state.Processes[desc] = proc
+		}
+	}
+	sm.onProcesses(state)
+}
+
 // sweepStaleXAuth removes Xauthority files left over from a previous daemon
 // instance whose X servers are no longer running.
 func (sm *sessionManager) sweepStaleXAuth() {
@@ -790,6 +830,7 @@ func (sm *sessionManager) GetOrCreate(username string, width, height uint16) (vn
 		sm.log.Infof("replacing dead virtual session for %s", username)
 		vs.Stop()
 		delete(sm.sessions, username)
+		sm.publishProcessesLocked()
 	}
 
 	vs, err := StartVirtualSession(username, width, height, sm.log)
@@ -802,9 +843,11 @@ func (sm *sessionManager) GetOrCreate(username string, width, height uint16) (vn
 		if cur, ok := sm.sessions[username]; ok && cur == vs {
 			delete(sm.sessions, username)
 			sm.log.Infof("removed idle virtual session for %s", username)
+			sm.publishProcessesLocked()
 		}
 	}
 	sm.sessions[username] = vs
+	sm.publishProcessesLocked()
 	return vs, nil
 }
 
@@ -834,4 +877,5 @@ func (sm *sessionManager) StopAll() {
 		delete(sm.sessions, username)
 		sm.log.Infof("stopped virtual session for %s", username)
 	}
+	sm.publishProcessesLocked()
 }