From 30788e1a70969f541bfa699d93f675b24775a5bb Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Wed, 13 Oct 2021 18:13:26 +0200
Subject: [PATCH] Add AlreadyUsed field to Auth Keys (fixes #157 and #158)

---
 api.go                           |  3 +++
 cmd/headscale/cli/preauthkeys.go |  3 ++-
 preauth_keys.go                  |  4 ++--
 preauth_keys_test.go             | 13 +++++++++++++
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/api.go b/api.go
index 13955d3..f68852d 100644
--- a/api.go
+++ b/api.go
@@ -395,6 +395,9 @@ func (h *Headscale) handleAuthKey(c *gin.Context, db *gorm.DB, idKey wgkey.Key,
 	m.RegisterMethod = "authKey"
 	db.Save(&m)
 
+	pak.AlreadyUsed = true
+	db.Save(&pak)
+
 	resp.MachineAuthorized = true
 	resp.User = *pak.Namespace.toUser()
 	respBody, err := encode(resp, &idKey, h.privateKey)
diff --git a/cmd/headscale/cli/preauthkeys.go b/cmd/headscale/cli/preauthkeys.go
index 1340267..28bbb7e 100644
--- a/cmd/headscale/cli/preauthkeys.go
+++ b/cmd/headscale/cli/preauthkeys.go
@@ -57,7 +57,7 @@ var listPreAuthKeys = &cobra.Command{
 			return
 		}
 
-		d := pterm.TableData{{"ID", "Key", "Reusable", "Ephemeral", "Expiration", "Created"}}
+		d := pterm.TableData{{"ID", "Key", "Reusable", "Ephemeral", "AlreadyUsed", "Expiration", "Created"}}
 		for _, k := range *keys {
 			expiration := "-"
 			if k.Expiration != nil {
@@ -76,6 +76,7 @@ var listPreAuthKeys = &cobra.Command{
 				k.Key,
 				reusable,
 				strconv.FormatBool(k.Ephemeral),
+				fmt.Sprintf("%v", k.AlreadyUsed),
 				expiration,
 				k.CreatedAt.Format("2006-01-02 15:04:05"),
 			})
diff --git a/preauth_keys.go b/preauth_keys.go
index cc849fc..babc2ff 100644
--- a/preauth_keys.go
+++ b/preauth_keys.go
@@ -21,6 +21,7 @@ type PreAuthKey struct {
 	Namespace   Namespace
 	Reusable    bool
 	Ephemeral   bool `gorm:"default:false"`
+	AlreadyUsed bool `gorm:"default:false"`
 
 	CreatedAt  *time.Time
 	Expiration *time.Time
@@ -110,11 +111,10 @@ func (h *Headscale) checkKeyValidity(k string) (*PreAuthKey, error) {
 		return nil, err
 	}
 
-	if len(machines) != 0 {
+	if len(machines) != 0 || pak.AlreadyUsed {
 		return nil, errorAuthKeyNotReusableAlreadyUsed
 	}
 
-	// missing here validation on current usage
 	return &pak, nil
 }
 
diff --git a/preauth_keys_test.go b/preauth_keys_test.go
index 37f2e4d..80d24e3 100644
--- a/preauth_keys_test.go
+++ b/preauth_keys_test.go
@@ -180,3 +180,16 @@ func (*Suite) TestExpirePreauthKey(c *check.C) {
 	c.Assert(err, check.Equals, errorAuthKeyExpired)
 	c.Assert(p, check.IsNil)
 }
+
+func (*Suite) TestNotReusableMarkedAsAlreadyUsed(c *check.C) {
+	n, err := h.CreateNamespace("test6")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, false, false, nil)
+	c.Assert(err, check.IsNil)
+	pak.AlreadyUsed = true
+	h.db.Save(&pak)
+
+	_, err = h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.Equals, errorAuthKeyNotReusableAlreadyUsed)
+}