From 1f4b645d5bb2413fbfeb139478e96b9c45b2da29 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 1 Nov 2025 09:05:37 +0000
Subject: [PATCH] Refactor: Extract route filtering logic into helper function

Addressed code review feedback:
- Extracted duplicated exit route filtering logic into buildRouteFilterFunc
- Reused exitNode lookup in integration test to avoid duplication
- Added missing matcher package import

This improves code maintainability while preserving the same functionality.

Co-authored-by: kradalby <98431+kradalby@users.noreply.github.com>
---
 hscontrol/mapper/builder.go | 68 ++++++++++++++++---------------------
 integration/route_test.go   | 14 +++-----
 2 files changed, 34 insertions(+), 48 deletions(-)

diff --git a/hscontrol/mapper/builder.go b/hscontrol/mapper/builder.go
index 69171f7a..6fe6fc84 100644
--- a/hscontrol/mapper/builder.go
+++ b/hscontrol/mapper/builder.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/juanfont/headscale/hscontrol/policy"
+	"github.com/juanfont/headscale/hscontrol/policy/matcher"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/views"
@@ -67,6 +68,33 @@ func (b *MapResponseBuilder) WithCapabilityVersion(capVer tailcfg.CapabilityVers
 	return b
 }
 
+// buildRouteFilterFunc creates a route filter function that includes both primary and exit routes.
+// It filters routes based on ACL policy to ensure only authorized routes are visible.
+func (b *MapResponseBuilder) buildRouteFilterFunc(
+	viewingNode types.NodeView,
+	matchers []matcher.Match,
+) routeFilterFunc {
+	return func(id types.NodeID) []netip.Prefix {
+		// Get the peer node to check for exit routes
+		peer, ok := b.mapper.state.GetNodeByID(id)
+		if !ok {
+			return nil
+		}
+		
+		// Start with primary routes (subnet routes, but not exit routes)
+		routes := policy.ReduceRoutes(viewingNode, b.mapper.state.GetNodePrimaryRoutes(id), matchers)
+		
+		// Also filter exit routes through policy
+		// Only add exit routes if the viewing node has permission to use them
+		if exitRoutes := peer.ExitRoutes(); len(exitRoutes) > 0 {
+			filteredExitRoutes := policy.ReduceRoutes(viewingNode, exitRoutes, matchers)
+			routes = append(routes, filteredExitRoutes...)
+		}
+		
+		return routes
+	}
+}
+
 // WithSelfNode adds the requesting node to the response.
 func (b *MapResponseBuilder) WithSelfNode() *MapResponseBuilder {
 	nv, ok := b.mapper.state.GetNodeByID(b.nodeID)
@@ -78,25 +106,7 @@ func (b *MapResponseBuilder) WithSelfNode() *MapResponseBuilder {
 	_, matchers := b.mapper.state.Filter()
 	tailnode, err := tailNode(
 		nv, b.capVer, b.mapper.state,
-		func(id types.NodeID) []netip.Prefix {
-			// Get the peer node to check for exit routes
-			peer, ok := b.mapper.state.GetNodeByID(id)
-			if !ok {
-				return nil
-			}
-			
-			// Start with primary routes (subnet routes, but not exit routes)
-			routes := policy.ReduceRoutes(nv, b.mapper.state.GetNodePrimaryRoutes(id), matchers)
-			
-			// Also filter exit routes through policy
-			// Only add exit routes if the viewing node (self in this case) has permission to use them
-			if exitRoutes := peer.ExitRoutes(); len(exitRoutes) > 0 {
-				filteredExitRoutes := policy.ReduceRoutes(nv, exitRoutes, matchers)
-				routes = append(routes, filteredExitRoutes...)
-			}
-			
-			return routes
-		},
+		b.buildRouteFilterFunc(nv, matchers),
 		b.mapper.cfg)
 	if err != nil {
 		b.addError(err)
@@ -269,25 +279,7 @@ func (b *MapResponseBuilder) buildTailPeers(peers views.Slice[types.NodeView]) (
 
 	tailPeers, err := tailNodes(
 		changedViews, b.capVer, b.mapper.state,
-		func(id types.NodeID) []netip.Prefix {
-			// Get the peer node to check for exit routes
-			peer, ok := b.mapper.state.GetNodeByID(id)
-			if !ok {
-				return nil
-			}
-			
-			// Start with primary routes (subnet routes, but not exit routes)
-			routes := policy.ReduceRoutes(node, b.mapper.state.GetNodePrimaryRoutes(id), matchers)
-			
-			// Also filter exit routes through policy
-			// Only add exit routes if the viewing node has permission to use them
-			if exitRoutes := peer.ExitRoutes(); len(exitRoutes) > 0 {
-				filteredExitRoutes := policy.ReduceRoutes(node, exitRoutes, matchers)
-				routes = append(routes, filteredExitRoutes...)
-			}
-			
-			return routes
-		},
+		b.buildRouteFilterFunc(node, matchers),
 		b.mapper.cfg)
 	if err != nil {
 		return nil, err
diff --git a/integration/route_test.go b/integration/route_test.go
index df5379a0..ed11d5e5 100644
--- a/integration/route_test.go
+++ b/integration/route_test.go
@@ -3126,14 +3126,16 @@ require.NoErrorf(t, err, "failed to advertise exit node: %s", err)
 
 // Wait for the exit node to be registered
 var nodes []*v1.Node
+var exitNode *v1.Node
+exitStatus := exitClient.MustStatus()
+
 assert.EventuallyWithT(t, func(c *assert.CollectT) {
 nodes, err = headscale.ListNodes()
 assert.NoError(c, err)
 assert.Len(c, nodes, 3)
 
 // Find the exit node
-var exitNode *v1.Node
-exitStatus := exitClient.MustStatus()
+exitNode = nil
 for _, node := range nodes {
 if node.GetName() == exitStatus.Self.HostName {
 exitNode = node
@@ -3148,14 +3150,6 @@ assert.Len(c, exitNode.GetAvailableRoutes(), 2, "exit node should advertise 2 ro
 }, 10*time.Second, 500*time.Millisecond, "waiting for exit node advertisement")
 
 // Approve the exit routes
-var exitNode *v1.Node
-exitStatus := exitClient.MustStatus()
-for _, node := range nodes {
-if node.GetName() == exitStatus.Self.HostName {
-exitNode = node
-break
-}
-}
 require.NotNil(t, exitNode, "exit node not found after advertisement")
 
 _, err = headscale.ApproveRoutes(exitNode.GetId(), []netip.Prefix{tsaddr.AllIPv4(), tsaddr.AllIPv6()})