Added fields in Machine to store authkey + validation tests

Author: juanfont

machine.go

@@ -25,9 +25,13 @@ type Machine struct {
 	NamespaceID uint
 	Namespace   Namespace
 
-	Registered bool // temp
-	LastSeen   *time.Time
-	Expiry     *time.Time
+	Registered     bool // temp
+	RegisterMethod string
+	AuthKeyID      uint
+	AuthKey        *PreAuthKey
+
+	LastSeen *time.Time
+	Expiry   *time.Time
 
 	HostInfo      postgres.Jsonb
 	Endpoints     postgres.Jsonb

preauth_keys.go

@@ -9,6 +9,7 @@ import (
 
 const errorAuthKeyNotFound = Error("AuthKey not found")
 const errorAuthKeyExpired = Error("AuthKey expired")
+const errorAuthKeyNotReusableAlreadyUsed = Error("AuthKey not reusable already used")
 
 // PreAuthKey describes a pre-authorization key usable in a particular namespace
 type PreAuthKey struct {
@@ -93,6 +94,19 @@ func (h *Headscale) checkKeyValidity(k string) (*PreAuthKey, error) {
 		return nil, errorAuthKeyExpired
 	}
 
+	if pak.Reusable { // we don't need to check if has been used before
+		return &pak, nil
+	}
+
+	machines := []Machine{}
+	if err := db.Preload("AuthKey").Where(&Machine{AuthKeyID: uint(pak.ID)}).Find(&machines).Error; err != nil {
+		return nil, err
+	}
+
+	if len(machines) != 0 {
+		return nil, errorAuthKeyNotReusableAlreadyUsed
+	}
+
 	// missing here validation on current usage
 	return &pak, nil
 }

preauth_keys_test.go

@@ -94,3 +94,87 @@ func (*Suite) TestPreAuthKeyDoesNotExist(c *check.C) {
 	c.Assert(err, check.Equals, errorAuthKeyNotFound)
 	c.Assert(p, check.IsNil)
 }
+
+func (*Suite) TestValidateKeyOk(c *check.C) {
+	n, err := h.CreateNamespace("test3")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, true, nil)
+	c.Assert(err, check.IsNil)
+
+	p, err := h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.IsNil)
+	c.Assert(p.ID, check.Equals, pak.ID)
+}
+
+func (*Suite) TestAlreadyUsedKey(c *check.C) {
+	n, err := h.CreateNamespace("test4")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, false, nil)
+	c.Assert(err, check.IsNil)
+
+	db, err := h.db()
+	if err != nil {
+		c.Fatal(err)
+	}
+	defer db.Close()
+	m := Machine{
+		ID:             0,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Name:           "testest",
+		NamespaceID:    n.ID,
+		Registered:     true,
+		RegisterMethod: "authKey",
+		AuthKeyID:      uint(pak.ID),
+	}
+	db.Save(&m)
+
+	p, err := h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.Equals, errorAuthKeyNotReusableAlreadyUsed)
+	c.Assert(p, check.IsNil)
+}
+
+func (*Suite) TestReusableBeingUsedKey(c *check.C) {
+	n, err := h.CreateNamespace("test5")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, true, nil)
+	c.Assert(err, check.IsNil)
+
+	db, err := h.db()
+	if err != nil {
+		c.Fatal(err)
+	}
+	defer db.Close()
+	m := Machine{
+		ID:             1,
+		MachineKey:     "foo",
+		NodeKey:        "bar",
+		DiscoKey:       "faa",
+		Name:           "testest",
+		NamespaceID:    n.ID,
+		Registered:     true,
+		RegisterMethod: "authKey",
+		AuthKeyID:      uint(pak.ID),
+	}
+	db.Save(&m)
+
+	p, err := h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.IsNil)
+	c.Assert(p.ID, check.Equals, pak.ID)
+}
+
+func (*Suite) TestNotReusableNotBeingUsedKey(c *check.C) {
+	n, err := h.CreateNamespace("test6")
+	c.Assert(err, check.IsNil)
+
+	pak, err := h.CreatePreAuthKey(n.Name, false, nil)
+	c.Assert(err, check.IsNil)
+
+	p, err := h.checkKeyValidity(pak.Key)
+	c.Assert(err, check.IsNil)
+	c.Assert(p.ID, check.Equals, pak.ID)
+}