From a613e8a8ce47c3e141219d7d5942a8fb2bbe5723 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Sat, 1 Nov 2025 22:46:10 +0100
Subject: [PATCH] state: add test for endpoint storage and propagation

Add TestEndpointStorageInNodeStore to verify that endpoints
sent in MapRequest are correctly stored in NodeStore and
propagated to peers. This test reproduces the issue from #2846.
---
 hscontrol/state/endpoint_test.go | 108 +++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 hscontrol/state/endpoint_test.go

diff --git a/hscontrol/state/endpoint_test.go b/hscontrol/state/endpoint_test.go
new file mode 100644
index 00000000..119933e6
--- /dev/null
+++ b/hscontrol/state/endpoint_test.go
@@ -0,0 +1,108 @@
+package state
+
+import (
+	"net/netip"
+	"testing"
+
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"tailscale.com/tailcfg"
+)
+
+// TestEndpointStorageInNodeStore verifies that endpoints sent in MapRequest via ApplyPeerChange
+// are correctly stored in the NodeStore and can be retrieved for sending to peers.
+// This test reproduces the issue reported in https://github.com/juanfont/headscale/issues/2846
+func TestEndpointStorageInNodeStore(t *testing.T) {
+	// Create two test nodes
+	node1 := createTestNode(1, 1, "test-user", "node1")
+	node2 := createTestNode(2, 1, "test-user", "node2")
+
+	// Create NodeStore with allow-all peers function
+	store := NewNodeStore(nil, allowAllPeersFunc)
+	store.Start()
+	defer store.Stop()
+
+	// Add both nodes to NodeStore
+	store.PutNode(node1)
+	store.PutNode(node2)
+
+	// Create a MapRequest with endpoints for node1
+	endpoints := []netip.AddrPort{
+		netip.MustParseAddrPort("192.168.1.1:41641"),
+		netip.MustParseAddrPort("10.0.0.1:41641"),
+	}
+
+	mapReq := tailcfg.MapRequest{
+		NodeKey:   node1.NodeKey,
+		DiscoKey:  node1.DiscoKey,
+		Endpoints: endpoints,
+		Hostinfo: &tailcfg.Hostinfo{
+			Hostname: "node1",
+		},
+	}
+
+	// Simulate what UpdateNodeFromMapRequest does: create PeerChange and apply it
+	peerChange := node1.PeerChangeFromMapRequest(mapReq)
+
+	// Verify PeerChange has endpoints
+	require.NotNil(t, peerChange.Endpoints, "PeerChange should contain endpoints")
+	assert.Equal(t, len(endpoints), len(peerChange.Endpoints),
+		"PeerChange should have same number of endpoints as MapRequest")
+
+	// Apply the PeerChange via NodeStore.UpdateNode
+	updatedNode, ok := store.UpdateNode(node1.ID, func(n *types.Node) {
+		n.ApplyPeerChange(&peerChange)
+	})
+	require.True(t, ok, "UpdateNode should succeed")
+	require.True(t, updatedNode.Valid(), "Updated node should be valid")
+
+	// Verify endpoints are in the updated node view
+	storedEndpoints := updatedNode.Endpoints().AsSlice()
+	assert.Equal(t, len(endpoints), len(storedEndpoints),
+		"NodeStore should have same number of endpoints as sent")
+
+	if len(storedEndpoints) == len(endpoints) {
+		for i, ep := range endpoints {
+			assert.Equal(t, ep, storedEndpoints[i],
+				"Endpoint %d should match", i)
+		}
+	}
+
+	// Verify we can retrieve the node again and endpoints are still there
+	retrievedNode, found := store.GetNode(node1.ID)
+	require.True(t, found, "node1 should exist in NodeStore")
+
+	retrievedEndpoints := retrievedNode.Endpoints().AsSlice()
+	assert.Equal(t, len(endpoints), len(retrievedEndpoints),
+		"Retrieved node should have same number of endpoints")
+
+	// Verify that when we get node1 as a peer of node2, it has endpoints
+	// This is the critical part that was failing in the bug report
+	peers := store.ListPeers(node2.ID)
+	require.Greater(t, peers.Len(), 0, "node2 should have at least one peer")
+
+	// Find node1 in the peer list
+	var node1Peer types.NodeView
+	foundPeer := false
+	for _, peer := range peers.All() {
+		if peer.ID() == node1.ID {
+			node1Peer = peer
+			foundPeer = true
+			break
+		}
+	}
+	require.True(t, foundPeer, "node1 should be in node2's peer list")
+
+	// Check that node1's endpoints are available in the peer view
+	peerEndpoints := node1Peer.Endpoints().AsSlice()
+	assert.Equal(t, len(endpoints), len(peerEndpoints),
+		"Peer view should have same number of endpoints as sent")
+
+	if len(peerEndpoints) == len(endpoints) {
+		for i, ep := range endpoints {
+			assert.Equal(t, ep, peerEndpoints[i],
+				"Peer endpoint %d should match", i)
+		}
+	}
+}