mirror of
https://github.com/juanfont/headscale.git
synced 2024-12-27 06:35:52 -05:00
83769ba715
This commits removes the locks used to guard data integrity for the database and replaces them with Transactions, turns out that SQL had a way to deal with this all along. This reduces the complexity we had with multiple locks that might stack or recurse (database, nofitifer, mapper). All notifications and state updates are now triggered _after_ a database change. Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
207 lines
5.5 KiB
Go
207 lines
5.5 KiB
Go
package db
|
|
|
|
import (
|
|
"time"
|
|
|
|
"github.com/juanfont/headscale/hscontrol/types"
|
|
"github.com/juanfont/headscale/hscontrol/util"
|
|
"gopkg.in/check.v1"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
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().Add(-5 * time.Second)
|
|
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)
|
|
|
|
node := types.Node{
|
|
ID: 0,
|
|
Hostname: "testest",
|
|
UserID: user.ID,
|
|
RegisterMethod: util.RegisterMethodAuthKey,
|
|
AuthKeyID: uint(pak.ID),
|
|
}
|
|
db.DB.Save(&node)
|
|
|
|
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)
|
|
|
|
node := types.Node{
|
|
ID: 1,
|
|
Hostname: "testest",
|
|
UserID: user.ID,
|
|
RegisterMethod: util.RegisterMethodAuthKey,
|
|
AuthKeyID: uint(pak.ID),
|
|
}
|
|
db.DB.Save(&node)
|
|
|
|
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)
|
|
node := types.Node{
|
|
ID: 0,
|
|
Hostname: "testest",
|
|
UserID: user.ID,
|
|
RegisterMethod: util.RegisterMethodAuthKey,
|
|
LastSeen: &now,
|
|
AuthKeyID: uint(pak.ID),
|
|
}
|
|
db.DB.Save(&node)
|
|
|
|
_, err = db.ValidatePreAuthKey(pak.Key)
|
|
// Ephemeral keys are by definition reusable
|
|
c.Assert(err, check.IsNil)
|
|
|
|
_, err = db.getNode("test7", "testest")
|
|
c.Assert(err, check.IsNil)
|
|
|
|
db.DB.Transaction(func(tx *gorm.DB) error {
|
|
ExpireEphemeralNodes(tx, time.Second*20)
|
|
return nil
|
|
})
|
|
|
|
// The machine record should have been deleted
|
|
_, err = db.getNode("test7", "testest")
|
|
c.Assert(err, check.NotNil)
|
|
}
|
|
|
|
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().GetAclTags(), check.DeepEquals, tags)
|
|
}
|