Detect dead VNC peers on both ends and report session stats

client/vnc/server/metrics_conn.go

@@ -0,0 +1,187 @@
+//go:build !js && !ios && !android
+
+package server
+
+import (
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// SessionTick is one sampling slice of a VNC session's wire activity.
+// BytesOut / Writes / FBUs are deltas observed during this tick;
+// Max* fields are the high-water marks observed during this tick (reset
+// at the start of the next). Period is the wall-clock duration covered
+// (typically sessionTickInterval, shorter for the final flush).
+type SessionTick struct {
+	Period        time.Duration
+	BytesOut      uint64
+	Writes        uint64
+	FBUs          uint64
+	MaxFBUBytes   uint64
+	MaxFBURects   uint64
+	MaxWriteBytes uint64
+	WriteNanos    uint64
+}
+
+// sessionTickInterval is how often metricsConn emits a SessionTick. One
+// second matches noVNC's request cadence so each tick covers roughly one
+// FBU round-trip during steady-state activity.
+const sessionTickInterval = time.Second
+
+// metricsConn wraps a net.Conn and tracks per-session byte / write / FBU
+// counters. Updates are atomic so the cost is a few atomic ops per Write
+// (well under 100 ns), negligible against the syscall itself, so the wrap
+// is always installed. A goroutine emits a SessionTick to the recorder
+// every sessionTickInterval (only when the tick has activity to report);
+// a final partial-tick flush runs on Close.
+type metricsConn struct {
+	net.Conn
+
+	recorder func(SessionTick)
+
+	bytesOut    uint64
+	writes      uint64
+	writeNanos  uint64
+	largestPkt  uint64
+	fbus        uint64
+	fbuBytes    uint64
+	fbuRects    uint64
+	maxFBUBytes uint64
+	maxFBURects uint64
+
+	tickMu       sync.Mutex
+	tickStart    time.Time
+	tickPrevB    uint64
+	tickPrevW    uint64
+	tickPrevF    uint64
+	tickPrevNS   uint64
+
+	closeOnce sync.Once
+	done      chan struct{}
+}
+
+func newMetricsConn(c net.Conn, recorder func(SessionTick)) net.Conn {
+	m := &metricsConn{
+		Conn:      c,
+		recorder:  recorder,
+		tickStart: time.Now(),
+		done:      make(chan struct{}),
+	}
+	if recorder != nil {
+		go m.tickLoop()
+	}
+	return m
+}
+
+// tickLoop emits a SessionTick every sessionTickInterval until done.
+// Empty ticks (no writes since the last tick) are skipped.
+func (m *metricsConn) tickLoop() {
+	t := time.NewTicker(sessionTickInterval)
+	defer t.Stop()
+	for {
+		select {
+		case <-m.done:
+			return
+		case <-t.C:
+			m.flushTick(false)
+		}
+	}
+}
+
+// flushTick computes deltas since the last tick, resets the per-tick max
+// trackers, and emits a SessionTick to the recorder. final=true forces
+// emission even if no writes happened (used at session close to record
+// the trailing partial period).
+func (m *metricsConn) flushTick(final bool) {
+	m.tickMu.Lock()
+	defer m.tickMu.Unlock()
+
+	b := atomic.LoadUint64(&m.bytesOut)
+	w := atomic.LoadUint64(&m.writes)
+	f := atomic.LoadUint64(&m.fbus)
+	ns := atomic.LoadUint64(&m.writeNanos)
+
+	db := b - m.tickPrevB
+	dw := w - m.tickPrevW
+	df := f - m.tickPrevF
+	dns := ns - m.tickPrevNS
+	m.tickPrevB, m.tickPrevW, m.tickPrevF, m.tickPrevNS = b, w, f, ns
+
+	maxFBU := atomic.SwapUint64(&m.maxFBUBytes, 0)
+	maxRects := atomic.SwapUint64(&m.maxFBURects, 0)
+	maxPkt := atomic.SwapUint64(&m.largestPkt, 0)
+
+	period := time.Since(m.tickStart)
+	m.tickStart = time.Now()
+
+	if dw == 0 && !final {
+		return
+	}
+	m.recorder(SessionTick{
+		Period:        period,
+		BytesOut:      db,
+		Writes:        dw,
+		FBUs:          df,
+		MaxFBUBytes:   maxFBU,
+		MaxFBURects:   maxRects,
+		MaxWriteBytes: maxPkt,
+		WriteNanos:    dns,
+	})
+}
+
+// isFBUHeader reports whether the given Write payload is the 4-byte
+// FramebufferUpdate header (message type 0, padding 0, rect-count high
+// byte). Rect bodies are written separately by sendDirtyAndMoves, so the
+// FBU/rect boundary lines up with Write boundaries.
+func isFBUHeader(p []byte) bool {
+	return len(p) == 4 && p[0] == serverFramebufferUpdate
+}
+
+func (m *metricsConn) Write(p []byte) (int, error) {
+	if isFBUHeader(p) {
+		if b := atomic.SwapUint64(&m.fbuBytes, 0); b > 0 {
+			if b > atomic.LoadUint64(&m.maxFBUBytes) {
+				atomic.StoreUint64(&m.maxFBUBytes, b)
+			}
+		}
+		if r := atomic.SwapUint64(&m.fbuRects, 0); r > 0 {
+			if r > atomic.LoadUint64(&m.maxFBURects) {
+				atomic.StoreUint64(&m.maxFBURects, r)
+			}
+		}
+		atomic.AddUint64(&m.fbus, 1)
+	}
+
+	t0 := time.Now()
+	n, err := m.Conn.Write(p)
+	atomic.AddUint64(&m.writeNanos, uint64(time.Since(t0).Nanoseconds()))
+	atomic.AddUint64(&m.bytesOut, uint64(n))
+	atomic.AddUint64(&m.writes, 1)
+	if !isFBUHeader(p) {
+		atomic.AddUint64(&m.fbuBytes, uint64(n))
+		atomic.AddUint64(&m.fbuRects, 1)
+	}
+	if uint64(n) > atomic.LoadUint64(&m.largestPkt) {
+		atomic.StoreUint64(&m.largestPkt, uint64(n))
+	}
+	return n, err
+}
+
+func (m *metricsConn) Close() error {
+	m.closeOnce.Do(func() {
+		close(m.done)
+		if m.recorder == nil {
+			return
+		}
+		if b := atomic.SwapUint64(&m.fbuBytes, 0); b > atomic.LoadUint64(&m.maxFBUBytes) {
+			atomic.StoreUint64(&m.maxFBUBytes, b)
+		}
+		if r := atomic.SwapUint64(&m.fbuRects, 0); r > atomic.LoadUint64(&m.maxFBURects) {
+			atomic.StoreUint64(&m.maxFBURects, r)
+		}
+		m.flushTick(true)
+	})
+	return m.Conn.Close()
+}

client/vnc/server/server.go

@@ -156,9 +156,15 @@ type Server struct {
 	netstackNet  *netstack.Net
 	agentToken   []byte // raw token bytes for agent-mode auth
 
-	sessionsMu sync.Mutex
-	sessionSeq uint64
-	sessions   map[uint64]ActiveSessionInfo
+	sessionsMu   sync.Mutex
+	sessionSeq   uint64
+	sessions     map[uint64]ActiveSessionInfo
+	sessionConns map[uint64]net.Conn
+
+	// sessionRecorder, when non-nil, receives a SessionTick periodically
+	// during each VNC session and on session close. The engine wires
+	// this to its metrics framework.
+	sessionRecorder func(SessionTick)
 }
 
 // ActiveSessionInfo describes a currently connected VNC client.
@@ -195,7 +201,8 @@ func New(capturer ScreenCapturer, injector InputInjector) *Server {
 		injector:   injector,
 		authorizer: sshauth.NewAuthorizer(),
 		log:        log.WithField("component", "vnc-server"),
-		sessions:   make(map[uint64]ActiveSessionInfo),
+		sessions:     make(map[uint64]ActiveSessionInfo),
+		sessionConns: make(map[uint64]net.Conn),
 	}
 }
 
@@ -210,12 +217,13 @@ func (s *Server) ActiveSessions() []ActiveSessionInfo {
 	return out
 }
 
-func (s *Server) addSession(info ActiveSessionInfo) uint64 {
+func (s *Server) addSession(info ActiveSessionInfo, conn net.Conn) uint64 {
 	s.sessionsMu.Lock()
 	defer s.sessionsMu.Unlock()
 	s.sessionSeq++
 	id := s.sessionSeq
 	s.sessions[id] = info
+	s.sessionConns[id] = conn
 	return id
 }
 
@@ -223,6 +231,24 @@ func (s *Server) removeSession(id uint64) {
 	s.sessionsMu.Lock()
 	defer s.sessionsMu.Unlock()
 	delete(s.sessions, id)
+	delete(s.sessionConns, id)
+}
+
+// closeActiveSessions closes every active session's connection so the
+// per-session serve goroutines unblock from their Read loops and exit.
+// Called from Stop to make sure clients see an immediate disconnect when
+// the server is brought down, instead of waiting for the OS to reclaim
+// the sockets after process exit.
+func (s *Server) closeActiveSessions() {
+	s.sessionsMu.Lock()
+	conns := make([]net.Conn, 0, len(s.sessionConns))
+	for _, c := range s.sessionConns {
+		conns = append(conns, c)
+	}
+	s.sessionsMu.Unlock()
+	for _, c := range conns {
+		_ = c.Close()
+	}
 }
 
 // SetServiceMode enables proxy-to-agent mode for Windows service operation.
@@ -230,6 +256,14 @@ func (s *Server) SetServiceMode(enabled bool) {
 	s.serviceMode = enabled
 }
 
+// SetSessionRecorder installs a callback that receives a SessionTick
+// each sessionTickInterval during a VNC session and one final tick on
+// session close. Pass nil to disable. Empty ticks (no wire activity)
+// are skipped.
+func (s *Server) SetSessionRecorder(recorder func(SessionTick)) {
+	s.sessionRecorder = recorder
+}
+
 // SetJWTConfig configures JWT authentication for VNC connections.
 // Pass nil to disable JWT (public mode).
 func (s *Server) SetJWTConfig(config *JWTConfig) {
@@ -340,6 +374,13 @@ 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.
+	s.closeActiveSessions()
+
 	if s.vmgr != nil {
 		s.vmgr.StopAll()
 	}
@@ -378,10 +419,36 @@ func (s *Server) acceptLoop() {
 			continue
 		}
 
+		enableTCPKeepAlive(conn, s.log)
 		go s.handleConnection(conn)
 	}
 }
 
+// vncKeepAlivePeriod controls how often TCP layer probes are sent on an
+// idle connection. Default OS settings (2 hours) are too long for an
+// interactive session: when the server-side host dies without sending FIN
+// (power loss, network partition, hung kernel), the client only learns of
+// the dead connection when the OS gives up on a probe. 30 s here means
+// most clients notice within ~3 minutes worst case.
+const vncKeepAlivePeriod = 30 * time.Second
+
+// enableTCPKeepAlive turns on SO_KEEPALIVE on the underlying TCP socket.
+// Non-TCP conns (e.g. the netstack-backed listener) are skipped silently;
+// keepalive there is the netstack's concern.
+func enableTCPKeepAlive(c net.Conn, log *log.Entry) {
+	tc, ok := c.(*net.TCPConn)
+	if !ok {
+		return
+	}
+	if err := tc.SetKeepAlive(true); err != nil {
+		log.Debugf("set keepalive: %v", err)
+		return
+	}
+	if err := tc.SetKeepAlivePeriod(vncKeepAlivePeriod); err != nil {
+		log.Debugf("set keepalive period: %v", err)
+	}
+}
+
 func (s *Server) validateCapturer(capturer ScreenCapturer) error {
 	// Quick check first: if already ready, return immediately.
 	if capturer.Width() > 0 && capturer.Height() > 0 {
@@ -472,7 +539,7 @@ func (s *Server) handleConnection(conn net.Conn) {
 		Mode:          modeString(header.mode),
 		Username:      header.username,
 		JWTUsername:   jwtUserID,
-	})
+	}, conn)
 	defer s.removeSession(sessionID)
 
 	if err := s.validateCapturer(capturer); err != nil {
@@ -481,6 +548,7 @@ func (s *Server) handleConnection(conn net.Conn) {
 		return
 	}
 
+	conn = newMetricsConn(conn, s.sessionRecorder)
 	sess := &session{
 		conn:     conn,
 		capturer: capturer,

client/vnc/server/server_windows.go

@@ -255,6 +255,8 @@ func (s *Server) serviceAcceptLoop() {
 			continue
 		}
 
+		enableTCPKeepAlive(conn, s.log)
+		conn = newMetricsConn(conn, s.sessionRecorder)
 		go s.handleServiceConnection(conn, sm)
 	}
 }