From 1c6bfc503c9b8f7adbaaae662138c9a0fb89b213 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@dalby.cc>
Date: Thu, 2 May 2024 11:53:16 +0200
Subject: [PATCH] fix preauth key logging in as previous user (#1920)

* add test case to reproduce #1885

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

* fix preauth key issue logging in as wrong user

Fixes #1885

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

* add test to gh

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

---------

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 .github/workflows/test-integration.yaml |  1 +
 hscontrol/auth.go                       |  9 ++-
 integration/cli_test.go                 | 95 +++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test-integration.yaml b/.github/workflows/test-integration.yaml
index d93aaca2..9581bada 100644
--- a/.github/workflows/test-integration.yaml
+++ b/.github/workflows/test-integration.yaml
@@ -26,6 +26,7 @@ jobs:
           - TestPreAuthKeyCommand
           - TestPreAuthKeyCommandWithoutExpiry
           - TestPreAuthKeyCommandReusableEphemeral
+          - TestPreAuthKeyCorrectUserLoggedInCommand
           - TestApiKeyCommand
           - TestNodeTagCommand
           - TestNodeAdvertiseTagNoACLCommand
diff --git a/hscontrol/auth.go b/hscontrol/auth.go
index 0679d72e..dab9ff42 100644
--- a/hscontrol/auth.go
+++ b/hscontrol/auth.go
@@ -315,13 +315,16 @@ func (h *Headscale) handleAuthKey(
 
 		node.NodeKey = nodeKey
 		node.AuthKeyID = uint(pak.ID)
-		err := h.db.NodeSetExpiry(node.ID, registerRequest.Expiry)
+		node.Expiry = &registerRequest.Expiry
+		node.User = pak.User
+		node.UserID = pak.UserID
+		err := h.db.DB.Save(node).Error
 		if err != nil {
 			log.Error().
 				Caller().
 				Str("node", node.Hostname).
 				Err(err).
-				Msg("Failed to refresh node")
+				Msg("failed to save node after logging in with auth key")
 
 			return
 		}
@@ -344,7 +347,7 @@ func (h *Headscale) handleAuthKey(
 		}
 
 		ctx := types.NotifyCtx(context.Background(), "handle-authkey", "na")
-		h.nodeNotifier.NotifyWithIgnore(ctx, types.StateUpdateExpire(node.ID, registerRequest.Expiry), node.ID)
+		h.nodeNotifier.NotifyAll(ctx, types.StateUpdate{Type: types.StatePeerChanged, ChangeNodes: []types.NodeID{node.ID}})
 	} else {
 		now := time.Now().UTC()
 
diff --git a/integration/cli_test.go b/integration/cli_test.go
index 24e3b19b..57edf58e 100644
--- a/integration/cli_test.go
+++ b/integration/cli_test.go
@@ -388,6 +388,101 @@ func TestPreAuthKeyCommandReusableEphemeral(t *testing.T) {
 	assert.Len(t, listedPreAuthKeys, 3)
 }
 
+func TestPreAuthKeyCorrectUserLoggedInCommand(t *testing.T) {
+	IntegrationSkip(t)
+	t.Parallel()
+
+	user1 := "user1"
+	user2 := "user2"
+
+	scenario, err := NewScenario(dockertestMaxWait())
+	assertNoErr(t, err)
+	defer scenario.Shutdown()
+
+	spec := map[string]int{
+		user1: 1,
+		user2: 0,
+	}
+
+	err = scenario.CreateHeadscaleEnv(spec, []tsic.Option{}, hsic.WithTestName("clipak"))
+	assertNoErr(t, err)
+
+	headscale, err := scenario.Headscale()
+	assertNoErr(t, err)
+
+	var user2Key v1.PreAuthKey
+
+	err = executeAndUnmarshal(
+		headscale,
+		[]string{
+			"headscale",
+			"preauthkeys",
+			"--user",
+			user2,
+			"create",
+			"--reusable",
+			"--expiration",
+			"24h",
+			"--output",
+			"json",
+			"--tags",
+			"tag:test1,tag:test2",
+		},
+		&user2Key,
+	)
+	assertNoErr(t, err)
+
+	allClients, err := scenario.ListTailscaleClients()
+	assertNoErrListClients(t, err)
+
+	assert.Len(t, allClients, 1)
+
+	client := allClients[0]
+
+	// Log out from user1
+	err = client.Logout()
+	assertNoErr(t, err)
+
+	err = scenario.WaitForTailscaleLogout()
+	assertNoErr(t, err)
+
+	status, err := client.Status()
+	assertNoErr(t, err)
+	if status.BackendState == "Starting" || status.BackendState == "Running" {
+		t.Fatalf("expected node to be logged out, backend state: %s", status.BackendState)
+	}
+
+	err = client.Login(headscale.GetEndpoint(), user2Key.GetKey())
+	assertNoErr(t, err)
+
+	status, err = client.Status()
+	assertNoErr(t, err)
+	if status.BackendState != "Running" {
+		t.Fatalf("expected node to be logged in, backend state: %s", status.BackendState)
+	}
+
+	if status.Self.UserID.String() != "userid:2" {
+		t.Fatalf("expected node to be logged in as userid:2, got: %s", status.Self.UserID.String())
+	}
+
+	var listNodes []v1.Node
+	err = executeAndUnmarshal(
+		headscale,
+		[]string{
+			"headscale",
+			"nodes",
+			"list",
+			"--output",
+			"json",
+		},
+		&listNodes,
+	)
+	assert.Nil(t, err)
+	assert.Len(t, listNodes, 1)
+
+	assert.Equal(t, "user2", listNodes[0].User.Name)
+}
+
 func TestApiKeyCommand(t *testing.T) {
 	IntegrationSkip(t)
 	t.Parallel()