Keep DXGI on an idle desktop, honour the FreeBSD pitch when swizzling, close the handler-drain race

client/internal/approval/broker.go

@@ -82,6 +82,13 @@ var timeoutValue = func() time.Duration { return DefaultTimeout }
 // The caller must reject the underlying connection (fail-closed).
 var ErrNoSubscriber = errors.New("no UI subscriber connected for approval")
 
+// ErrPromptNotShown indicates a UI is connected but never received the prompt,
+// because its event queue was full when the prompt was published. Distinct from
+// ErrNoSubscriber so an operator can tell "nobody was listening" from "somebody
+// was listening and we could not reach them", which have different remedies.
+// Fail-closed either way.
+var ErrPromptNotShown = errors.New("approval prompt was not delivered to the UI")
+
 // ErrTimeout indicates the user did not respond within DefaultTimeout.
 var ErrTimeout = errors.New("approval timed out")
 
@@ -179,7 +186,7 @@ func (b *Broker) Request(ctx context.Context, p Prompt) (Decision, error) {
 	// than that nothing ever asked them.
 	if !b.pub.PublishEvent(proto.SystemEvent_INFO, proto.SystemEvent_APPROVAL, subject, subject, meta) {
 		log.Warnf("approval request %s (%s) reached no subscriber; denying without waiting", id, p.Kind)
-		return zero, ErrNoSubscriber
+		return zero, ErrPromptNotShown
 	}
 	log.Debugf("approval request %s (%s) emitted: %s", id, p.Kind, subject)
 

client/internal/approval/broker_test.go

@@ -75,7 +75,9 @@ func TestRequestUndeliveredPromptFailsFast(t *testing.T) {
 
 	start := time.Now()
 	_, err := b.Request(context.Background(), Prompt{Kind: KindVNC, Subject: "test"})
-	assert.ErrorIs(t, err, ErrNoSubscriber)
+	assert.ErrorIs(t, err, ErrPromptNotShown)
+	assert.NotErrorIs(t, err, ErrNoSubscriber,
+		"a UI was connected; the prompt just never reached it, which is a different problem")
 	assert.Less(t, time.Since(start), time.Second, "must not wait out the approval timeout")
 
 	b.mu.Lock()

client/vnc/server/agent_handshake.go

@@ -7,9 +7,11 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"crypto/subtle"
+	"errors"
 	"fmt"
 	"io"
 	"net"
+	"syscall"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -145,6 +147,28 @@ func agentServerHandshake(conn net.Conn, token []byte) (bool, error) {
 	return flag[0] != 0, nil
 }
 
+// isProbeDisconnect reports whether err is a peer that connected and left
+// without speaking.
+//
+// The daemon's own readiness check dials the agent socket and closes it
+// immediately, and it is not alone: anything probing the socket for liveness
+// does the same. Since the agent now writes its challenge first, such a probe
+// surfaces as a failed write (a reset or broken pipe) as often as a failed
+// read, and logging either at warning level would fill the daemon log with
+// entries for something entirely expected.
+func isProbeDisconnect(err error) bool {
+	switch {
+	case errors.Is(err, io.EOF), errors.Is(err, io.ErrUnexpectedEOF):
+		return true
+	case errors.Is(err, net.ErrClosed):
+		return true
+	case errors.Is(err, syscall.EPIPE), errors.Is(err, syscall.ECONNRESET):
+		return true
+	default:
+		return false
+	}
+}
+
 // viewOnlyByte renders the flag as the single byte both tags cover.
 func viewOnlyByte(viewOnly bool) []byte {
 	if viewOnly {

client/vnc/server/agent_ipc.go

@@ -153,17 +153,21 @@ func generateAuthToken() (string, error) {
 	return hex.EncodeToString(b), nil
 }
 
-// proxyToAgent dials the per-session agent's Unix socket, validates the
-// peer's kernel-asserted uid (so the daemon never hands its per-spawn
-// token to an impostor that won the listen race), writes the raw token
-// bytes plus a single view-only flag byte, then copies bytes both ways
-// until either side closes. The token + flag prefix must precede any RFB
-// byte so the agent's verifyAgentToken can run first. Returns nil once a
-// stream is established; the caller is responsible for sending an
-// RFB-level rejection on error so the client sees a reason instead of a
-// bare timeout. authedLog receives one audit line per dispatched
-// preamble so an operator can correlate daemon→agent traffic with the
-// remote session that triggered it.
+// proxyToAgent dials the per-session agent's Unix socket, checks the peer's
+// kernel-asserted uid, runs the mutual challenge-response that proves both ends
+// hold the per-spawn token, then copies bytes both ways until either side
+// closes.
+//
+// The handshake precedes any RFB byte, and the token stays on both ends: it is
+// only ever used as an HMAC key, so a process that squatted the socket learns
+// nothing it could replay, and the daemon stops before proxying to one that
+// cannot answer. See agent_handshake.go.
+//
+// Returns nil once a stream is established; the caller is responsible for
+// sending an RFB-level rejection on error so the client sees a reason instead
+// of a bare timeout. authedLog receives one audit line per established session
+// so an operator can correlate daemon→agent traffic with the remote session
+// that triggered it.
 func proxyToAgent(ctx context.Context, client net.Conn, socketPath, authToken string, peerUID uint32, viewOnly bool, authedLog *log.Entry) error {
 	tokenBytes, err := hex.DecodeString(authToken)
 	if err != nil || len(tokenBytes) != agentTokenLen {
@@ -185,7 +189,7 @@ func proxyToAgent(ctx context.Context, client net.Conn, socketPath, authToken st
 		return fmt.Errorf("agent handshake: %w", err)
 	}
 
-	// Audit: one line per successfully-dispatched daemon→agent preamble.
+	// Audit: one line per daemon→agent session that completed the handshake.
 	// Token printed as its first 8 hex chars (enough to correlate, not
 	// enough to use). Kept at Info so the default deployment captures it.
 	tokenFp := authToken
@@ -193,7 +197,7 @@ func proxyToAgent(ctx context.Context, client net.Conn, socketPath, authToken st
 		tokenFp = tokenFp[:8]
 	}
 	if authedLog != nil {
-		authedLog.Infof("VNC IPC: dispatched preamble to agent socket=%s peer_uid=%d view_only=%v token_fp=%s", socketPath, peerUID, viewOnly, tokenFp)
+		authedLog.Infof("VNC IPC: agent authenticated socket=%s peer_uid=%d view_only=%v token_fp=%s", socketPath, peerUID, viewOnly, tokenFp)
 	}
 
 	defer client.Close()

client/vnc/server/capture_dxgi_windows.go

@@ -6,7 +6,8 @@ import (
 	"errors"
 	"fmt"
 	"image"
-	"time"
+
+	log "github.com/sirupsen/logrus"
 
 	"github.com/kirides/go-d3d/d3d11"
 	"github.com/kirides/go-d3d/outputduplication"
@@ -28,6 +29,9 @@ type dxgiCapturer struct {
 	outIdx int
 	width  int
 	height int
+	// gotFrame records whether img has ever held a real desktop frame. Until it
+	// does, img is all zeroes and must not be handed out.
+	gotFrame bool
 }
 
 func newDXGICapturer() (*dxgiCapturer, error) {
@@ -62,49 +66,36 @@ func newDXGICapturer() (*dxgiCapturer, error) {
 		height: h,
 	}
 
-	if err := c.grabFirstFrame(); err != nil {
-		c.close()
-		return nil, err
+	// Best effort: an idle desktop legitimately has no new frame to give, and
+	// waiting here for one would only delay a capturer that is otherwise ready.
+	// capture() declines to hand out img until a frame has actually arrived.
+	if err := c.dup.GetImage(c.img, firstFrameTimeoutMS); err == nil {
+		c.gotFrame = true
+	} else if !errors.Is(err, outputduplication.ErrNoImageYet) {
+		log.Debugf("first DXGI frame: %v", err)
 	}
 
 	return c, nil
 }
 
-// Bounds the wait for the very first frame. Each attempt blocks for
-// firstFrameAttempt, so the deadline allows a few of them.
-const (
-	firstFrameAttemptMS = 2000
-	firstFrameDeadline  = 6 * time.Second
-)
-
-// grabFirstFrame fills c.img with a real desktop frame before the capturer is
-// handed out.
-//
-// capture() deliberately tolerates ErrNoImageYet, because on an idle desktop
-// DXGI reports "nothing new" and the right answer is the frame already in hand.
-// At construction there is no such frame: accepting the timeout there would
-// publish the all-zero buffer, and every session attaching in that window would
-// be served a black screen that looks like a successful capture. Failing
-// instead lets createCapturer fall back to GDI.
-func (c *dxgiCapturer) grabFirstFrame() error {
-	deadline := time.Now().Add(firstFrameDeadline)
-	for {
-		err := c.dup.GetImage(c.img, firstFrameAttemptMS)
-		if err == nil {
-			return nil
-		}
-		if !errors.Is(err, outputduplication.ErrNoImageYet) {
-			return fmt.Errorf("acquire first desktop frame: %w", err)
-		}
-		if time.Now().After(deadline) {
-			return fmt.Errorf("no desktop frame within %s", firstFrameDeadline)
-		}
-	}
-}
+// firstFrameTimeoutMS is how long the constructor waits for an initial frame
+// before handing back a capturer that has not produced one yet.
+const firstFrameTimeoutMS = 2000
 
 func (c *dxgiCapturer) capture() (*image.RGBA, error) {
 	err := c.dup.GetImage(c.img, 100)
-	if err != nil && !errors.Is(err, outputduplication.ErrNoImageYet) {
+	switch {
+	case err == nil:
+		c.gotFrame = true
+	case !errors.Is(err, outputduplication.ErrNoImageYet):
+		return nil, err
+	case !c.gotFrame:
+		// "No new frame" is the right answer on an idle desktop, but only once
+		// there is a frame to repeat. Before that img is all zeroes, and
+		// returning it would serve a black desktop that looks like a successful
+		// capture. Reported as an error instead: the worker retries shortly and
+		// stays on DXGI, where failing the constructor would have dropped the
+		// whole session to the much slower GDI path.
 		return nil, err
 	}
 

client/vnc/server/capture_fb_freebsd.go

@@ -147,7 +147,12 @@ func (c *FBCapturer) CaptureInto(dst *image.RGBA) error {
 	switch c.bpp {
 	case 32:
 		// vt(4) on KMS framebuffers is BGRA: byte 0=B, 1=G, 2=R.
-		swizzleBGRAtoRGBA(dst.Pix, c.mmap[:c.h*c.stride])
+		//
+		// Row-aware, like the 24- and 16-bit paths: a padded pitch means the
+		// bytes after each row's pixels are not pixels, and a flat swizzle over
+		// h*stride would feed them to the encoder and slide every subsequent
+		// row left by the padding.
+		swizzleFB32BGRA(dst.Pix, dst.Stride, c.mmap, c.stride, c.w, c.h)
 	case 24:
 		swizzleFB24(dst.Pix, dst.Stride, c.mmap, c.stride, c.w, c.h)
 	case 16:
@@ -156,6 +161,21 @@ func (c *FBCapturer) CaptureInto(dst *image.RGBA) error {
 	return nil
 }
 
+// swizzleFB32BGRA converts 32bpp BGRA rows into RGBA, honouring the source
+// pitch so padding between rows is skipped rather than read as pixels.
+func swizzleFB32BGRA(dst []byte, dstStride int, src []byte, srcStride, w, h int) {
+	for y := 0; y < h; y++ {
+		srcRow := src[y*srcStride : y*srcStride+w*4]
+		dstRow := dst[y*dstStride:]
+		for x := 0; x < w; x++ {
+			dstRow[x*4+0] = srcRow[x*4+2]
+			dstRow[x*4+1] = srcRow[x*4+1]
+			dstRow[x*4+2] = srcRow[x*4+0]
+			dstRow[x*4+3] = 0xff
+		}
+	}
+}
+
 // Close releases the framebuffer mmap and file descriptor. Serialized with
 // CaptureInto via c.mu so an in-flight capture can't read freed memory.
 func (c *FBCapturer) Close() {

client/vnc/server/handshake.go

@@ -223,19 +223,16 @@ func (s *Server) maybeRunNoiseHandshake(conn net.Conn, magic [4]byte, headerMode
 	return peerStatic, true, true, nil
 }
 
-// verifyAgentToken validates the agent token prefix when configured and
-// reads the trailing view-only flag byte the daemon writes alongside it.
-// Returns (ok, viewOnly). ok=false closes the connection.
+// verifyAgentToken runs the agent's half of the mutual challenge-response with
+// the daemon when a token is configured, and reports the view-only flag the
+// daemon authenticated. Returns (ok, viewOnly). ok=false closes the connection.
 func (s *Server) verifyAgentToken(conn net.Conn, connLog *log.Entry) (bool, bool) {
 	if len(s.agentToken) == 0 {
 		return true, false
 	}
 	viewOnly, err := agentServerHandshake(conn, s.agentToken)
 	if err != nil {
-		if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
-			// Connect-then-close probes (the daemon's own readiness check
-			// among them) hit this path on every dial; logging them would
-			// just flood the daemon log without surfacing a real failure.
+		if isProbeDisconnect(err) {
 			connLog.Tracef("agent auth: %v", err)
 		} else {
 			connLog.Warnf("agent auth: %v", err)

client/vnc/server/server.go

@@ -185,7 +185,16 @@ type Server struct {
 	vmgr           virtualSessionManager
 	// handlers counts the in-flight connection handlers so Stop can wait for
 	// them before tearing down the capturer and injector they use.
-	handlers sync.WaitGroup
+	//
+	// handlersMu guards both the counter's growth and stopping. A WaitGroup
+	// forbids an Add that starts from zero while a Wait is in flight, and the
+	// accept loops run in their own goroutines, so without a barrier a
+	// connection accepted just before the listener closed could Add after Wait
+	// had already returned. It cannot be s.mu: Stop holds that for its whole
+	// body, so an accept loop taking it would deadlock against the wait.
+	handlersMu sync.Mutex
+	handlers   sync.WaitGroup
+	stopping   bool
 	// onVirtualProcesses forwards live virtual-session process records to the
 	// daemon for crash recovery; nil when nothing is listening.
 	onVirtualProcesses func(*ShutdownState)
@@ -469,9 +478,29 @@ func (s *Server) closeActiveSessions() {
 // wedged in a syscall; shutdown must not hang on it.
 const handlerDrainTimeout = 5 * time.Second
 
+// beginHandler registers one in-flight connection handler, reporting false when
+// the server is already stopping and the caller should drop the connection
+// instead of starting one.
+func (s *Server) beginHandler() bool {
+	s.handlersMu.Lock()
+	defer s.handlersMu.Unlock()
+
+	if s.stopping {
+		return false
+	}
+	s.handlers.Add(1)
+	return true
+}
+
 // awaitHandlers waits for the in-flight connection handlers, giving up after
 // handlerDrainTimeout so a stuck one cannot hold shutdown open.
 func (s *Server) awaitHandlers() {
+	// Closes the door before waiting: past this point beginHandler refuses, so
+	// no Add can begin while Wait is running.
+	s.handlersMu.Lock()
+	s.stopping = true
+	s.handlersMu.Unlock()
+
 	done := make(chan struct{})
 	go func() {
 		s.handlers.Wait()
@@ -659,6 +688,12 @@ 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.
+	s.handlersMu.Lock()
+	s.stopping = false
+	s.handlersMu.Unlock()
+
 	s.ctx, s.cancel = context.WithCancel(ctx)
 	s.vmgr = s.platformSessionManager()
 
@@ -835,7 +870,12 @@ func (s *Server) acceptLoop(ln net.Listener) {
 			continue
 		}
 		enableTCPKeepAlive(conn, s.log)
-		s.handlers.Add(1)
+		if !s.beginHandler() {
+			s.releaseConnSlot()
+			s.untrackConn(conn)
+			_ = conn.Close()
+			continue
+		}
 		go func(c net.Conn) {
 			defer s.handlers.Done()
 			defer s.releaseConnSlot()

client/vnc/server/server_darwin.go

@@ -62,7 +62,12 @@ func (s *Server) serviceAcceptLoop(ln net.Listener) {
 		enableTCPKeepAlive(conn, s.log)
 		metered := newMetricsConn(conn, s.sessionRecorder)
 		s.retrackConn(conn, metered)
-		s.handlers.Add(1)
+		if !s.beginHandler() {
+			s.releaseConnSlot()
+			s.untrackConn(metered)
+			_ = conn.Close()
+			continue
+		}
 		go func(c net.Conn) {
 			defer s.handlers.Done()
 			defer s.releaseConnSlot()

client/vnc/server/server_test.go

@@ -297,6 +297,10 @@ func TestAgentToken_MismatchClosesConnection(t *testing.T) {
 		// never reached the greeting.
 		_ = err
 	}
+	// The handshake clears the deadline on its way out, so re-arm it: without
+	// one, a server that neither closes nor greets would hang this test until
+	// the CI timeout instead of failing here.
+	require.NoError(t, conn.SetDeadline(time.Now().Add(10*time.Second)))
 
 	// Server must close without sending the RFB greeting.
 	var version [12]byte
@@ -329,6 +333,8 @@ func TestAgentToken_MatchAllowsHandshake(t *testing.T) {
 	require.NoError(t, conn.SetDeadline(time.Now().Add(10*time.Second)))
 
 	require.NoError(t, agentClientHandshake(conn, token, false))
+	// Re-armed because the handshake clears it on the way out.
+	require.NoError(t, conn.SetDeadline(time.Now().Add(10*time.Second)))
 
 	// Send session header so handleConnection can proceed past readConnectionHeader.
 	header := make([]byte, 11) // ModeAttach + usernameLen=0 + sessionID=0 + width=0 + height=0

client/vnc/server/server_windows.go

@@ -349,7 +349,12 @@ func (s *Server) serviceAcceptLoop(ln net.Listener) {
 		enableTCPKeepAlive(conn, s.log)
 		metered := newMetricsConn(conn, s.sessionRecorder)
 		s.retrackConn(conn, metered)
-		s.handlers.Add(1)
+		if !s.beginHandler() {
+			s.releaseConnSlot()
+			s.untrackConn(metered)
+			_ = conn.Close()
+			continue
+		}
 		go func(c net.Conn) {
 			defer s.handlers.Done()
 			defer s.releaseConnSlot()