From 9ede661bab372826d45b3f4c47b42458eb2fb8fe Mon Sep 17 00:00:00 2001
From: Viktor Liu <viktor@netbird.io>
Date: Wed, 23 Sep 2026 15:05:30 +0200
Subject: [PATCH] Retrack the agent connection once it is wrapped with the
 granted peer address

---
 client/vnc/server/agent_ipc.go   | 11 -----------
 client/vnc/server/server.go      | 18 ++++++++++++++++--
 client/vnc/server/server_test.go | 10 ++++++++++
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/client/vnc/server/agent_ipc.go b/client/vnc/server/agent_ipc.go
index 698e379e5..1b6642ca3 100644
--- a/client/vnc/server/agent_ipc.go
+++ b/client/vnc/server/agent_ipc.go
@@ -293,14 +293,3 @@ func dialAgentWithRetry(ctx context.Context, addr string) (net.Conn, error) {
 	}
 	return nil, lastErr
 }
-
-// retrackConn replaces a tracked raw connection with the wrapper its handler
-// will actually hold, so shutdown and the handler's own untrackConn agree on
-// which object is registered. Service mode only: that is where an accepted
-// connection is wrapped before the handler sees it.
-func (s *Server) retrackConn(raw, wrapped net.Conn) {
-	s.sessionsMu.Lock()
-	delete(s.acceptedConns, raw)
-	s.acceptedConns[wrapped] = struct{}{}
-	s.sessionsMu.Unlock()
-}
diff --git a/client/vnc/server/server.go b/client/vnc/server/server.go
index cae818346..df1b1e92a 100644
--- a/client/vnc/server/server.go
+++ b/client/vnc/server/server.go
@@ -618,6 +618,16 @@ func (s *Server) trackConn(c net.Conn) {
 	s.sessionsMu.Unlock()
 }
 
+// retrackConn replaces a tracked raw connection with the wrapper its handler
+// will actually hold, so shutdown, connAuth registration and the handler's own
+// untrackConn all agree on which object is registered.
+func (s *Server) retrackConn(raw, wrapped net.Conn) {
+	s.sessionsMu.Lock()
+	delete(s.acceptedConns, raw)
+	s.acceptedConns[wrapped] = struct{}{}
+	s.sessionsMu.Unlock()
+}
+
 // untrackConn forgets a connection once its handler is returning.
 func (s *Server) untrackConn(c net.Conn) {
 	s.sessionsMu.Lock()
@@ -1126,8 +1136,12 @@ func (s *Server) handleConnection(conn net.Conn) {
 	}
 	// Behind the daemon the accepted address is only the local socket; the
 	// remote peer is the one the daemon vouched for in the grant.
-	conn = withGrantPeer(conn, grant)
-	connLog = s.log.WithField("remote", conn.RemoteAddr().String())
+	if wrapped := withGrantPeer(conn, grant); wrapped != conn {
+		s.retrackConn(conn, wrapped)
+		defer s.untrackConn(wrapped)
+		conn = wrapped
+		connLog = s.log.WithField("remote", conn.RemoteAddr().String())
+	}
 	header, err := s.readConnectionHeader(conn)
 	if err != nil {
 		connLog.Infof("VNC connection rejected: header read failed: %v", err)
diff --git a/client/vnc/server/server_test.go b/client/vnc/server/server_test.go
index 640a72c58..913c77e0c 100644
--- a/client/vnc/server/server_test.go
+++ b/client/vnc/server/server_test.go
@@ -357,6 +357,16 @@ func TestAgentToken_MatchAllowsHandshake(t *testing.T) {
 	sessions := srv.ActiveSessions()
 	require.Len(t, sessions, 1, "one session must be active")
 	assert.Equal(t, peer, sessions[0].RemoteAddress, "the session must report the grant's peer")
+
+	// Shutdown and connAuth cleanup key on the tracked connection, so it must
+	// be the wrapper the handler registers everything under, not the raw one.
+	srv.sessionsMu.Lock()
+	var tracked []string
+	for c := range srv.acceptedConns {
+		tracked = append(tracked, c.RemoteAddr().String())
+	}
+	srv.sessionsMu.Unlock()
+	assert.Equal(t, []string{peer}, tracked, "the wrapped connection must be the tracked one")
 }
 
 func TestSessionMode_RejectedWhenNoVMGR(t *testing.T) {