From a0cc20bef99e435261bb081a6f59b81d42c4c858 Mon Sep 17 00:00:00 2001
From: Juan Font <juanfontalonso@gmail.com>
Date: Sun, 3 Mar 2024 10:11:10 +0000
Subject: [PATCH] Ephemeral keys can now be reusable and non-reusable

Fixes the issue reported in #1712. In Tailscale SaaS, ephemeral keys can be single-user or reusable. Until now, our ephemerals were only reusable. This PR makes us adhere to the .com behaviour.
---
 cmd/headscale/cli/preauthkeys.go  |  9 +------
 hscontrol/db/preauth_keys.go      |  2 +-
 hscontrol/db/preauth_keys_test.go | 39 ++++++++++++++++++++++++++++---
 3 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go
index c8dd2adc..cc3b1b76 100644
--- a/cmd/headscale/cli/preauthkeys.go
+++ b/cmd/headscale/cli/preauthkeys.go
@@ -107,13 +107,6 @@ var listPreAuthKeys = &cobra.Command{
 				expiration = ColourTime(key.GetExpiration().AsTime())
 			}
 
-			var reusable string
-			if key.GetEphemeral() {
-				reusable = "N/A"
-			} else {
-				reusable = fmt.Sprintf("%v", key.GetReusable())
-			}
-
 			aclTags := ""
 
 			for _, tag := range key.GetAclTags() {
@@ -125,7 +118,7 @@ var listPreAuthKeys = &cobra.Command{
 			tableData = append(tableData, []string{
 				key.GetId(),
 				key.GetKey(),
-				reusable,
+				strconv.FormatBool(key.GetReusable()),
 				strconv.FormatBool(key.GetEphemeral()),
 				strconv.FormatBool(key.GetUsed()),
 				expiration,
diff --git a/hscontrol/db/preauth_keys.go b/hscontrol/db/preauth_keys.go
index 0fdb8221..d1d94bbe 100644
--- a/hscontrol/db/preauth_keys.go
+++ b/hscontrol/db/preauth_keys.go
@@ -196,7 +196,7 @@ func ValidatePreAuthKey(tx *gorm.DB, k string) (*types.PreAuthKey, error) {
 		return nil, ErrPreAuthKeyExpired
 	}
 
-	if pak.Reusable || pak.Ephemeral { // we don't need to check if has been used before
+	if pak.Reusable { // we don't need to check if has been used before
 		return &pak, nil
 	}
 
diff --git a/hscontrol/db/preauth_keys_test.go b/hscontrol/db/preauth_keys_test.go
index 003a396f..53cf37c4 100644
--- a/hscontrol/db/preauth_keys_test.go
+++ b/hscontrol/db/preauth_keys_test.go
@@ -123,7 +123,41 @@ func (*Suite) TestNotReusableNotBeingUsedKey(c *check.C) {
 	c.Assert(key.ID, check.Equals, pak.ID)
 }
 
-func (*Suite) TestEphemeralKey(c *check.C) {
+func (*Suite) TestEphemeralKeyReusable(c *check.C) {
+	user, err := db.CreateUser("test7")
+	c.Assert(err, check.IsNil)
+
+	pak, err := db.CreatePreAuthKey(user.Name, true, true, nil, nil)
+	c.Assert(err, check.IsNil)
+
+	now := time.Now().Add(-time.Second * 30)
+	node := types.Node{
+		ID:             0,
+		Hostname:       "testest",
+		UserID:         user.ID,
+		RegisterMethod: util.RegisterMethodAuthKey,
+		LastSeen:       &now,
+		AuthKeyID:      uint(pak.ID),
+	}
+	db.DB.Save(&node)
+
+	_, err = db.ValidatePreAuthKey(pak.Key)
+	c.Assert(err, check.IsNil)
+
+	_, err = db.getNode("test7", "testest")
+	c.Assert(err, check.IsNil)
+
+	db.DB.Transaction(func(tx *gorm.DB) error {
+		ExpireEphemeralNodes(tx, time.Second*20)
+		return nil
+	})
+
+	// The machine record should have been deleted
+	_, err = db.getNode("test7", "testest")
+	c.Assert(err, check.NotNil)
+}
+
+func (*Suite) TestEphemeralKeyNotReusable(c *check.C) {
 	user, err := db.CreateUser("test7")
 	c.Assert(err, check.IsNil)
 
@@ -142,8 +176,7 @@ func (*Suite) TestEphemeralKey(c *check.C) {
 	db.DB.Save(&node)
 
 	_, err = db.ValidatePreAuthKey(pak.Key)
-	// Ephemeral keys are by definition reusable
-	c.Assert(err, check.IsNil)
+	c.Assert(err, check.NotNil)
 
 	_, err = db.getNode("test7", "testest")
 	c.Assert(err, check.IsNil)