From 495ff53573d15eae9d403018035e17a2a4275992 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Fri, 1 Dec 2023 14:14:47 +0100
Subject: [PATCH] update nodes if services are changed

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/auth.go           |  2 +-
 hscontrol/mapper/tail.go    |  1 +
 hscontrol/poll.go           | 23 +++++++++++++++++++++++
 integration/general_test.go | 37 +++++++++++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/hscontrol/auth.go b/hscontrol/auth.go
index 573c4cb02d..749ad1ee23 100644
--- a/hscontrol/auth.go
+++ b/hscontrol/auth.go
@@ -100,7 +100,7 @@ func (h *Headscale) handleRegister(
 			}
 		}
 
-		logInfo("Node not found in databas, creating new")
+		logInfo("Node not found in database, creating new")
 
 		givenName, err := h.db.GenerateGivenName(
 			machineKey,
diff --git a/hscontrol/mapper/tail.go b/hscontrol/mapper/tail.go
index 6aa688ad93..fba613c10f 100644
--- a/hscontrol/mapper/tail.go
+++ b/hscontrol/mapper/tail.go
@@ -99,6 +99,7 @@ func tailNode(
 			strconv.FormatUint(node.ID, util.Base10),
 		), // in headscale, unlike tailcontrol server, IDs are permanent
 		Name: hostname,
+		Cap:  capVer,
 
 		User: tailcfg.UserID(node.UserID),
 
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
index 0be9179e1b..79a584859c 100644
--- a/hscontrol/poll.go
+++ b/hscontrol/poll.go
@@ -101,8 +101,17 @@ func (h *Headscale) handlePoll(
 			oldRoutes := node.Hostinfo.RoutableIPs
 			newRoutes := mapRequest.Hostinfo.RoutableIPs
 
+			oldServicesCount := len(node.Hostinfo.Services)
+			newServicesCount := len(mapRequest.Hostinfo.Services)
+
 			node.Hostinfo = mapRequest.Hostinfo
 
+			sendUpdate := false
+
+			// Route changes come as part of Hostinfo, which means that
+			// when an update comes, the Node Route logic need to run.
+			// This will require a "change" in comparison to a "patch",
+			// which is more costly.
 			if !xslices.Equal(oldRoutes, newRoutes) {
 				err := h.db.SaveNodeRoutes(node)
 				if err != nil {
@@ -119,6 +128,20 @@ func (h *Headscale) handlePoll(
 					return
 				}
 
+				sendUpdate = true
+			}
+
+			// Services is mostly useful for discovery and not critical,
+			// except for peerapi, which is how nodes talk to eachother.
+			// If peerapi was not part of the initial mapresponse, we
+			// need to make sure its sent out later as it is needed for
+			// Taildrop.
+			// TODO(kradalby): Length comparison is a bit naive, replace.
+			if oldServicesCount != newServicesCount {
+				sendUpdate = true
+			}
+
+			if sendUpdate {
 				stateUpdate := types.StateUpdate{
 					Type:        types.StatePeerChanged,
 					ChangeNodes: types.Nodes{node},
diff --git a/integration/general_test.go b/integration/general_test.go
index 60ffd83848..dd32a35f0e 100644
--- a/integration/general_test.go
+++ b/integration/general_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/rs/zerolog/log"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
+	"tailscale.com/client/tailscale/apitype"
 )
 
 func TestPingAllByIP(t *testing.T) {
@@ -311,6 +312,42 @@ func TestTaildrop(t *testing.T) {
 	_, err = scenario.ListTailscaleClientsFQDNs()
 	assertNoErrListFQDN(t, err)
 
+	for _, client := range allClients {
+		if !strings.Contains(client.Hostname(), "head") {
+			command := []string{"apk", "add", "curl"}
+			_, _, err := client.Execute(command)
+			if err != nil {
+				t.Fatalf("failed to install curl on %s, err: %s", client.Hostname(), err)
+			}
+
+		}
+		curlCommand := []string{"curl", "--unix-socket", "/var/run/tailscale/tailscaled.sock", "http://local-tailscaled.sock/localapi/v0/file-targets"}
+		err = retry(10, 1*time.Second, func() error {
+			result, _, err := client.Execute(curlCommand)
+			if err != nil {
+				return err
+			}
+			var fts []apitype.FileTarget
+			err = json.Unmarshal([]byte(result), &fts)
+			if err != nil {
+				return err
+			}
+
+			if len(fts) != len(allClients)-1 {
+				ftStr := fmt.Sprintf("FileTargets for %s:\n", client.Hostname())
+				for _, ft := range fts {
+					ftStr += fmt.Sprintf("\t%s\n", ft.Node.Name)
+				}
+				return fmt.Errorf("client %s does not have all its peers as FileTargets, got %d, want: %d\n%s", client.Hostname(), len(fts), len(allClients)-1, ftStr)
+			}
+
+			return err
+		})
+		if err != nil {
+			t.Errorf("failed to query localapi for filetarget on %s, err: %s", client.Hostname(), err)
+		}
+	}
+
 	for _, client := range allClients {
 		command := []string{"touch", fmt.Sprintf("/tmp/file_from_%s", client.Hostname())}