From 096ac31bb3c9e6aa041975ce1695f364db22509d Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Mon, 11 Sep 2023 11:45:46 -0500
Subject: [PATCH] handle route updates correctly

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/db/routes.go  |  6 ++++++
 hscontrol/poll.go       | 18 +++++++++++++++++-
 integration/cli_test.go | 13 ++++++-------
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/hscontrol/db/routes.go b/hscontrol/db/routes.go
index af6b744f5b..2d51d18311 100644
--- a/hscontrol/db/routes.go
+++ b/hscontrol/db/routes.go
@@ -278,6 +278,12 @@ func (hsdb *HSDatabase) saveMachineRoutes(machine *types.Machine) error {
 		advertisedRoutes[prefix] = false
 	}
 
+	log.Trace().
+		Str("machine", machine.Hostname).
+		Interface("advertisedRoutes", advertisedRoutes).
+		Interface("currentRoutes", currentRoutes).
+		Msg("updating routes")
+
 	for pos, route := range currentRoutes {
 		if _, ok := advertisedRoutes[netip.Prefix(route.Prefix)]; ok {
 			if !route.Advertised {
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
index 4df3865e5a..45abe56424 100644
--- a/hscontrol/poll.go
+++ b/hscontrol/poll.go
@@ -66,6 +66,9 @@ func (h *Headscale) handlePoll(
 ) {
 	logInfo, logErr := logPollFunc(mapRequest, machine, isNoise)
 
+	// This is the mechanism where the node gives us inforamtion about its
+	// current configuration.
+	//
 	// If OmitPeers is true, Stream is false, and ReadOnly is false,
 	// then te server will let clients update their endpoints without
 	// breaking existing long-polling (Stream == true) connections.
@@ -84,8 +87,11 @@ func (h *Headscale) handlePoll(
 			Msg("Received endpoint update")
 
 		now := time.Now().UTC()
-		machine.Endpoints = mapRequest.Endpoints
 		machine.LastSeen = &now
+		machine.Hostname = mapRequest.Hostinfo.Hostname
+		machine.HostInfo = types.HostInfo(*mapRequest.Hostinfo)
+		machine.DiscoKey = util.DiscoPublicKeyStripPrefix(mapRequest.DiscoKey)
+		machine.Endpoints = mapRequest.Endpoints
 
 		if err := h.db.MachineSave(machine); err != nil {
 			logErr(err, "Failed to persist/update machine in the database")
@@ -94,6 +100,14 @@ func (h *Headscale) handlePoll(
 			return
 		}
 
+		err := h.db.SaveMachineRoutes(machine)
+		if err != nil {
+			logErr(err, "Error processing machine routes")
+			http.Error(writer, "", http.StatusInternalServerError)
+
+			return
+		}
+
 		h.nodeNotifier.NotifyWithIgnore(
 			types.StateUpdate{
 				Type:    types.StatePeerChanged,
@@ -134,6 +148,8 @@ func (h *Headscale) handlePoll(
 		return
 	}
 
+	now := time.Now().UTC()
+	machine.LastSeen = &now
 	machine.Hostname = mapRequest.Hostinfo.Hostname
 	machine.HostInfo = types.HostInfo(*mapRequest.Hostinfo)
 	machine.DiscoKey = util.DiscoPublicKeyStripPrefix(mapRequest.DiscoKey)
diff --git a/integration/cli_test.go b/integration/cli_test.go
index 5a27d51726..ef0ed70175 100644
--- a/integration/cli_test.go
+++ b/integration/cli_test.go
@@ -413,14 +413,12 @@ func TestEnablingRoutes(t *testing.T) {
 	// advertise routes using the up command
 	for i, client := range allClients {
 		routeStr := fmt.Sprintf("10.0.%d.0/24", i)
-		hostname, _ := client.FQDN()
-		_, _, err = client.Execute([]string{
+		command := []string{
 			"tailscale",
-			"up",
-			fmt.Sprintf("--advertise-routes=%s", routeStr),
-			"-login-server", headscale.GetEndpoint(),
-			"--hostname", hostname,
-		})
+			"set",
+			"--advertise-routes=" + routeStr,
+		}
+		_, _, err := client.Execute(command)
 		assertNoErrf(t, "failed to advertise route: %s", err)
 	}
 
@@ -474,6 +472,7 @@ func TestEnablingRoutes(t *testing.T) {
 		&enablingRoutes,
 	)
 	assertNoErr(t, err)
+	assert.Len(t, enablingRoutes, 3)
 
 	for _, route := range enablingRoutes {
 		assert.Equal(t, route.Advertised, true)