From 486faa9656e622c82e768169dd2d5c6cc063c0a0 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Wed, 5 May 2021 23:00:04 +0200
Subject: [PATCH] WIP Working on authkeys + tests

---
 preauth_keys.go      | 25 +++++++++++++++++++++++++
 preauth_keys_test.go | 21 +++++++++++++++++++++
 utils.go             |  4 ++++
 3 files changed, 50 insertions(+)

diff --git a/preauth_keys.go b/preauth_keys.go
index de89b04d..4f6ec4d7 100644
--- a/preauth_keys.go
+++ b/preauth_keys.go
@@ -7,6 +7,9 @@ import (
 	"time"
 )
 
+const errorAuthKeyNotFound = Error("AuthKey not found")
+const errorAuthKeyExpired = Error("AuthKey expired")
+
 // PreAuthKey describes a pre-authorization key usable in a particular namespace
 type PreAuthKey struct {
 	ID          uint64 `gorm:"primary_key"`
@@ -72,6 +75,28 @@ func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error)
 	return &keys, 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) {
+	db, err := h.db()
+	if err != nil {
+		return nil, err
+	}
+	defer db.Close()
+
+	pak := PreAuthKey{}
+	if db.First(&pak, "key = ?", k).RecordNotFound() {
+		return nil, errorAuthKeyNotFound
+	}
+
+	if pak.Expiration != nil && pak.Expiration.Before(time.Now()) {
+		return nil, errorAuthKeyExpired
+	}
+
+	// missing here validation on current usage
+	return &pak, nil
+}
+
 func (h *Headscale) generateKey() (string, error) {
 	size := 24
 	bytes := make([]byte, size)
diff --git a/preauth_keys_test.go b/preauth_keys_test.go
index 5ac3bcfb..4a98dbeb 100644
--- a/preauth_keys_test.go
+++ b/preauth_keys_test.go
@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"os"
 	"testing"
+	"time"
 
 	_ "github.com/jinzhu/gorm/dialects/sqlite" // sql driver
 
@@ -48,6 +49,7 @@ func (s *Suite) TearDownSuite(c *check.C) {
 
 func (*Suite) TestCreatePreAuthKey(c *check.C) {
 	_, err := h.CreatePreAuthKey("bogus", true, nil)
+
 	c.Assert(err, check.NotNil)
 
 	n, err := h.CreateNamespace("test")
@@ -73,3 +75,22 @@ func (*Suite) TestCreatePreAuthKey(c *check.C) {
 	// Make sure the Namespace association is populated
 	c.Assert((*keys)[0].Namespace.Name, check.Equals, n.Name)
 }
+
+func (*Suite) TestExpiredPreAuthKey(c *check.C) {
+	n, err := h.CreateNamespace("test2")
+	c.Assert(err, check.IsNil)
+
+	now := time.Now()
+	pak, err := h.CreatePreAuthKey(n.Name, true, &now)
+	c.Assert(err, check.IsNil)
+
+	p, err := h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.Equals, errorAuthKeyExpired)
+	c.Assert(p, check.IsNil)
+}
+
+func (*Suite) TestPreAuthKeyDoesNotExist(c *check.C) {
+	p, err := h.checkKeyValidity("potatoKey")
+	c.Assert(err, check.Equals, errorAuthKeyNotFound)
+	c.Assert(p, check.IsNil)
+}
diff --git a/utils.go b/utils.go
index e787b1d7..f85a4ffa 100644
--- a/utils.go
+++ b/utils.go
@@ -21,6 +21,10 @@ import (
 	"tailscale.com/wgengine/wgcfg"
 )
 
+type Error string
+
+func (e Error) Error() string { return string(e) }
+
 func decode(msg []byte, v interface{}, pubKey *wgcfg.Key, privKey *wgcfg.PrivateKey) error {
 	return decodeMsg(msg, v, pubKey, privKey)
 }