From 3b0749a320e442fe03ff6721fefed64c4551bdb7 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Wed, 9 Aug 2023 20:37:41 +0200
Subject: [PATCH] Update packetfilter when peers change

Previously we did not update the packet filter
when nodes changed, which would cause new nodes
to be missing from packet filters of old nodes.

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/mapper/mapper.go | 31 ++++++++++++++++++-------------
 hscontrol/types/machine.go | 10 ++++++++++
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/hscontrol/mapper/mapper.go b/hscontrol/mapper/mapper.go
index 1b6219a0..97b08797 100644
--- a/hscontrol/mapper/mapper.go
+++ b/hscontrol/mapper/mapper.go
@@ -382,28 +382,31 @@ func (m *Mapper) DERPMapResponse(
 func (m *Mapper) PeerChangedResponse(
 	mapRequest tailcfg.MapRequest,
 	machine *types.Machine,
-	machineKeys []uint64,
+	machineIDs []uint64,
 	pol *policy.ACLPolicy,
 ) ([]byte, error) {
 	var err error
-	changed := make(types.Machines, len(machineKeys))
+	changed := make(types.Machines, len(machineIDs))
 	lastSeen := make(map[tailcfg.NodeID]bool)
-	for idx, machineKey := range machineKeys {
-		peer, err := m.db.GetMachineByID(machineKey)
-		if err != nil {
-			return nil, err
-		}
 
-		changed[idx] = *peer
-
-		// We have just seen the node, let the peers update their list.
-		lastSeen[tailcfg.NodeID(peer.ID)] = true
+	peersList, err := m.db.ListPeers(machine)
+	if err != nil {
+		return nil, err
 	}
 
-	rules, _, err := policy.GenerateFilterAndSSHRules(
+	peers := peersList.IDMap()
+
+	for idx, machineID := range machineIDs {
+		changed[idx] = peers[machineID]
+
+		// We have just seen the node, let the peers update their list.
+		lastSeen[tailcfg.NodeID(machineID)] = true
+	}
+
+	rules, sshPolicy, err := policy.GenerateFilterAndSSHRules(
 		pol,
 		machine,
-		changed,
+		peersList,
 	)
 	if err != nil {
 		return nil, err
@@ -434,6 +437,8 @@ func (m *Mapper) PeerChangedResponse(
 
 	resp := m.baseMapResponse(machine)
 	resp.PeersChanged = tailPeers
+	resp.PacketFilter = policy.ReduceFilterRules(machine, rules)
+	resp.SSHPolicy = sshPolicy
 	// resp.PeerSeenChange = lastSeen
 
 	return m.marshalMapResponse(mapRequest, &resp, machine, mapRequest.Compress)
diff --git a/hscontrol/types/machine.go b/hscontrol/types/machine.go
index 4e5a940f..04522868 100644
--- a/hscontrol/types/machine.go
+++ b/hscontrol/types/machine.go
@@ -353,3 +353,13 @@ func (machines MachinesP) String() string {
 
 	return fmt.Sprintf("[ %s ](%d)", strings.Join(temp, ", "), len(temp))
 }
+
+func (machines Machines) IDMap() map[uint64]Machine {
+	ret := map[uint64]Machine{}
+
+	for _, machine := range machines {
+		ret[machine.ID] = machine
+	}
+
+	return ret
+}