Make the token-leak, VNC port-scoping and bidirectional-SSH tests able to fail

client/vnc/server/agent_handshake_test.go

@@ -5,6 +5,7 @@ package server
 import (
 	"bytes"
 	"net"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -84,17 +85,25 @@ func TestAgentHandshake_TokenNeverSent(t *testing.T) {
 	defer daemonSide.Close()
 	defer agentSide.Close()
 
-	// Tee everything the daemon writes so it can be searched afterwards.
-	var sent bytes.Buffer
+	// Tee both directions on the agent's end: what it reads is everything the
+	// daemon sent, what it writes is its challenge and reply. Either side
+	// leaking the token is the failure this test is for.
+	var wire bytes.Buffer
+	var mu sync.Mutex
+	done := make(chan struct{})
 	go func() {
-		_, _ = agentServerHandshake(&teeConn{Conn: agentSide, read: &sent}, token)
+		defer close(done)
+		_, _ = agentServerHandshake(&teeConn{Conn: agentSide, mu: &mu, read: &wire, written: &wire}, token)
 	}()
 
 	require.NoError(t, agentClientHandshake(daemonSide, token, false))
+	<-done
+	mu.Lock()
+	defer mu.Unlock()
 	// bytes.Contains, not assert.NotContains: testify compares a []byte
 	// haystack element-wise, and a []byte is never an element of a []byte, so
 	// the assertion held whatever crossed the wire — including the whole token.
-	assert.False(t, bytes.Contains(sent.Bytes(), token), "the token must not cross the socket")
+	assert.False(t, bytes.Contains(wire.Bytes(), token), "the token must not cross the socket")
 }
 
 // A tag is bound to the nonce it answered, so replaying one against a fresh
@@ -115,16 +124,30 @@ func TestAgentMAC_IsBoundToNonceAndLabel(t *testing.T) {
 		"the two directions must not share a tag, or one could be replayed as the other")
 }
 
-// teeConn records everything read from the wrapped connection.
+// teeConn records every byte read from and written to the wrapped connection.
 type teeConn struct {
 	net.Conn
-	read *bytes.Buffer
+	mu      *sync.Mutex
+	read    *bytes.Buffer
+	written *bytes.Buffer
 }
 
 func (c *teeConn) Read(b []byte) (int, error) {
 	n, err := c.Conn.Read(b)
-	if n > 0 {
+	if n > 0 && c.read != nil {
+		c.mu.Lock()
 		c.read.Write(b[:n])
+		c.mu.Unlock()
+	}
+	return n, err
+}
+
+func (c *teeConn) Write(b []byte) (int, error) {
+	n, err := c.Conn.Write(b)
+	if n > 0 && c.written != nil {
+		c.mu.Lock()
+		c.written.Write(b[:n])
+		c.mu.Unlock()
 	}
 	return n, err
 }

management/server/types/networkmap_components_correctness_test.go

@@ -1421,6 +1421,10 @@ func TestPeerSSHEnabledFromPolicies_MatchesMap_Sweep(t *testing.T) {
 func TestPeerSSHEnabledFromPolicies_MatchesMap_BidirectionalNetbirdSSH(t *testing.T) {
 	account, validatedPeers := scalableTestAccountWithoutDefaultPolicy(20, 2)
 	account.Groups["ssh-users"] = &types.Group{ID: "ssh-users", Name: "SSH Users", Peers: []string{}}
+	// A member, so the rule authorizes somebody: with an empty group every map
+	// carries no users either way, and the source-side assertion below could
+	// not fail.
+	account.Users["user-dev"] = &types.User{Id: "user-dev", Role: types.UserRoleUser, AccountID: "test-account", AutoGroups: []string{"ssh-users"}}
 	account.Policies = append(account.Policies, &types.Policy{
 		ID: "policy-ssh-bidi", Name: "SSH Access", Enabled: true, AccountID: "test-account",
 		Rules: []*types.PolicyRule{{
@@ -1446,4 +1450,6 @@ func TestPeerSSHEnabledFromPolicies_MatchesMap_BidirectionalNetbirdSSH(t *testin
 	destination := componentsNetworkMap(account, "peer-10", validatedPeers)
 	require.NotNil(t, destination)
 	assert.True(t, destination.EnableSSH, "the destination-side peer must have SSH enabled")
+	assert.NotEmpty(t, destination.AuthorizedUsers["root"],
+		"the destination-side peer must carry the group's user, or the source-side check proves nothing")
 }

shared/management/networkmap/envelope_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net"
 	"net/netip"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -17,6 +18,7 @@ import (
 	nbnetworkmap "github.com/netbirdio/netbird/shared/management/networkmap"
 	"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
 	"github.com/netbirdio/netbird/shared/management/proto"
+	sharedtypes "github.com/netbirdio/netbird/shared/management/types"
 )
 
 // TestEnvelopeToNetworkMap_RoundTrip exercises the full client-side pipeline:
@@ -123,12 +125,32 @@ func TestEnvelopeToNetworkMap_VNCPolicyProducesVncAuth(t *testing.T) {
 
 	// The VNC marker protocol must not reach the firewall rules, for the same
 	// reason NetbirdSSH must not: agents fall into UNKNOWN-protocol handling.
+	// And the rewrite must stay scoped to the VNC port: a portless netbird-vnc
+	// rule that degraded into plain TCP would open every port on the peer.
+	// Requiring rules at all is what keeps the loop from passing vacuously.
+	require.NotEmpty(t, result.NetworkMap.FirewallRules, "a netbird-vnc policy must produce firewall rules")
+	wantPort := strconv.Itoa(sharedtypes.VNCInternalPort)
 	for i, fr := range result.NetworkMap.FirewallRules {
-		require.NotEqualf(t, proto.RuleProtocol_NETBIRD_VNC, fr.Protocol,
+		require.Equalf(t, proto.RuleProtocol_TCP, fr.Protocol,
 			"FirewallRules[%d].Protocol must be the rewritten TCP, not NETBIRD_VNC", i)
+		require.Equalf(t, wantPort, firewallRulePort(fr),
+			"FirewallRules[%d] must be scoped to the VNC port", i)
 	}
 }
 
+// firewallRulePort returns the single port a firewall rule opens, from either
+// the legacy Port field or a single-port PortInfo, or "" when it opens a range
+// or every port.
+func firewallRulePort(fr *proto.FirewallRule) string {
+	if fr.Port != "" {
+		return fr.Port
+	}
+	if p, ok := fr.GetPortInfo().GetPortSelection().(*proto.PortInfo_Port); ok {
+		return strconv.FormatUint(uint64(p.Port), 10)
+	}
+	return ""
+}
+
 // A policy that authorizes nothing VNC-related must leave VncAuth unset, so the
 // authorizer keeps refusing rather than being handed an empty allow-list to
 // interpret.