Skip to content

Commit

Permalink
denormalise PreAuthKey tags
Browse files Browse the repository at this point in the history
this commit denormalises the Tags related to a Pre auth key
back onto the preauthkey table and struct as a string list.

There was not really any real normalisation here as we just added
a bunch of duplicate tags with new IDs and preauthkeyIDs, lots of
GORM cermony but no actual advantage.

This work is the start to fixup tags which currently are not working
as they should.

Updates juanfont#1369

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
  • Loading branch information
kradalby committed Sep 27, 2024
1 parent 204a102 commit e49e9d3
Show file tree
Hide file tree
Showing 5 changed files with 127 additions and 38 deletions.
57 changes: 56 additions & 1 deletion hscontrol/db/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package db
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"net/netip"
Expand All @@ -19,6 +20,7 @@ import (
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"tailscale.com/util/set"
)

var errDatabaseNotSupported = errors.New("database type not supported")
Expand Down Expand Up @@ -291,7 +293,12 @@ func NewHeadscaleDatabase(
return err
}

err = tx.AutoMigrate(&types.PreAuthKeyACLTag{})
type preAuthKeyACLTag struct {
ID uint64 `gorm:"primary_key"`
PreAuthKeyID uint64
Tag string
}
err = tx.AutoMigrate(&preAuthKeyACLTag{})
if err != nil {
return err
}
Expand Down Expand Up @@ -413,6 +420,54 @@ func NewHeadscaleDatabase(
},
Rollback: func(db *gorm.DB) error { return nil },
},
// denormalise the ACL tags for preauth keys back onto
// the preauth key table. We dont normalise or reuse and
// it is just a bunch of work for extra work.
{
ID: "202409271400",
Migrate: func(tx *gorm.DB) error {
preauthkeyTags := map[uint64]set.Set[string]{}

type preAuthKeyACLTag struct {
ID uint64 `gorm:"primary_key"`
PreAuthKeyID uint64
Tag string
}

var aclTags []preAuthKeyACLTag
if err := tx.Find(&aclTags).Error; err != nil {
return err
}

// Store the current tags.
for _, tag := range aclTags {
if preauthkeyTags[tag.PreAuthKeyID] == nil {
preauthkeyTags[tag.PreAuthKeyID] = set.SetOf([]string{tag.Tag})
} else {
preauthkeyTags[tag.PreAuthKeyID].Add(tag.Tag)
}
}

// Add tags column and restore the tags.
_ = tx.Migrator().AddColumn(&types.PreAuthKey{}, "tags")
for keyID, tags := range preauthkeyTags {
s := tags.Slice()
j, err := json.Marshal(s)
if err != nil {
return err
}
if err := tx.Model(&types.PreAuthKey{}).Where("id = ?", keyID).Update("tags", string(j)).Error; err != nil {
return err
}
}

// Drop the old table.
_ = tx.Migrator().DropTable(&preAuthKeyACLTag{})

return nil
},
Rollback: func(db *gorm.DB) error { return nil },
},
},
)

Expand Down
64 changes: 64 additions & 0 deletions hscontrol/db/db_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ import (
"net/netip"
"os"
"path/filepath"
"slices"
"sort"
"testing"

"github.com/google/go-cmp/cmp"
Expand Down Expand Up @@ -108,6 +110,68 @@ func TestMigrations(t *testing.T) {
}
},
},
// at 14:15:06 ❯ go run ./cmd/headscale preauthkeys list
// ID | Key | Reusable | Ephemeral | Used | Expiration | Created | Tags
// 1 | 09b28f.. | false | false | false | 2024-09-27 | 2024-09-27 | tag:derp
// 2 | 3112b9.. | false | false | false | 2024-09-27 | 2024-09-27 | tag:derp
// 3 | 7c23b9.. | false | false | false | 2024-09-27 | 2024-09-27 | tag:derp,tag:merp
// 4 | f20155.. | false | false | false | 2024-09-27 | 2024-09-27 | tag:test
// 5 | b212b9.. | false | false | false | 2024-09-27 | 2024-09-27 | tag:test,tag:woop,tag:dedu
{
dbPath: "testdata/0-23-0-to-0-24-0-preauthkey-tags-table.sqlite",
wantFunc: func(t *testing.T, h *HSDatabase) {
keys, err := Read(h.DB, func(rx *gorm.DB) ([]types.PreAuthKey, error) {
kratest, err := ListPreAuthKeys(rx, "kratest")
if err != nil {
return nil, err
}

testkra, err := ListPreAuthKeys(rx, "testkra")
if err != nil {
return nil, err
}

return append(kratest, testkra...), nil
})
assert.NoError(t, err)

assert.Len(t, keys, 5)
want := []types.PreAuthKey{
{
ID: 1,
Tags: []string{"tag:derp"},
},
{
ID: 2,
Tags: []string{"tag:derp"},
},
{
ID: 3,
Tags: []string{"tag:derp", "tag:merp"},
},
{
ID: 4,
Tags: []string{"tag:test"},
},
{
ID: 5,
Tags: []string{"tag:test", "tag:woop", "tag:dedu"},
},
}

if diff := cmp.Diff(want, keys, cmp.Comparer(func(a, b []string) bool {
sort.Sort(sort.StringSlice(a))
sort.Sort(sort.StringSlice(b))
return slices.Equal(a, b)
}), cmpopts.IgnoreFields(types.PreAuthKey{}, "Key", "UserID", "User", "CreatedAt", "Expiration")); diff != "" {
t.Errorf("TestMigrations() mismatch (-want +got):\n%s", diff)
}

if h.DB.Migrator().HasTable("pre_auth_key_acl_tags") {
t.Errorf("TestMigrations() table pre_auth_key_acl_tags should not exist")
}
},
},
}

for _, tt := range tests {
Expand Down
25 changes: 3 additions & 22 deletions hscontrol/db/preauth_keys.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,28 +71,13 @@ func CreatePreAuthKey(
Ephemeral: ephemeral,
CreatedAt: &now,
Expiration: expiration,
Tags: types.StringList(aclTags),
}

if err := tx.Save(&key).Error; err != nil {
return nil, fmt.Errorf("failed to create key in the database: %w", err)
}

if len(aclTags) > 0 {
seenTags := map[string]bool{}

for _, tag := range aclTags {
if !seenTags[tag] {
if err := tx.Save(&types.PreAuthKeyACLTag{PreAuthKeyID: key.ID, Tag: tag}).Error; err != nil {
return nil, fmt.Errorf(
"failed to create key tag in the database: %w",
err,
)
}
seenTags[tag] = true
}
}
}

return &key, nil
}

Expand All @@ -110,7 +95,7 @@ func ListPreAuthKeys(tx *gorm.DB, userName string) ([]types.PreAuthKey, error) {
}

keys := []types.PreAuthKey{}
if err := tx.Preload("User").Preload("ACLTags").Where(&types.PreAuthKey{UserID: user.ID}).Find(&keys).Error; err != nil {
if err := tx.Preload("User").Where(&types.PreAuthKey{UserID: user.ID}).Find(&keys).Error; err != nil {
return nil, err
}

Expand All @@ -135,10 +120,6 @@ func GetPreAuthKey(tx *gorm.DB, user string, key string) (*types.PreAuthKey, err
// does not exist.
func DestroyPreAuthKey(tx *gorm.DB, pak types.PreAuthKey) error {
return tx.Transaction(func(db *gorm.DB) error {
if result := db.Unscoped().Where(types.PreAuthKeyACLTag{PreAuthKeyID: pak.ID}).Delete(&types.PreAuthKeyACLTag{}); result.Error != nil {
return result.Error
}

if result := db.Unscoped().Delete(pak); result.Error != nil {
return result.Error
}
Expand Down Expand Up @@ -182,7 +163,7 @@ func (hsdb *HSDatabase) ValidatePreAuthKey(k string) (*types.PreAuthKey, error)
// If returns no error and a PreAuthKey, it can be used.
func ValidatePreAuthKey(tx *gorm.DB, k string) (*types.PreAuthKey, error) {
pak := types.PreAuthKey{}
if result := tx.Preload("User").Preload("ACLTags").First(&pak, "key = ?", k); errors.Is(
if result := tx.Preload("User").First(&pak, "key = ?", k); errors.Is(
result.Error,
gorm.ErrRecordNotFound,
) {
Expand Down
Binary file not shown.
19 changes: 4 additions & 15 deletions hscontrol/types/preauth_key.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,21 +16,14 @@ type PreAuthKey struct {
UserID uint
User User `gorm:"constraint:OnDelete:CASCADE;"`
Reusable bool
Ephemeral bool `gorm:"default:false"`
Used bool `gorm:"default:false"`
ACLTags []PreAuthKeyACLTag `gorm:"constraint:OnDelete:CASCADE;"`
Ephemeral bool `gorm:"default:false"`
Used bool `gorm:"default:false"`
Tags []string `gorm:"serializer:json"`

CreatedAt *time.Time
Expiration *time.Time
}

// PreAuthKeyACLTag describes an autmatic tag applied to a node when registered with the associated PreAuthKey.
type PreAuthKeyACLTag struct {
ID uint64 `gorm:"primary_key"`
PreAuthKeyID uint64
Tag string
}

func (key *PreAuthKey) Proto() *v1.PreAuthKey {
protoKey := v1.PreAuthKey{
User: key.User.Name,
Expand All @@ -39,7 +32,7 @@ func (key *PreAuthKey) Proto() *v1.PreAuthKey {
Ephemeral: key.Ephemeral,
Reusable: key.Reusable,
Used: key.Used,
AclTags: make([]string, len(key.ACLTags)),
AclTags: key.Tags,
}

if key.Expiration != nil {
Expand All @@ -50,9 +43,5 @@ func (key *PreAuthKey) Proto() *v1.PreAuthKey {
protoKey.CreatedAt = timestamppb.New(*key.CreatedAt)
}

for idx := range key.ACLTags {
protoKey.AclTags[idx] = key.ACLTags[idx].Tag
}

return &protoKey
}

0 comments on commit e49e9d3

Please sign in to comment.