juanfont.headscale/hscontrol/db/users.go

package db

import (
	"errors"

	"github.com/juanfont/headscale/hscontrol/types"
	"github.com/juanfont/headscale/hscontrol/util"
	"github.com/rs/zerolog/log"
	"gorm.io/gorm"
)

var (
	ErrUserExists        = errors.New("user already exists")
	ErrUserNotFound      = errors.New("user not found")
	ErrUserStillHasNodes = errors.New("user not empty: node(s) found")
)

func (hsdb *HSDatabase) CreateUser(name string) (*types.User, error) {
	return Write(hsdb.DB, func(tx *gorm.DB) (*types.User, error) {
		return CreateUser(tx, name)
	})
}

// CreateUser creates a new User. Returns error if could not be created
// or another user already exists.
func CreateUser(tx *gorm.DB, name string) (*types.User, error) {
	err := util.CheckForFQDNRules(name)
	if err != nil {
		return nil, err
	}
	user := types.User{}
	if err := tx.Where("name = ?", name).First(&user).Error; err == nil {
		return nil, ErrUserExists
	}
	user.Name = name
	if err := tx.Create(&user).Error; err != nil {
		log.Error().
			Str("func", "CreateUser").
			Err(err).
			Msg("Could not create row")

		return nil, err
	}

	return &user, nil
}

func (hsdb *HSDatabase) DestroyUser(name string) error {
	return hsdb.Write(func(tx *gorm.DB) error {
		return DestroyUser(tx, name)
	})
}

// DestroyUser destroys a User. Returns error if the User does
// not exist or if there are nodes associated with it.
func DestroyUser(tx *gorm.DB, name string) error {
	user, err := GetUser(tx, name)
	if err != nil {
		return ErrUserNotFound
	}

	nodes, err := ListNodesByUser(tx, name)
	if err != nil {
		return err
	}
	if len(nodes) > 0 {
		return ErrUserStillHasNodes
	}

	keys, err := ListPreAuthKeys(tx, name)
	if err != nil {
		return err
	}
	for _, key := range keys {
		err = DestroyPreAuthKey(tx, key)
		if err != nil {
			return err
		}
	}

	if result := tx.Unscoped().Delete(&user); result.Error != nil {
		return result.Error
	}

	return nil
}

func (hsdb *HSDatabase) RenameUser(oldName, newName string) error {
	return hsdb.Write(func(tx *gorm.DB) error {
		return RenameUser(tx, oldName, newName)
	})
}

// RenameUser renames a User. Returns error if the User does
// not exist or if another User exists with the new name.
func RenameUser(tx *gorm.DB, oldName, newName string) error {
	var err error
	oldUser, err := GetUser(tx, oldName)
	if err != nil {
		return err
	}
	err = util.CheckForFQDNRules(newName)
	if err != nil {
		return err
	}
	_, err = GetUser(tx, newName)
	if err == nil {
		return ErrUserExists
	}
	if !errors.Is(err, ErrUserNotFound) {
		return err
	}

	oldUser.Name = newName

	if result := tx.Save(&oldUser); result.Error != nil {
		return result.Error
	}

	return nil
}

func (hsdb *HSDatabase) GetUser(name string) (*types.User, error) {
	return Read(hsdb.DB, func(rx *gorm.DB) (*types.User, error) {
		return GetUser(rx, name)
	})
}

func GetUser(tx *gorm.DB, name string) (*types.User, error) {
	user := types.User{}
	if result := tx.First(&user, "name = ?", name); errors.Is(
		result.Error,
		gorm.ErrRecordNotFound,
	) {
		return nil, ErrUserNotFound
	}

	return &user, nil
}

func (hsdb *HSDatabase) ListUsers() ([]types.User, error) {
	return Read(hsdb.DB, func(rx *gorm.DB) ([]types.User, error) {
		return ListUsers(rx)
	})
}

// ListUsers gets all the existing users.
func ListUsers(tx *gorm.DB) ([]types.User, error) {
	users := []types.User{}
	if err := tx.Find(&users).Error; err != nil {
		return nil, err
	}

	return users, nil
}

// ListNodesByUser gets all the nodes in a given user.
func ListNodesByUser(tx *gorm.DB, name string) (types.Nodes, error) {
	err := util.CheckForFQDNRules(name)
	if err != nil {
		return nil, err
	}
	user, err := GetUser(tx, name)
	if err != nil {
		return nil, err
	}

	nodes := types.Nodes{}
	if err := tx.Preload("AuthKey").Preload("AuthKey.User").Preload("User").Where(&types.Node{UserID: user.ID}).Find(&nodes).Error; err != nil {
		return nil, err
	}

	return nodes, nil
}

func (hsdb *HSDatabase) AssignNodeToUser(node *types.Node, username string) error {
	return hsdb.Write(func(tx *gorm.DB) error {
		return AssignNodeToUser(tx, node, username)
	})
}

// AssignNodeToUser assigns a Node to a user.
func AssignNodeToUser(tx *gorm.DB, node *types.Node, username string) error {
	err := util.CheckForFQDNRules(username)
	if err != nil {
		return err
	}
	user, err := GetUser(tx, username)
	if err != nil {
		return err
	}
	node.User = *user
	if result := tx.Save(&node); result.Error != nil {
		return result.Error
	}

	return nil
}