db/pak: make authkeys have identifiable hash

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

hscontrol/db/preauth_keys.go

@@ -1,16 +1,17 @@
 package db
 
 import (
-	"crypto/rand"
-	"encoding/hex"
 	"errors"
 	"fmt"
+	"slices"
 	"strings"
 	"time"
 
+	v2 "github.com/juanfont/headscale/hscontrol/policy/v2"
 	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/juanfont/headscale/hscontrol/util"
+	"golang.org/x/crypto/bcrypt"
 	"gorm.io/gorm"
-	"tailscale.com/util/set"
 )
 
 var (
@@ -19,72 +20,105 @@ var (
 	ErrSingleUseAuthKeyHasBeenUsed = errors.New("AuthKey has already been used")
 	ErrUserMismatch                = errors.New("user mismatch")
 	ErrPreAuthKeyACLTagInvalid     = errors.New("AuthKey tag is invalid")
+	ErrPreAuthKeyFailedToParse     = errors.New("failed to parse AuthKey")
 )
 
+const authKeyPrefix = "hskey-auth-"
+const authKeyPrefixLength = 12
+const authKeyLength = 64
+
 func (hsdb *HSDatabase) CreatePreAuthKey(
-	uid types.UserID,
+	uid *types.UserID,
 	reusable bool,
 	ephemeral bool,
 	expiration *time.Time,
-	aclTags []string,
-) (*types.PreAuthKey, error) {
-	return Write(hsdb.DB, func(tx *gorm.DB) (*types.PreAuthKey, error) {
-		return CreatePreAuthKey(tx, uid, reusable, ephemeral, expiration, aclTags)
+	tags []string,
+) (string, error) {
+	return Write(hsdb.DB, func(tx *gorm.DB) (string, error) {
+		return CreatePreAuthKey(tx, uid, reusable, ephemeral, expiration, tags)
 	})
 }
 
 // CreatePreAuthKey creates a new PreAuthKey in a user, and returns it.
+// A PreAuthKey can be tagged or owned by a user, but not both.
 func CreatePreAuthKey(
 	tx *gorm.DB,
-	uid types.UserID,
+	uid *types.UserID,
 	reusable bool,
 	ephemeral bool,
 	expiration *time.Time,
-	aclTags []string,
-) (*types.PreAuthKey, error) {
-	user, err := GetUserByID(tx, uid)
-	if err != nil {
-		return nil, err
+	tags []string,
+) (string, error) {
+	var err error
+	var user *types.User
+	var userID *uint
+
+	if uid == nil && len(tags) == 0 {
+		return "", errors.New("preauthkey must be either tagged or owned by user")
 	}
 
-	// Remove duplicates
-	aclTags = set.SetOf(aclTags).Slice()
+	if uid != nil && len(tags) > 0 {
+		return "", errors.New("preauthkey cannot be both tagged and owned by user")
+	}
 
-	// TODO(kradalby): factor out and create a reusable tag validation,
-	// check if there is one in Tailscale's lib.
-	for _, tag := range aclTags {
-		if !strings.HasPrefix(tag, "tag:") {
-			return nil, fmt.Errorf(
-				"%w: '%s' did not begin with 'tag:'",
-				ErrPreAuthKeyACLTagInvalid,
-				tag,
-			)
+	if uid != nil {
+		user, err = GetUserByID(tx, *uid)
+		if err != nil {
+			return "", err
+		}
+
+		userID = &user.ID
+	}
+
+	if len(tags) > 0 {
+		slices.Sort(tags)
+		tags = slices.Compact(tags)
+
+		for _, tag := range tags {
+			t := v2.Tag(tag)
+			if err := t.Validate(); err != nil {
+				return "", fmt.Errorf("invalid tag: %w", tag, err)
+			}
 		}
 	}
 
 	now := time.Now().UTC()
-	// TODO(kradalby): unify the key generations spread all over the code.
-	kstr, err := generateKey()
+
+	prefix, err := util.GenerateRandomStringURLSafe(apiPrefixLength)
 	if err != nil {
-		return nil, err
+		return "", err
+	}
+
+	toBeHashed, err := util.GenerateRandomStringURLSafe(apiKeyLength)
+	if err != nil {
+		return "", err
+	}
+
+	// Key to return to user, this will only be visible _once_
+	keyStr := authKeyPrefix + "-" + prefix + "-" + toBeHashed
+
+	hash, err := bcrypt.GenerateFromPassword([]byte(toBeHashed), bcrypt.DefaultCost)
+	if err != nil {
+		return "", err
 	}
 
 	key := types.PreAuthKey{
-		Key:        kstr,
-		UserID:     user.ID,
-		User:       *user,
 		Reusable:   reusable,
 		Ephemeral:  ephemeral,
 		CreatedAt:  &now,
 		Expiration: expiration,
-		Tags:       aclTags,
+		UserID:     userID,
+		User:       user,
+		Tags:       tags,
+		Prefix:     prefix,
+		Hash:       hash,
 	}
 
 	if err := tx.Save(&key).Error; err != nil {
-		return nil, fmt.Errorf("failed to create key in the database: %w", err)
+		return "", fmt.Errorf("failed to create key in the database: %w", err)
 	}
 
-	return &key, nil
+	return keyStr, nil
 }
 
 func (hsdb *HSDatabase) ListPreAuthKeys(uid types.UserID) ([]types.PreAuthKey, error) {
@@ -101,28 +135,68 @@ func ListPreAuthKeysByUser(tx *gorm.DB, uid types.UserID) ([]types.PreAuthKey, e
 	}
 
 	keys := []types.PreAuthKey{}
-	if err := tx.Preload("User").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
 	}
 
 	return keys, nil
 }
 
-func (hsdb *HSDatabase) GetPreAuthKey(key string) (*types.PreAuthKey, error) {
-	return Read(hsdb.DB, func(rx *gorm.DB) (*types.PreAuthKey, error) {
-		return GetPreAuthKey(rx, key)
+// GetPreAuthKey returns a PreAuthKey by its key string.
+// It will return an error if the key is not found, or if it is expired, used or invalid.
+func (hsdb *HSDatabase) GetPreAuthKey(keyStr string) (*types.PreAuthKey, error) {
+	return Write(hsdb.DB, func(tx *gorm.DB) (*types.PreAuthKey, error) {
+		return GetPreAuthKey(tx, keyStr)
 	})
 }
 
-// GetPreAuthKey returns a PreAuthKey for a given key. The caller is responsible
-// for checking if the key is usable (expired or used).
-func GetPreAuthKey(tx *gorm.DB, key string) (*types.PreAuthKey, error) {
-	pak := types.PreAuthKey{}
-	if err := tx.Preload("User").First(&pak, "key = ?", key).Error; err != nil {
+// GetPreAuthKey returns a PreAuthKey by its key string.
+// It will return an error if the key is not found, or if it is expired, used or invalid.
+func GetPreAuthKey(tx *gorm.DB, keyStr string) (*types.PreAuthKey, error) {
+	pak, err := findAuthKey(tx, keyStr)
+	if err != nil {
+		return nil, err
+	}
+
+	if pak.Expiration != nil && pak.Expiration.Before(time.Now()) {
+		return nil, ErrPreAuthKeyExpired
+	}
+
+	if pak.Used {
+		return nil, ErrSingleUseAuthKeyHasBeenUsed
+	}
+
+	return pak, nil
+}
+
+func findAuthKey(tx *gorm.DB, keyStr string) (*types.PreAuthKey, error) {
+	var pak *types.PreAuthKey
+	_, prefixAndHash, found := strings.Cut(keyStr, authKeyPrefix)
+
+	if !found {
+		if err := tx.Preload("User").First(pak, "key = ?", keyStr).Error; err != nil {
+			return nil, ErrPreAuthKeyNotFound
+		}
+	} else {
+		prefix, hash, found := strings.Cut(prefixAndHash, "-")
+		if !found {
+			return nil, ErrPreAuthKeyFailedToParse
+		}
+
+		if err := tx.Preload("User").First(pak, "prefix = ?", prefix).Error; err != nil {
+			return nil, ErrPreAuthKeyNotFound
+		}
+
+		if err := bcrypt.CompareHashAndPassword(pak.Hash, []byte(hash)); err != nil {
+			return nil, err
+		}
+	}
+
+	if pak == nil {
 		return nil, ErrPreAuthKeyNotFound
 	}
 
-	return &pak, nil
+	return pak, nil
 }
 
 // DestroyPreAuthKey destroys a preauthkey. Returns error if the PreAuthKey
@@ -161,13 +235,3 @@ func ExpirePreAuthKey(tx *gorm.DB, k *types.PreAuthKey) error {
 
 	return nil
 }
-
-func generateKey() (string, error) {
-	size := 24
-	bytes := make([]byte, size)
-	if _, err := rand.Read(bytes); err != nil {
-		return "", err
-	}
-
-	return hex.EncodeToString(bytes), nil
-}