mirror of
https://github.com/juanfont/headscale.git
synced 2025-11-09 21:49:39 -05:00
Split code into modules
This is a massive commit that restructures the code into modules:
db/
All functions related to modifying the Database
types/
All type definitions and methods that can be exclusivly used on
these types without dependencies
policy/
All Policy related code, now without dependencies on the Database.
policy/matcher/
Dedicated code to match machines in a list of FilterRules
Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
This commit is contained in:
Kristoffer Dalby
committed by
Kristoffer Dalby
Kristoffer Dalby
parent
14e29a7bee
commit
feb15365b5
213
hscontrol/db/preauth_keys_test.go
Normal file
213
hscontrol/db/preauth_keys_test.go
Normal file
@@ -0,0 +1,213 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/juanfont/headscale/hscontrol/types"
|
||||
"github.com/juanfont/headscale/hscontrol/util"
|
||||
"gopkg.in/check.v1"
|
||||
)
|
||||
|
||||
func (*Suite) TestCreatePreAuthKey(c *check.C) {
|
||||
_, err := db.CreatePreAuthKey("bogus", true, false, nil, nil)
|
||||
|
||||
c.Assert(err, check.NotNil)
|
||||
|
||||
user, err := db.CreateUser("test")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
key, err := db.CreatePreAuthKey(user.Name, true, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
// Did we get a valid key?
|
||||
c.Assert(key.Key, check.NotNil)
|
||||
c.Assert(len(key.Key), check.Equals, 48)
|
||||
|
||||
// Make sure the User association is populated
|
||||
c.Assert(key.User.Name, check.Equals, user.Name)
|
||||
|
||||
_, err = db.ListPreAuthKeys("bogus")
|
||||
c.Assert(err, check.NotNil)
|
||||
|
||||
keys, err := db.ListPreAuthKeys(user.Name)
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(len(keys), check.Equals, 1)
|
||||
|
||||
// Make sure the User association is populated
|
||||
c.Assert((keys)[0].User.Name, check.Equals, user.Name)
|
||||
}
|
||||
|
||||
func (*Suite) TestExpiredPreAuthKey(c *check.C) {
|
||||
user, err := db.CreateUser("test2")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
now := time.Now()
|
||||
pak, err := db.CreatePreAuthKey(user.Name, true, false, &now, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
key, err := db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.Equals, ErrPreAuthKeyExpired)
|
||||
c.Assert(key, check.IsNil)
|
||||
}
|
||||
|
||||
func (*Suite) TestPreAuthKeyDoesNotExist(c *check.C) {
|
||||
key, err := db.ValidatePreAuthKey("potatoKey")
|
||||
c.Assert(err, check.Equals, ErrPreAuthKeyNotFound)
|
||||
c.Assert(key, check.IsNil)
|
||||
}
|
||||
|
||||
func (*Suite) TestValidateKeyOk(c *check.C) {
|
||||
user, err := db.CreateUser("test3")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, true, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
key, err := db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(key.ID, check.Equals, pak.ID)
|
||||
}
|
||||
|
||||
func (*Suite) TestAlreadyUsedKey(c *check.C) {
|
||||
user, err := db.CreateUser("test4")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, false, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
machine := types.Machine{
|
||||
ID: 0,
|
||||
MachineKey: "foo",
|
||||
NodeKey: "bar",
|
||||
DiscoKey: "faa",
|
||||
Hostname: "testest",
|
||||
UserID: user.ID,
|
||||
RegisterMethod: util.RegisterMethodAuthKey,
|
||||
AuthKeyID: uint(pak.ID),
|
||||
}
|
||||
db.db.Save(&machine)
|
||||
|
||||
key, err := db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.Equals, ErrSingleUseAuthKeyHasBeenUsed)
|
||||
c.Assert(key, check.IsNil)
|
||||
}
|
||||
|
||||
func (*Suite) TestReusableBeingUsedKey(c *check.C) {
|
||||
user, err := db.CreateUser("test5")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, true, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
machine := types.Machine{
|
||||
ID: 1,
|
||||
MachineKey: "foo",
|
||||
NodeKey: "bar",
|
||||
DiscoKey: "faa",
|
||||
Hostname: "testest",
|
||||
UserID: user.ID,
|
||||
RegisterMethod: util.RegisterMethodAuthKey,
|
||||
AuthKeyID: uint(pak.ID),
|
||||
}
|
||||
db.db.Save(&machine)
|
||||
|
||||
key, err := db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(key.ID, check.Equals, pak.ID)
|
||||
}
|
||||
|
||||
func (*Suite) TestNotReusableNotBeingUsedKey(c *check.C) {
|
||||
user, err := db.CreateUser("test6")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, false, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
key, err := db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(key.ID, check.Equals, pak.ID)
|
||||
}
|
||||
|
||||
func (*Suite) TestEphemeralKey(c *check.C) {
|
||||
user, err := db.CreateUser("test7")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, false, true, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
now := time.Now().Add(-time.Second * 30)
|
||||
machine := types.Machine{
|
||||
ID: 0,
|
||||
MachineKey: "foo",
|
||||
NodeKey: "bar",
|
||||
DiscoKey: "faa",
|
||||
Hostname: "testest",
|
||||
UserID: user.ID,
|
||||
RegisterMethod: util.RegisterMethodAuthKey,
|
||||
LastSeen: &now,
|
||||
AuthKeyID: uint(pak.ID),
|
||||
}
|
||||
db.db.Save(&machine)
|
||||
|
||||
_, err = db.ValidatePreAuthKey(pak.Key)
|
||||
// Ephemeral keys are by definition reusable
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
_, err = db.GetMachine("test7", "testest")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
db.ExpireEphemeralMachines(time.Second * 20)
|
||||
|
||||
// The machine record should have been deleted
|
||||
_, err = db.GetMachine("test7", "testest")
|
||||
c.Assert(err, check.NotNil)
|
||||
|
||||
c.Assert(channelUpdates, check.Equals, int32(1))
|
||||
}
|
||||
|
||||
func (*Suite) TestExpirePreauthKey(c *check.C) {
|
||||
user, err := db.CreateUser("test3")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, true, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(pak.Expiration, check.IsNil)
|
||||
|
||||
err = db.ExpirePreAuthKey(pak)
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(pak.Expiration, check.NotNil)
|
||||
|
||||
key, err := db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.Equals, ErrPreAuthKeyExpired)
|
||||
c.Assert(key, check.IsNil)
|
||||
}
|
||||
|
||||
func (*Suite) TestNotReusableMarkedAsUsed(c *check.C) {
|
||||
user, err := db.CreateUser("test6")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
pak, err := db.CreatePreAuthKey(user.Name, false, false, nil, nil)
|
||||
c.Assert(err, check.IsNil)
|
||||
pak.Used = true
|
||||
db.db.Save(&pak)
|
||||
|
||||
_, err = db.ValidatePreAuthKey(pak.Key)
|
||||
c.Assert(err, check.Equals, ErrSingleUseAuthKeyHasBeenUsed)
|
||||
}
|
||||
|
||||
func (*Suite) TestPreAuthKeyACLTags(c *check.C) {
|
||||
user, err := db.CreateUser("test8")
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
_, err = db.CreatePreAuthKey(user.Name, false, false, nil, []string{"badtag"})
|
||||
c.Assert(err, check.NotNil) // Confirm that malformed tags are rejected
|
||||
|
||||
tags := []string{"tag:test1", "tag:test2"}
|
||||
tagsWithDuplicate := []string{"tag:test1", "tag:test2", "tag:test2"}
|
||||
_, err = db.CreatePreAuthKey(user.Name, false, false, nil, tagsWithDuplicate)
|
||||
c.Assert(err, check.IsNil)
|
||||
|
||||
listedPaks, err := db.ListPreAuthKeys("test8")
|
||||
c.Assert(err, check.IsNil)
|
||||
c.Assert(listedPaks[0].Proto().AclTags, check.DeepEquals, tags)
|
||||
}
|
||||
Reference in New Issue
Block a user