Count in-flight writes in the final metrics tick and scope the agent token constants to their platforms

client/vnc/server/agent_darwin.go

@@ -22,6 +22,15 @@ import (
 	"github.com/netbirdio/netbird/client/configs"
 )
 
+// agentTokenStdinFlag tells the agent to read its token from stdin instead of
+// the environment. Must match the flag cmd.vncAgentCmd registers. The agent
+// runs as the console user, and the environment it was started with stays
+// readable by that user's other processes for the agent's whole life (macOS
+// KERN_PROCARGS2 keeps the original strings even after unsetenv): holding the
+// token lets a process drive the agent directly, past the daemon's gates and
+// with the agent's Screen Recording grant.
+const agentTokenStdinFlag = "--token-stdin" // #nosec G101 -- flag name, not a credential
+
 // darwinAgentManager spawns a per-user VNC agent on demand and keeps it alive
 // only while connections are using it. Concurrent connections share one agent;
 // the last one to finish takes it down again.

client/vnc/server/agent_ipc.go

@@ -128,28 +128,10 @@ func (s *Server) handleServiceConnection(conn net.Conn, sa sessionAgent) {
 	authedLog.Infof("VNC connection closed (%dms)", time.Since(start).Milliseconds())
 }
 
-const (
-	// agentTokenEnvVar names the environment variable the daemon uses to
-	// hand the per-spawn token to the agent child. Out-of-band channels
-	// like this keep the secret out of the command line, where listings
-	// such as `ps` or Windows tasklist would expose it.
-	agentTokenEnvVar = "NB_VNC_AGENT_TOKEN" // #nosec G101 -- env var name, not a credential
-
-	// agentTokenStdinFlag tells the agent to read its token from stdin
-	// instead of agentTokenEnvVar. Must match the flag cmd.vncAgentCmd
-	// registers. Used where the agent runs as the console user: there the
-	// environment it was started with stays readable by that user's other
-	// processes for the agent's whole life (macOS KERN_PROCARGS2 keeps the
-	// original strings even after unsetenv), and holding the token lets a
-	// process drive the agent directly, past the daemon's gates and with the
-	// agent's Screen Recording grant.
-	agentTokenStdinFlag = "--token-stdin" // #nosec G101 -- flag name, not a credential
-
-	// vncAgentSubcommand is the CLI subcommand the daemon invokes to start
-	// the per-session agent process. Must match cmd.vncAgentCmd.Use in
-	// client/cmd/vnc_agent.go.
-	vncAgentSubcommand = "vnc-agent"
-)
+// vncAgentSubcommand is the CLI subcommand the daemon invokes to start the
+// per-session agent process. Must match cmd.vncAgentCmd.Use in
+// client/cmd/vnc_agent.go.
+const vncAgentSubcommand = "vnc-agent"
 
 // generateAuthToken returns a fresh hex-encoded random token for one
 // daemon→agent session. The daemon hands this to the spawned agent

client/vnc/server/agent_windows.go

@@ -21,6 +21,12 @@ import (
 	"golang.org/x/sys/windows"
 )
 
+// agentTokenEnvVar names the environment variable the daemon uses to hand the
+// per-spawn token to the agent child. Out-of-band channels like this keep the
+// secret out of the command line, where tasklist would expose it. The Windows
+// agent runs as SYSTEM, so its environment is not readable by user processes.
+const agentTokenEnvVar = "NB_VNC_AGENT_TOKEN" // #nosec G101 -- env var name, not a credential
+
 const (
 	stillActive = 259
 

client/vnc/server/metrics_conn.go

@@ -75,6 +75,10 @@ type metricsConn struct {
 	busyLastNanos uint64
 	busyFraction  float64
 
+	// writeMu is held shared by each Write for its whole duration and taken
+	// exclusively by Close before the final snapshot, so that snapshot counts
+	// every write that was in flight when the connection was closed.
+	writeMu   sync.RWMutex
 	closeOnce sync.Once
 	done      chan struct{}
 }
@@ -224,6 +228,9 @@ func (m *metricsConn) endFBU() {
 }
 
 func (m *metricsConn) Write(p []byte) (int, error) {
+	m.writeMu.RLock()
+	defer m.writeMu.RUnlock()
+
 	t0 := time.Now()
 	n, err := m.Conn.Write(p)
 	m.writeNanos.Add(uint64(time.Since(t0).Nanoseconds()))
@@ -252,14 +259,29 @@ func (m *metricsConn) flushFBUMax() {
 	}
 }
 
+// Close closes the connection and records the final partial tick. The socket is
+// closed first, which is what unblocks a Write stuck on a peer that stopped
+// reading; the snapshot is taken only once those writes have returned, so the
+// last frame and the last write of a session are counted in it rather than
+// lost.
 func (m *metricsConn) Close() error {
+	closed := false
+	var err error
 	m.closeOnce.Do(func() {
+		closed = true
 		close(m.done)
+		err = m.Conn.Close()
+
+		m.writeMu.Lock()
+		defer m.writeMu.Unlock()
 		if m.recorder == nil {
 			return
 		}
 		m.flushFBUMax()
 		m.flushTick(true)
 	})
-	return m.Conn.Close()
+	if !closed {
+		return m.Conn.Close()
+	}
+	return err
 }

client/vnc/server/metrics_conn_test.go

@@ -0,0 +1,76 @@
+//go:build !js && !ios && !android
+
+package server
+
+import (
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// A session's last write is often one stuck on a peer that stopped reading,
+// and Close is what unblocks it. The final tick has to be taken after that
+// write returns, or the session's closing metrics leave it out.
+func TestMetricsConn_FinalTickCountsInFlightWrite(t *testing.T) {
+	local, remote := net.Pipe()
+	t.Cleanup(func() { _ = remote.Close() })
+
+	var mu sync.Mutex
+	var ticks []SessionTick
+	conn := newMetricsConn(local, func(tick SessionTick) {
+		mu.Lock()
+		ticks = append(ticks, tick)
+		mu.Unlock()
+	})
+
+	// Nothing reads remote, so this Write blocks until the pipe is closed.
+	writing := make(chan struct{})
+	wrote := make(chan struct{})
+	go func() {
+		close(writing)
+		_, _ = conn.Write([]byte("frame"))
+		close(wrote)
+	}()
+	<-writing
+	time.Sleep(20 * time.Millisecond)
+
+	require.NoError(t, conn.Close())
+	select {
+	case <-wrote:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Close did not unblock the in-flight write")
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+	require.NotEmpty(t, ticks, "Close must emit a final tick")
+	final := ticks[len(ticks)-1]
+	assert.Equal(t, uint64(1), final.Writes, "the write in flight at Close must be counted in the final tick")
+}
+
+// A proxied service-mode connection never sees FBU boundaries, so its ticks
+// must say so rather than report zero updates.
+func TestMetricsConn_ProxyReportsFBUsUntracked(t *testing.T) {
+	local, remote := net.Pipe()
+	t.Cleanup(func() { _ = remote.Close() })
+	go func() { buf := make([]byte, 64); _, _ = remote.Read(buf) }()
+
+	var got []SessionTick
+	var mu sync.Mutex
+	conn := newProxyMetricsConn(local, func(tick SessionTick) {
+		mu.Lock()
+		got = append(got, tick)
+		mu.Unlock()
+	})
+	_, _ = conn.Write([]byte("x"))
+	require.NoError(t, conn.Close())
+
+	mu.Lock()
+	defer mu.Unlock()
+	require.NotEmpty(t, got)
+	assert.False(t, got[len(got)-1].FBUsTracked, "a proxied connection's FBU fields are unknown, not zero")
+}