Make the VNC server single-use and recognise crash-leftover desktop processes by their recorded command

client/vnc/server/server.go

@@ -102,6 +102,9 @@ type cursorPositionSource interface {
 	CursorPos() (x, y int, err error)
 }
 
+// errServerStopped is returned by Start on a Server that has been stopped.
+var errServerStopped = errors.New("VNC server was stopped and cannot be restarted; build a new one")
+
 // errFrameUnchanged is returned by capturers that hash the raw source
 // bytes (currently macOS) when the new frame is byte-identical to the
 // last one, so the encoder can short-circuit to an empty update.
@@ -186,7 +189,11 @@ type Server struct {
 	network6   netip.Prefix
 	log        *log.Entry
 
-	mu       sync.Mutex
+	mu sync.Mutex
+	// stopped is set by Stop. A Server is single-use: Stop closes the listener
+	// it may have been handed, the capturer and the injector, none of which
+	// it can recreate, so a second Start would run on closed resources.
+	stopped  bool
 	listener net.Listener
 	// extraListeners holds additional listeners (e.g. the v6 overlay), closed
 	// alongside listener on Stop.
@@ -777,11 +784,15 @@ func (s *Server) VNCAuth() *sshauth.Config {
 // Start begins listening for VNC connections on the given address.
 // network is the NetBird overlay prefix used to validate connection sources.
 // When Config.Listener was supplied, addr and network are ignored and the
-// pre-built listener is used (the per-session agent path).
+// pre-built listener is used (the per-session agent path). A Server cannot be
+// started again after Stop; build a new one instead.
 func (s *Server) Start(ctx context.Context, addr netip.AddrPort, network netip.Prefix) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	if s.stopped {
+		return errServerStopped
+	}
 	if s.listener != nil {
 		return fmt.Errorf("server already running")
 	}
@@ -789,15 +800,6 @@ func (s *Server) Start(ctx context.Context, addr netip.AddrPort, network netip.P
 		return fmt.Errorf("invalid agent token configuration")
 	}
 
-	// Reopen the door a previous Stop closed, so a restarted server accepts
-	// handlers again, and drop that Stop's drain signal so the next one waits
-	// on a fresh channel rather than on one already closed.
-	s.handlersMu.Lock()
-	s.stopping = false
-	s.handlersDrained = nil
-	s.handlersMu.Unlock()
-	s.resetServiceAgent()
-
 	s.ctx, s.cancel = context.WithCancel(ctx)
 	s.vmgr = s.platformSessionManager()
 
@@ -887,6 +889,7 @@ func (s *Server) openOverlayListener(addr netip.AddrPort, network netip.Prefix)
 func (s *Server) Stop() error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	s.stopped = true
 
 	// Before anything is closed, so no connection is admitted into a teardown
 	// already in progress.

client/vnc/server/server_test.go

@@ -681,44 +681,23 @@ func TestCloseAdmission_RefusesAndDrains(t *testing.T) {
 	}
 }
 
-// TestCloseAdmission_HandlerOutlivingDrain reproduces the restart hazard: Stop
-// gives up after handlerDrainTimeout, so a handler can still be running when a
-// later Start reopens admission. The counter must survive that (a sync.WaitGroup
-// panics on an Add racing the Wait it left behind), and the next Stop must wait
-// for both the straggler and the handlers admitted after the restart.
-func TestCloseAdmission_HandlerOutlivingDrain(t *testing.T) {
-	srv := &Server{log: log.WithField("test", t.Name())}
+// A Server is single-use. Stop closes the listener, capturer and injector it was
+// handed and cannot recreate, so a second Start would accept on a closed socket
+// and capture from a closed display. It must be refused instead.
+func TestStart_AfterStopIsRejected(t *testing.T) {
+	srv := New(Config{
+		Capturer:    &StubCapturer{},
+		Injector:    &StubInputInjector{},
+		DisableAuth: true,
+	})
+	addr := netip.MustParseAddrPort("127.0.0.1:0")
+	network := netip.MustParsePrefix("127.0.0.0/8")
 
-	require.True(t, srv.beginHandler())
-	firstDrain := srv.closeAdmission()
-	select {
-	case <-firstDrain:
-		t.Fatal("drained while a handler was still in flight")
-	default:
-	}
+	require.NoError(t, srv.Start(t.Context(), addr, network))
+	require.NoError(t, srv.Stop())
 
-	// What Start does after a drain that timed out.
-	srv.handlersMu.Lock()
-	srv.stopping = false
-	srv.handlersDrained = nil
-	srv.handlersMu.Unlock()
-
-	require.True(t, srv.beginHandler(), "a restarted server must admit handlers again")
-
-	secondDrain := srv.closeAdmission()
-	srv.endHandler() // the straggler from before the restart
-	select {
-	case <-secondDrain:
-		t.Fatal("drained before the post-restart handler finished")
-	default:
-	}
-
-	srv.endHandler()
-	select {
-	case <-secondDrain:
-	case <-time.After(time.Second):
-		t.Fatal("drain signal never fired after restart")
-	}
+	err := srv.Start(t.Context(), addr, network)
+	require.ErrorIs(t, err, errServerStopped, "a stopped server must not start again")
 }
 
 func TestAcceptRetryable(t *testing.T) {

client/vnc/server/service_agent.go

@@ -16,17 +16,6 @@ type sessionAgent interface {
 	Release()
 }
 
-// resetServiceAgent reopens the latch stopServiceAgent closed, so a restarted
-// server can build a manager again. Without it every accept loop of the new
-// lifecycle keeps getting a nil agent and service mode stays dead for the rest
-// of the process. Owned by Start, mirroring the other stop-time latches it
-// clears.
-func (s *Server) resetServiceAgent() {
-	s.serviceAgentMu.Lock()
-	defer s.serviceAgentMu.Unlock()
-	s.serviceAgentStopped = false
-}
-
 // stopServiceAgent tears down the shared manager, if one was ever built, and
 // latches the server so a still-draining accept loop cannot build another.
 // Owned by Stop rather than by an accept loop: the loops share the manager, so

client/vnc/server/shutdown_state.go

@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"fmt"
 	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"syscall"
@@ -27,6 +28,11 @@ type sessionProcess struct {
 	// StartTime is field 22 of /proc/<pid>/stat, in clock ticks since boot.
 	StartTime uint64 `json:"startTime,omitempty"`
 	UID       uint32 `json:"uid,omitempty"`
+	// Command is the base name of the process's argv[0] when it was started.
+	// Matching on it recognises whatever the launcher ran, including a desktop
+	// picked from xsessions or the xterm fallback, which a fixed list of
+	// names misses. Empty on records written before it was recorded.
+	Command string `json:"command,omitempty"`
 }
 
 // ShutdownState tracks VNC virtual session processes for crash recovery.
@@ -115,9 +121,20 @@ func describeProcess(pid int) sessionProcess {
 	} else {
 		log.Debugf("read uid for pid %d: %v", pid, err)
 	}
+	if cmdline, err := os.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid)); err == nil {
+		proc.Command = commandName(cmdline)
+	} else {
+		log.Debugf("read cmdline for pid %d: %v", pid, err)
+	}
 	return proc
 }
 
+// commandName returns the base name of argv[0] from a /proc cmdline.
+func commandName(cmdline []byte) string {
+	argv0, _, _ := bytes.Cut(cmdline, []byte{0})
+	return filepath.Base(string(argv0))
+}
+
 // 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
@@ -157,15 +174,25 @@ func isOurProcess(proc sessionProcess, desc string) bool {
 		return false
 	}
 
-	cmd := string(cmdline)
-	// Match against expected process types.
+	// The recorded command covers whatever the launcher ran; the name list
+	// covers records written before it was recorded, and a launcher script
+	// that has since exec'd into the real session binary under another name.
+	if proc.Command != "" && commandName(cmdline) == proc.Command {
+		return true
+	}
+	return matchesKnownSessionProcess(desc, string(cmdline))
+}
+
+// matchesKnownSessionProcess reports whether cmd looks like the X server or
+// desktop process desc describes.
+func matchesKnownSessionProcess(desc, cmd string) bool {
 	if strings.Contains(desc, "xvfb") || strings.Contains(desc, "xorg") {
 		return strings.Contains(cmd, "Xvfb") || strings.Contains(cmd, "Xorg")
 	}
 	if strings.Contains(desc, "desktop") {
 		return strings.Contains(cmd, "session") || strings.Contains(cmd, "plasma") ||
 			strings.Contains(cmd, "gnome") || strings.Contains(cmd, "xfce") ||
-			strings.Contains(cmd, "dbus-launch")
+			strings.Contains(cmd, "dbus-launch") || strings.Contains(cmd, "xterm")
 	}
 	return false
 }

client/vnc/server/shutdown_state_test.go

@@ -77,3 +77,28 @@ func TestDescribeProcessRoundTrips(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, start, proc.StartTime)
 }
+
+// A process recorded with its own command is recognised by that command, even
+// when it is not one of the known desktop names. The test binary stands in for
+// an xsessions entry or the xterm fallback: its desc says "desktop" but its name
+// matches nothing on the list.
+func TestIsOurProcessMatchesRecordedCommand(t *testing.T) {
+	if _, err := os.Stat("/proc/self/cmdline"); err != nil {
+		t.Skip("no procfs")
+	}
+
+	proc := describeProcess(os.Getpid())
+	require.NotEmpty(t, proc.Command, "the command name must be recorded")
+	assert.True(t, isOurProcess(proc, "desktop:50"),
+		"a record naming its command must match the live process by that command")
+
+	proc.Command = "something-else"
+	assert.False(t, isOurProcess(proc, "xvfb:50"),
+		"a recorded command that does not match, and no known name, must be refused")
+}
+
+func TestCommandName(t *testing.T) {
+	assert.Equal(t, "xterm", commandName([]byte("/usr/bin/xterm\x00-geometry\x0080x24\x00")))
+	assert.Equal(t, "Xvfb", commandName([]byte("Xvfb\x00:50\x00")))
+	assert.Equal(t, ".", commandName(nil))
+}