From 226cb89d97da6ab77ec4afecda18c4e7b5b5bd3a Mon Sep 17 00:00:00 2001
From: Juan Font <juanfontalonso@gmail.com>
Date: Sat, 7 Aug 2021 23:57:52 +0200
Subject: [PATCH] Added func to expire PAKs

---
 preauth_keys.go      |  7 +++++++
 preauth_keys_test.go | 17 +++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/preauth_keys.go b/preauth_keys.go
index 7cffceae..460797c4 100644
--- a/preauth_keys.go
+++ b/preauth_keys.go
@@ -67,6 +67,13 @@ func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error)
 	return &keys, nil
 }
 
+func (h *Headscale) MarkExpirePreAuthKey(k *PreAuthKey) error {
+	if err := h.db.Model(&k).Update("Expiration", time.Now()).Error; err != nil {
+		return err
+	}
+	return nil
+}
+
 // checkKeyValidity does the heavy lifting for validation of the PreAuthKey coming from a node
 // If returns no error and a PreAuthKey, it can be used
 func (h *Headscale) checkKeyValidity(k string) (*PreAuthKey, error) {
diff --git a/preauth_keys_test.go b/preauth_keys_test.go
index 6f1369c5..37f2e4dd 100644
--- a/preauth_keys_test.go
+++ b/preauth_keys_test.go
@@ -163,3 +163,20 @@ func (*Suite) TestEphemeralKey(c *check.C) {
 	_, err = h.GetMachine("test7", "testest")
 	c.Assert(err, check.NotNil)
 }
+
+func (*Suite) TestExpirePreauthKey(c *check.C) {
+	n, err := h.CreateNamespace("test3")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, true, false, nil)
+	c.Assert(err, check.IsNil)
+	c.Assert(pak.Expiration, check.IsNil)
+
+	err = h.MarkExpirePreAuthKey(pak)
+	c.Assert(err, check.IsNil)
+	c.Assert(pak.Expiration, check.NotNil)
+
+	p, err := h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.Equals, errorAuthKeyExpired)
+	c.Assert(p, check.IsNil)
+}