From d6f46b79a3591fd26e85f649bd19280888549261 Mon Sep 17 00:00:00 2001 From: Viktor Liu Date: Wed, 23 Sep 2026 08:54:08 +0200 Subject: [PATCH] Count in-flight writes in the final metrics tick and scope the agent token constants to their platforms --- client/vnc/server/agent_darwin.go | 9 +++ client/vnc/server/agent_ipc.go | 26 ++------- client/vnc/server/agent_windows.go | 6 ++ client/vnc/server/metrics_conn.go | 24 +++++++- client/vnc/server/metrics_conn_test.go | 76 ++++++++++++++++++++++++++ 5 files changed, 118 insertions(+), 23 deletions(-) create mode 100644 client/vnc/server/metrics_conn_test.go diff --git a/client/vnc/server/agent_darwin.go b/client/vnc/server/agent_darwin.go index ebb87893b..e7c05c347 100644 --- a/client/vnc/server/agent_darwin.go +++ b/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. diff --git a/client/vnc/server/agent_ipc.go b/client/vnc/server/agent_ipc.go index 2c6c18aeb..f642310b3 100644 --- a/client/vnc/server/agent_ipc.go +++ b/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 diff --git a/client/vnc/server/agent_windows.go b/client/vnc/server/agent_windows.go index bd5e8e35e..45bfd9e59 100644 --- a/client/vnc/server/agent_windows.go +++ b/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 diff --git a/client/vnc/server/metrics_conn.go b/client/vnc/server/metrics_conn.go index b4b747a0e..2821713bf 100644 --- a/client/vnc/server/metrics_conn.go +++ b/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 } diff --git a/client/vnc/server/metrics_conn_test.go b/client/vnc/server/metrics_conn_test.go new file mode 100644 index 000000000..21f0c25af --- /dev/null +++ b/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") +}