From 55b9e3cfdace0a8da7cc358a019d8a9e2dfbfbb7 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Tue, 16 Jun 2026 08:45:11 +0000 Subject: [PATCH] db, sqliteconfig: consolidate query and config helpers --- cmd/headscale/cli/policy.go | 4 +- hscontrol/db/api_key.go | 59 ++----------- hscontrol/db/db.go | 67 ++++----------- hscontrol/db/db_test.go | 2 +- hscontrol/db/node.go | 75 ++++++++--------- hscontrol/db/node_test.go | 2 +- hscontrol/db/preauth_keys.go | 123 ++++++++++++++++------------ hscontrol/db/sqliteconfig/config.go | 55 ++++++------- hscontrol/db/text_serialiser.go | 72 ++++++++-------- hscontrol/db/users.go | 40 +++------ hscontrol/db/users_test.go | 4 +- hscontrol/db/versioncheck.go | 50 +++++------ hscontrol/state/state.go | 4 +- 13 files changed, 232 insertions(+), 325 deletions(-) diff --git a/cmd/headscale/cli/policy.go b/cmd/headscale/cli/policy.go index 039fedd0..09496054 100644 --- a/cmd/headscale/cli/policy.go +++ b/cmd/headscale/cli/policy.go @@ -135,7 +135,7 @@ var setPolicy = &cobra.Command{ } defer d.Close() - users, err := d.ListUsers() + users, err := d.ListUsers(nil) if err != nil { return fmt.Errorf("loading users for policy validation: %w", err) } @@ -194,7 +194,7 @@ var checkPolicy = &cobra.Command{ } defer d.Close() - users, err := d.ListUsers() + users, err := d.ListUsers(nil) if err != nil { return fmt.Errorf("loading users: %w", err) } diff --git a/hscontrol/db/api_key.go b/hscontrol/db/api_key.go index e64648c9..b395c38f 100644 --- a/hscontrol/db/api_key.go +++ b/hscontrol/db/api_key.go @@ -205,61 +205,20 @@ func validateAPIKey(db *gorm.DB, keyStr string) (*types.APIKey, error) { } // New format: parse and verify - const expectedMinLength = apiKeyPrefixLength + 1 + apiKeyHashLength - if len(prefixAndSecret) < expectedMinLength { - return nil, fmt.Errorf( - "%w: key too short, expected at least %d chars after prefix, got %d", - ErrAPIKeyFailedToParse, - expectedMinLength, - len(prefixAndSecret), - ) - } - - // Use fixed-length parsing - prefix := prefixAndSecret[:apiKeyPrefixLength] - - // Validate separator at expected position - if prefixAndSecret[apiKeyPrefixLength] != '-' { - return nil, fmt.Errorf( - "%w: expected separator '-' at position %d, got '%c'", - ErrAPIKeyFailedToParse, - apiKeyPrefixLength, - prefixAndSecret[apiKeyPrefixLength], - ) - } - - secret := prefixAndSecret[apiKeyPrefixLength+1:] - - // Validate secret length - if len(secret) != apiKeyHashLength { - return nil, fmt.Errorf( - "%w: secret length mismatch, expected %d chars, got %d", - ErrAPIKeyFailedToParse, - apiKeyHashLength, - len(secret), - ) - } - - // Validate prefix contains only base64 URL-safe characters - if !isValidBase64URLSafe(prefix) { - return nil, fmt.Errorf( - "%w: prefix contains invalid characters (expected base64 URL-safe: A-Za-z0-9_-)", - ErrAPIKeyFailedToParse, - ) - } - - // Validate secret contains only base64 URL-safe characters - if !isValidBase64URLSafe(secret) { - return nil, fmt.Errorf( - "%w: secret contains invalid characters (expected base64 URL-safe: A-Za-z0-9_-)", - ErrAPIKeyFailedToParse, - ) + prefix, secret, err := parsePrefixedKey( + prefixAndSecret, + apiKeyPrefixLength, + apiKeyHashLength, + ErrAPIKeyFailedToParse, + ) + if err != nil { + return nil, err } // Look up by prefix (indexed) var key types.APIKey - err := db.First(&key, "prefix = ?", prefix).Error + err = db.First(&key, "prefix = ?", prefix).Error if err != nil { return nil, fmt.Errorf("API key not found: %w", err) } diff --git a/hscontrol/db/db.go b/hscontrol/db/db.go index e963373a..5a18f783 100644 --- a/hscontrol/db/db.go +++ b/hscontrol/db/db.go @@ -220,7 +220,7 @@ AND auth_key_id NOT IN ( { ID: "202505141324", Migrate: func(tx *gorm.DB) error { - users, err := ListUsers(tx) + users, err := ListUsers(tx, nil) if err != nil { return fmt.Errorf("listing users: %w", err) } @@ -617,7 +617,7 @@ AND auth_key_id NOT IN ( } // 2. Load users and nodes to create PolicyManager - users, err := ListUsers(tx) + users, err := ListUsers(tx, nil) if err != nil { return fmt.Errorf("loading users for RequestTags migration: %w", err) } @@ -970,59 +970,20 @@ func openDB(cfg types.DatabaseConfig) (*gorm.DB, error) { func runMigrations(cfg types.DatabaseConfig, dbConn *gorm.DB, migrations *gormigrate.Gormigrate) error { if cfg.Type == types.DatabaseSqlite { - // SQLite: Run migrations step-by-step, only disabling foreign keys when necessary - - // List of migration IDs that require foreign keys to be disabled - // These are migrations that perform complex schema changes that GORM cannot handle safely with FK enabled - // NO NEW MIGRATIONS SHOULD BE ADDED HERE. ALL NEW MIGRATIONS MUST RUN WITH FOREIGN KEYS ENABLED. - migrationsRequiringFKDisabled := map[string]bool{ - "202501221827": true, // Route table automigration with FK constraint issues - "202501311657": true, // PreAuthKey table automigration with FK constraint issues - // Add other migration IDs here as they are identified to need FK disabled + // SQLite: Run the early migrations that GORM cannot handle safely with + // foreign keys enabled (route and pre-auth-key automigrations) with FK + // disabled, then run everything else with FK enabled. + // + // NO NEW MIGRATIONS SHOULD RUN WITH FK DISABLED. As of 2025-07-02, all + // new migrations must run with foreign keys enabled via the + // migrations.Migrate() call below. + if err := dbConn.Exec("PRAGMA foreign_keys = OFF").Error; err != nil { //nolint:noinlineerr + return fmt.Errorf("disabling foreign keys: %w", err) } - // Get all migration IDs in order from the actual migration definitions - // Only IDs that are in the migrationsRequiringFKDisabled map will be processed with FK disabled - // any other new migrations are ran after. - migrationIDs := []string{ - // v0.25.0 - "202501221827", - "202501311657", - "202502070949", - - // v0.26.0 - "202502131714", - "202502171819", - "202505091439", - "202505141324", - - // As of 2025-07-02, no new IDs should be added here. - // They will be ran by the migrations.Migrate() call below. - } - - for _, migrationID := range migrationIDs { - log.Trace().Caller().Str("migration_id", migrationID).Msg("running migration") - needsFKDisabled := migrationsRequiringFKDisabled[migrationID] - - if needsFKDisabled { - // Disable foreign keys for this migration - err := dbConn.Exec("PRAGMA foreign_keys = OFF").Error - if err != nil { - return fmt.Errorf("disabling foreign keys for migration %s: %w", migrationID, err) - } - } else { - // Ensure foreign keys are enabled for this migration - err := dbConn.Exec("PRAGMA foreign_keys = ON").Error - if err != nil { - return fmt.Errorf("enabling foreign keys for migration %s: %w", migrationID, err) - } - } - - // Run up to this specific migration (will only run the next pending migration) - err := migrations.MigrateTo(migrationID) - if err != nil { - return fmt.Errorf("running migration %s: %w", migrationID, err) - } + // Run up to and including the last migration that requires FK disabled. + if err := migrations.MigrateTo("202501311657"); err != nil { //nolint:noinlineerr + return fmt.Errorf("running migration 202501311657: %w", err) } if err := dbConn.Exec("PRAGMA foreign_keys = ON").Error; err != nil { //nolint:noinlineerr diff --git a/hscontrol/db/db_test.go b/hscontrol/db/db_test.go index 60abcfa0..303c4321 100644 --- a/hscontrol/db/db_test.go +++ b/hscontrol/db/db_test.go @@ -36,7 +36,7 @@ func TestSQLiteMigrationAndDataValidation(t *testing.T) { // Verify users data preservation users, err := Read(hsdb.DB, func(rx *gorm.DB) ([]types.User, error) { - return ListUsers(rx) + return ListUsers(rx, nil) }) require.NoError(t, err) assert.Len(t, users, 1, "should preserve all 1 user from original schema") diff --git a/hscontrol/db/node.go b/hscontrol/db/node.go index 192f05ad..226f7b1d 100644 --- a/hscontrol/db/node.go +++ b/hscontrol/db/node.go @@ -30,6 +30,15 @@ const ( // ErrNodeNameNotUnique is returned when a node name is not unique. var ErrNodeNameNotUnique = errors.New("node name is not unique") +// preloadNode returns a session that eager-loads a node's AuthKey, the +// AuthKey's User, and the node's User. +func preloadNode(tx *gorm.DB) *gorm.DB { + return tx. + Preload("AuthKey"). + Preload("AuthKey.User"). + Preload("User") +} + var ( ErrNodeNotFound = errors.New("node not found") ErrNodeRouteIsNotAvailable = errors.New("route is not available on node") @@ -52,10 +61,7 @@ func (hsdb *HSDatabase) ListPeers(nodeID types.NodeID, peerIDs ...types.NodeID) func ListPeers(tx *gorm.DB, nodeID types.NodeID, peerIDs ...types.NodeID) (types.Nodes, error) { nodes := types.Nodes{} - err := tx. - Preload("AuthKey"). - Preload("AuthKey.User"). - Preload("User"). + err := preloadNode(tx). Where("id <> ?", nodeID). Where(peerIDs).Find(&nodes).Error if err != nil { @@ -78,10 +84,7 @@ func (hsdb *HSDatabase) ListNodes(nodeIDs ...types.NodeID) (types.Nodes, error) func ListNodes(tx *gorm.DB, nodeIDs ...types.NodeID) (types.Nodes, error) { nodes := types.Nodes{} - err := tx. - Preload("AuthKey"). - Preload("AuthKey.User"). - Preload("User"). + err := preloadNode(tx). Where(nodeIDs).Find(&nodes).Error if err != nil { return nil, err @@ -111,18 +114,22 @@ func (hsdb *HSDatabase) getNode(uid types.UserID, name string) (*types.Node, err // getNode finds a [types.Node] by name and user and returns the [types.Node] struct. func getNode(tx *gorm.DB, uid types.UserID, name string) (*types.Node, error) { - nodes, err := ListNodesByUser(tx, uid) + uidPtr := uint(uid) + + node := types.Node{} + + err := preloadNode(tx). + Where(&types.Node{UserID: &uidPtr, Hostname: name}). + First(&node).Error if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrNodeNotFound + } + return nil, err } - for _, m := range nodes { - if m.Hostname == name { - return m, nil - } - } - - return nil, ErrNodeNotFound + return &node, nil } func (hsdb *HSDatabase) GetNodeByID(id types.NodeID) (*types.Node, error) { @@ -132,11 +139,8 @@ func (hsdb *HSDatabase) GetNodeByID(id types.NodeID) (*types.Node, error) { // GetNodeByID finds a [types.Node] by ID and returns the [types.Node] struct. func GetNodeByID(tx *gorm.DB, id types.NodeID) (*types.Node, error) { mach := types.Node{} - if result := tx. - Preload("AuthKey"). - Preload("AuthKey.User"). - Preload("User"). - Find(&types.Node{ID: id}).First(&mach); result.Error != nil { + if result := preloadNode(tx). + First(&mach, "id = ?", id); result.Error != nil { return nil, result.Error } @@ -153,10 +157,7 @@ func GetNodeByNodeKey( nodeKey key.NodePublic, ) (*types.Node, error) { mach := types.Node{} - if result := tx. - Preload("AuthKey"). - Preload("AuthKey.User"). - Preload("User"). + if result := preloadNode(tx). First(&mach, "node_key = ?", nodeKey.String()); result.Error != nil { return nil, result.Error } @@ -522,6 +523,15 @@ func (e *EphemeralGarbageCollector) Start() { } } +// firstOr returns the first non-empty option, or def if none is provided. +func firstOr(def string, opt []string) string { + if len(opt) > 0 && opt[0] != "" { + return opt[0] + } + + return def +} + func (hsdb *HSDatabase) CreateNodeForTest(user *types.User, hostname ...string) *types.Node { if !testing.Testing() { panic("CreateNodeForTest can only be called during tests") @@ -531,10 +541,7 @@ func (hsdb *HSDatabase) CreateNodeForTest(user *types.User, hostname ...string) panic("CreateNodeForTest requires a valid user") } - nodeName := defaultTestNodePrefix - if len(hostname) > 0 && hostname[0] != "" { - nodeName = hostname[0] - } + nodeName := firstOr(defaultTestNodePrefix, hostname) // Create a preauth key for the node pak, err := hsdb.CreatePreAuthKey(user.TypedID(), false, false, nil, nil) @@ -604,10 +611,7 @@ func (hsdb *HSDatabase) CreateNodesForTest(user *types.User, count int, hostname panic("CreateNodesForTest requires a valid user") } - prefix := defaultTestNodePrefix - if len(hostnamePrefix) > 0 && hostnamePrefix[0] != "" { - prefix = hostnamePrefix[0] - } + prefix := firstOr(defaultTestNodePrefix, hostnamePrefix) nodes := make([]*types.Node, count) for i := range count { @@ -627,10 +631,7 @@ func (hsdb *HSDatabase) CreateRegisteredNodesForTest(user *types.User, count int panic("CreateRegisteredNodesForTest requires a valid user") } - prefix := defaultTestNodePrefix - if len(hostnamePrefix) > 0 && hostnamePrefix[0] != "" { - prefix = hostnamePrefix[0] - } + prefix := firstOr(defaultTestNodePrefix, hostnamePrefix) nodes := make([]*types.Node, count) for i := range count { diff --git a/hscontrol/db/node_test.go b/hscontrol/db/node_test.go index 7bbcf3d9..b628bf71 100644 --- a/hscontrol/db/node_test.go +++ b/hscontrol/db/node_test.go @@ -346,7 +346,7 @@ func TestAutoApproveRoutes(t *testing.T) { err = adb.DB.Save(&nodeTagged).Error require.NoError(t, err) - users, err := adb.ListUsers() + users, err := adb.ListUsers(nil) require.NoError(t, err) nodes, err := adb.ListNodes() diff --git a/hscontrol/db/preauth_keys.go b/hscontrol/db/preauth_keys.go index c8369f79..fbdcd8c7 100644 --- a/hscontrol/db/preauth_keys.go +++ b/hscontrol/db/preauth_keys.go @@ -211,60 +211,18 @@ func findAuthKey(tx *gorm.DB, keyStr string) (*types.PreAuthKey, error) { } // New format: hskey-auth-{12-char-prefix}-{64-char-hash} - // Expected minimum length: 12 (prefix) + 1 (separator) + 64 (hash) = 77 - const expectedMinLength = authKeyPrefixLength + 1 + authKeyLength - if len(prefixAndHash) < expectedMinLength { - return nil, fmt.Errorf( - "%w: key too short, expected at least %d chars after prefix, got %d", - ErrPreAuthKeyFailedToParse, - expectedMinLength, - len(prefixAndHash), - ) - } - - // Use fixed-length parsing instead of separator-based to handle dashes in base64 URL-safe - prefix := prefixAndHash[:authKeyPrefixLength] - - // Validate separator at expected position - if prefixAndHash[authKeyPrefixLength] != '-' { - return nil, fmt.Errorf( - "%w: expected separator '-' at position %d, got '%c'", - ErrPreAuthKeyFailedToParse, - authKeyPrefixLength, - prefixAndHash[authKeyPrefixLength], - ) - } - - hash := prefixAndHash[authKeyPrefixLength+1:] - - // Validate hash length - if len(hash) != authKeyLength { - return nil, fmt.Errorf( - "%w: hash length mismatch, expected %d chars, got %d", - ErrPreAuthKeyFailedToParse, - authKeyLength, - len(hash), - ) - } - - // Validate prefix contains only base64 URL-safe characters - if !isValidBase64URLSafe(prefix) { - return nil, fmt.Errorf( - "%w: prefix contains invalid characters (expected base64 URL-safe: A-Za-z0-9_-)", - ErrPreAuthKeyFailedToParse, - ) - } - - // Validate hash contains only base64 URL-safe characters - if !isValidBase64URLSafe(hash) { - return nil, fmt.Errorf( - "%w: hash contains invalid characters (expected base64 URL-safe: A-Za-z0-9_-)", - ErrPreAuthKeyFailedToParse, - ) + prefix, hash, err := parsePrefixedKey( + prefixAndHash, + authKeyPrefixLength, + authKeyLength, + ErrPreAuthKeyFailedToParse, + ) + if err != nil { + return nil, err } // Look up key by prefix - err := tx.Preload("User").First(&pak, "prefix = ?", prefix).Error + err = tx.Preload("User").First(&pak, "prefix = ?", prefix).Error if err != nil { return nil, ErrPreAuthKeyNotFound } @@ -278,6 +236,69 @@ func findAuthKey(tx *gorm.DB, keyStr string) (*types.PreAuthKey, error) { return &pak, nil } +// parsePrefixedKey splits the prefix-and-secret portion of a new-format key +// (the part after the "hskey-*-" prefix) into its fixed-length prefix and +// secret components, validating the length, separator position, and that both +// components are base64 URL-safe. Fixed-length parsing is used instead of +// separator-based to handle dashes in base64 URL-safe characters. +func parsePrefixedKey( + prefixAndSecret string, + prefixLen, secretLen int, + parseErr error, +) (string, string, error) { + expectedMinLength := prefixLen + 1 + secretLen + if len(prefixAndSecret) < expectedMinLength { + return "", "", fmt.Errorf( + "%w: key too short, expected at least %d chars after prefix, got %d", + parseErr, + expectedMinLength, + len(prefixAndSecret), + ) + } + + prefix := prefixAndSecret[:prefixLen] + + // Validate separator at expected position + if prefixAndSecret[prefixLen] != '-' { + return "", "", fmt.Errorf( + "%w: expected separator '-' at position %d, got '%c'", + parseErr, + prefixLen, + prefixAndSecret[prefixLen], + ) + } + + secret := prefixAndSecret[prefixLen+1:] + + // Validate secret length + if len(secret) != secretLen { + return "", "", fmt.Errorf( + "%w: secret length mismatch, expected %d chars, got %d", + parseErr, + secretLen, + len(secret), + ) + } + + // Validate prefix contains only base64 URL-safe characters + if !isValidBase64URLSafe(prefix) { + return "", "", fmt.Errorf( + "%w: prefix contains invalid characters (expected base64 URL-safe: A-Za-z0-9_-)", + parseErr, + ) + } + + // Validate secret contains only base64 URL-safe characters + if !isValidBase64URLSafe(secret) { + return "", "", fmt.Errorf( + "%w: secret contains invalid characters (expected base64 URL-safe: A-Za-z0-9_-)", + parseErr, + ) + } + + return prefix, secret, nil +} + // isValidBase64URLSafe checks if a string contains only base64 URL-safe characters. func isValidBase64URLSafe(s string) bool { for _, c := range s { diff --git a/hscontrol/db/sqliteconfig/config.go b/hscontrol/db/sqliteconfig/config.go index 711e434a..66fec590 100644 --- a/hscontrol/db/sqliteconfig/config.go +++ b/hscontrol/db/sqliteconfig/config.go @@ -347,33 +347,6 @@ func (c *Config) ToURL() (string, error) { return "", fmt.Errorf("invalid config: %w", err) } - var pragmas []string - - // Add pragma parameters only if they're set (non-zero/non-empty) - if c.BusyTimeout > 0 { - pragmas = append(pragmas, fmt.Sprintf("busy_timeout=%d", c.BusyTimeout)) - } - - if c.JournalMode != "" { - pragmas = append(pragmas, fmt.Sprintf("journal_mode=%s", c.JournalMode)) - } - - if c.AutoVacuum != "" { - pragmas = append(pragmas, fmt.Sprintf("auto_vacuum=%s", c.AutoVacuum)) - } - - if c.WALAutocheckpoint >= 0 { - pragmas = append(pragmas, fmt.Sprintf("wal_autocheckpoint=%d", c.WALAutocheckpoint)) - } - - if c.Synchronous != "" { - pragmas = append(pragmas, fmt.Sprintf("synchronous=%s", c.Synchronous)) - } - - if c.ForeignKeys { - pragmas = append(pragmas, "foreign_keys=ON") - } - // Handle different database types var baseURL string if c.Path == ":memory:" { @@ -383,16 +356,36 @@ func (c *Config) ToURL() (string, error) { } // Build query parameters - queryParts := make([]string, 0, 1+len(pragmas)) + var queryParts []string // Add _txlock first (it's a connection parameter, not a pragma) if c.TxLock != "" { queryParts = append(queryParts, "_txlock="+string(c.TxLock)) } - // Add pragma parameters - for _, pragma := range pragmas { - queryParts = append(queryParts, "_pragma="+pragma) + // Add pragma parameters only if they're set (non-zero/non-empty) + if c.BusyTimeout > 0 { + queryParts = append(queryParts, fmt.Sprintf("_pragma=busy_timeout=%d", c.BusyTimeout)) + } + + if c.JournalMode != "" { + queryParts = append(queryParts, fmt.Sprintf("_pragma=journal_mode=%s", c.JournalMode)) + } + + if c.AutoVacuum != "" { + queryParts = append(queryParts, fmt.Sprintf("_pragma=auto_vacuum=%s", c.AutoVacuum)) + } + + if c.WALAutocheckpoint >= 0 { + queryParts = append(queryParts, fmt.Sprintf("_pragma=wal_autocheckpoint=%d", c.WALAutocheckpoint)) + } + + if c.Synchronous != "" { + queryParts = append(queryParts, fmt.Sprintf("_pragma=synchronous=%s", c.Synchronous)) + } + + if c.ForeignKeys { + queryParts = append(queryParts, "_pragma=foreign_keys=ON") } if len(queryParts) > 0 { diff --git a/hscontrol/db/text_serialiser.go b/hscontrol/db/text_serialiser.go index 7c39b8be..c315772b 100644 --- a/hscontrol/db/text_serialiser.go +++ b/hscontrol/db/text_serialiser.go @@ -47,45 +47,45 @@ func (TextSerialiser) Scan(ctx context.Context, field *schema.Field, dst reflect fieldValue = fieldValue.Elem() } - if dbValue != nil { - var bytes []byte + if dbValue == nil { + return nil + } - switch v := dbValue.(type) { - case []byte: - bytes = v - case string: - bytes = []byte(v) - default: - return fmt.Errorf("%w: %#v", errUnmarshalTextValue, dbValue) + var bytes []byte + + switch v := dbValue.(type) { + case []byte: + bytes = v + case string: + bytes = []byte(v) + default: + return fmt.Errorf("%w: %#v", errUnmarshalTextValue, dbValue) + } + + if !isTextUnmarshaler(fieldValue) { + return fmt.Errorf("%w: %T", errUnsupportedType, fieldValue.Interface()) + } + + maybeInstantiatePtr(fieldValue) + f := fieldValue.MethodByName("UnmarshalText") + args := []reflect.Value{reflect.ValueOf(bytes)} + + ret := f.Call(args) + if !ret[0].IsNil() { + if err, ok := ret[0].Interface().(error); ok { + return decodingError(field.Name, err) } + } - if isTextUnmarshaler(fieldValue) { - maybeInstantiatePtr(fieldValue) - f := fieldValue.MethodByName("UnmarshalText") - args := []reflect.Value{reflect.ValueOf(bytes)} - - ret := f.Call(args) - if !ret[0].IsNil() { - if err, ok := ret[0].Interface().(error); ok { - return decodingError(field.Name, err) - } - } - - // If the underlying field is to a pointer type, we need to - // assign the value as a pointer to it. - // If it is not a pointer, we need to assign the value to the - // field. - dstField := field.ReflectValueOf(ctx, dst) - if dstField.Kind() == reflect.Pointer { - dstField.Set(fieldValue) - } else { - dstField.Set(fieldValue.Elem()) - } - - return nil - } else { - return fmt.Errorf("%w: %T", errUnsupportedType, fieldValue.Interface()) - } + // If the underlying field is to a pointer type, we need to + // assign the value as a pointer to it. + // If it is not a pointer, we need to assign the value to the + // field. + dstField := field.ReflectValueOf(ctx, dst) + if dstField.Kind() == reflect.Pointer { + dstField.Set(fieldValue) + } else { + dstField.Set(fieldValue.Elem()) } return nil diff --git a/hscontrol/db/users.go b/hscontrol/db/users.go index 271c3e3d..86ee936f 100644 --- a/hscontrol/db/users.go +++ b/hscontrol/db/users.go @@ -12,11 +12,10 @@ import ( ) var ( - ErrUserExists = errors.New("user already exists") - ErrUserNotFound = errors.New("user not found") - ErrUserStillHasNodes = errors.New("user not empty: node(s) found") - ErrUserWhereInvalidCount = errors.New("expect 0 or 1 where User structs") - ErrUserNotUnique = errors.New("expected exactly one user") + ErrUserExists = errors.New("user already exists") + ErrUserNotFound = errors.New("user not found") + ErrUserStillHasNodes = errors.New("user not empty: node(s) found") + ErrUserNotUnique = errors.New("expected exactly one user") ) func (hsdb *HSDatabase) CreateUser(user types.User) (*types.User, error) { @@ -152,24 +151,15 @@ func GetUserByOIDCIdentifier(tx *gorm.DB, id string) (*types.User, error) { return &user, nil } -func (hsdb *HSDatabase) ListUsers(where ...*types.User) ([]types.User, error) { - return ListUsers(hsdb.DB, where...) +func (hsdb *HSDatabase) ListUsers(filter *types.User) ([]types.User, error) { + return ListUsers(hsdb.DB, filter) } -// ListUsers gets all the existing users. -func ListUsers(tx *gorm.DB, where ...*types.User) ([]types.User, error) { - if len(where) > 1 { - return nil, fmt.Errorf("%w, got %d", ErrUserWhereInvalidCount, len(where)) - } - - var user *types.User - if len(where) == 1 { - user = where[0] - } - +// ListUsers gets all the existing users, optionally filtered by a non-nil filter. +func ListUsers(tx *gorm.DB, filter *types.User) ([]types.User, error) { users := []types.User{} - err := tx.Where(user).Find(&users).Error + err := tx.Where(filter).Find(&users).Error if err != nil { return nil, err } @@ -202,7 +192,7 @@ func ListNodesByUser(tx *gorm.DB, uid types.UserID) (types.Nodes, error) { uidPtr := uint(uid) - err := tx.Preload("AuthKey").Preload("AuthKey.User").Preload("User").Where(&types.Node{UserID: &uidPtr}).Find(&nodes).Error + err := preloadNode(tx).Where(&types.Node{UserID: &uidPtr}).Find(&nodes).Error if err != nil { return nil, err } @@ -215,10 +205,7 @@ func (hsdb *HSDatabase) CreateUserForTest(name ...string) *types.User { panic("CreateUserForTest can only be called during tests") } - userName := "testuser" - if len(name) > 0 && name[0] != "" { - userName = name[0] - } + userName := firstOr("testuser", name) user, err := hsdb.CreateUser(types.User{Name: userName}) if err != nil { @@ -233,10 +220,7 @@ func (hsdb *HSDatabase) CreateUsersForTest(count int, namePrefix ...string) []*t panic("CreateUsersForTest can only be called during tests") } - prefix := "testuser" - if len(namePrefix) > 0 && namePrefix[0] != "" { - prefix = namePrefix[0] - } + prefix := firstOr("testuser", namePrefix) users := make([]*types.User, count) for i := range count { diff --git a/hscontrol/db/users_test.go b/hscontrol/db/users_test.go index 36d836f9..2a755ff3 100644 --- a/hscontrol/db/users_test.go +++ b/hscontrol/db/users_test.go @@ -17,7 +17,7 @@ func TestCreateAndDestroyUser(t *testing.T) { user := db.CreateUserForTest("test") assert.Equal(t, "test", user.Name) - users, err := db.ListUsers() + users, err := db.ListUsers(nil) require.NoError(t, err) assert.Len(t, users, 1) @@ -227,7 +227,7 @@ func TestRenameUser(t *testing.T) { userTest := db.CreateUserForTest("test") assert.Equal(t, "test", userTest.Name) - users, err := db.ListUsers() + users, err := db.ListUsers(nil) require.NoError(t, err) assert.Len(t, users, 1) diff --git a/hscontrol/db/versioncheck.go b/hscontrol/db/versioncheck.go index a071348c..199b5964 100644 --- a/hscontrol/db/versioncheck.go +++ b/hscontrol/db/versioncheck.go @@ -11,6 +11,7 @@ import ( "github.com/juanfont/headscale/hscontrol/types" "github.com/rs/zerolog/log" "gorm.io/gorm" + "gorm.io/gorm/clause" ) var errVersionUpgrade = errors.New("version upgrade not supported") @@ -66,22 +67,20 @@ func parseVersion(s string) (semver, error) { return semver{}, fmt.Errorf("%q: %w", s, errVersionFormat) } - major, err := strconv.Atoi(parts[0]) - if err != nil { - return semver{}, fmt.Errorf("invalid major version in %q: %w", s, err) + var out [3]int + + names := [...]string{"major", "minor", "patch"} + + for i, p := range parts { + n, err := strconv.Atoi(p) + if err != nil { + return semver{}, fmt.Errorf("invalid %s version in %q: %w", names[i], s, err) + } + + out[i] = n } - minor, err := strconv.Atoi(parts[1]) - if err != nil { - return semver{}, fmt.Errorf("invalid minor version in %q: %w", s, err) - } - - patch, err := strconv.Atoi(parts[2]) - if err != nil { - return semver{}, fmt.Errorf("invalid patch version in %q: %w", s, err) - } - - return semver{Major: major, Minor: minor, Patch: patch}, nil + return semver{Major: out[0], Minor: out[1], Patch: out[2]}, nil } // ensureDatabaseVersionTable creates the database_versions table if it @@ -118,23 +117,12 @@ func getDatabaseVersion(db *gorm.DB) (string, error) { func setDatabaseVersion(db *gorm.DB, version string) error { now := time.Now().UTC() - // Try update first, then insert if no rows affected. - result := db.Exec( - "UPDATE database_versions SET version = ?, updated_at = ? WHERE id = 1", - version, now, - ) - if result.Error != nil { - return fmt.Errorf("updating database version: %w", result.Error) - } - - if result.RowsAffected == 0 { - err := db.Exec( - "INSERT INTO database_versions (id, version, updated_at) VALUES (1, ?, ?)", - version, now, - ).Error - if err != nil { - return fmt.Errorf("inserting database version: %w", err) - } + err := db.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.AssignmentColumns([]string{"version", "updated_at"}), + }).Create(&DatabaseVersion{ID: 1, Version: version, UpdatedAt: now}).Error + if err != nil { + return fmt.Errorf("upserting database version: %w", err) } return nil diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index edbed2e5..a56ffff1 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -247,7 +247,7 @@ func NewState(cfg *types.Config) (*State, error) { node.IsOnline = new(false) } - users, err := db.ListUsers() + users, err := db.ListUsers(nil) if err != nil { return nil, fmt.Errorf("loading users: %w", err) } @@ -508,7 +508,7 @@ func (s *State) ListUsersWithFilter(filter *types.User) ([]types.User, error) { // ListAllUsers retrieves all users in the system. func (s *State) ListAllUsers() ([]types.User, error) { - return s.db.ListUsers() + return s.db.ListUsers(nil) } // persistNodeRowToDB writes the node's database row, re-reading the