mirror of
https://github.com/juanfont/headscale.git
synced 2025-10-30 23:51:03 +01:00
58c94d2bd3
This commit restructures the map session in to a struct holding the state of what is needed during its lifetime. For streaming sessions, the event loop is structured a bit differently not hammering the clients with updates but rather batching them over a short, configurable time which should significantly improve cpu usage, and potentially flakyness. The use of Patch updates has been dialed back a little as it does not look like its a 100% ready for prime time. Nodes are now updated with full changes, except for a few things like online status. Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
240 lines
6.3 KiB
Go
240 lines
6.3 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) TestEphemeralKeyReusable(c *check.C) {
|
|
user, err := db.CreateUser("test7")
|
|
c.Assert(err, check.IsNil)
|
|
|
|
pak, err := db.CreatePreAuthKey(user.Name, true, 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)
|
|
c.Assert(err, check.IsNil)
|
|
|
|
_, err = db.getNode("test7", "testest")
|
|
c.Assert(err, check.IsNil)
|
|
|
|
db.DB.Transaction(func(tx *gorm.DB) error {
|
|
DeleteExpiredEphemeralNodes(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) TestEphemeralKeyNotReusable(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)
|
|
c.Assert(err, check.NotNil)
|
|
|
|
_, err = db.getNode("test7", "testest")
|
|
c.Assert(err, check.IsNil)
|
|
|
|
db.DB.Transaction(func(tx *gorm.DB) error {
|
|
DeleteExpiredEphemeralNodes(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)
|
|
}
|