juanfont.headscale/hscontrol/db/preauth_keys_test.go

package db

import (
	"slices"
	"strings"
	"testing"
	"time"

	"github.com/juanfont/headscale/hscontrol/types"
	"github.com/juanfont/headscale/hscontrol/util"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"tailscale.com/types/ptr"
)

func TestCreatePreAuthKey(t *testing.T) {
	tests := []struct {
		name    string
		test    func(*testing.T, *HSDatabase)
	}{
		{
			name: "error_invalid_user_id",
			test: func(t *testing.T, db *HSDatabase) {
				_, err := db.CreatePreAuthKey(12345, true, false, nil, nil)
				assert.Error(t, err)
			},
		},
		{
			name: "success_create_and_list",
			test: func(t *testing.T, db *HSDatabase) {
				user, err := db.CreateUser(types.User{Name: "test"})
				require.NoError(t, err)

				key, err := db.CreatePreAuthKey(types.UserID(user.ID), true, false, nil, nil)
				require.NoError(t, err)
				assert.NotEmpty(t, key.Key)

				// List keys for the user
				keys, err := db.ListPreAuthKeys(types.UserID(user.ID))
				require.NoError(t, err)
				assert.Len(t, keys, 1)

				// Verify User association is populated
				assert.Equal(t, user.ID, keys[0].User.ID)
			},
		},
		{
			name: "error_list_invalid_user_id",
			test: func(t *testing.T, db *HSDatabase) {
				_, err := db.ListPreAuthKeys(1000000)
				assert.Error(t, err)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			db, err := newSQLiteTestDB()
			require.NoError(t, err)

			tt.test(t, db)
		})
	}
}

func TestPreAuthKeyACLTags(t *testing.T) {
	tests := []struct {
		name    string
		test    func(*testing.T, *HSDatabase)
	}{
		{
			name: "reject_malformed_tags",
			test: func(t *testing.T, db *HSDatabase) {
				user, err := db.CreateUser(types.User{Name: "test-tags-1"})
				require.NoError(t, err)

				_, err = db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, []string{"badtag"})
				assert.Error(t, err)
			},
		},
		{
			name: "deduplicate_and_sort_tags",
			test: func(t *testing.T, db *HSDatabase) {
				user, err := db.CreateUser(types.User{Name: "test-tags-2"})
				require.NoError(t, err)

				expectedTags := []string{"tag:test1", "tag:test2"}
				tagsWithDuplicate := []string{"tag:test1", "tag:test2", "tag:test2"}

				_, err = db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, tagsWithDuplicate)
				require.NoError(t, err)

				listedPaks, err := db.ListPreAuthKeys(types.UserID(user.ID))
				require.NoError(t, err)
				require.Len(t, listedPaks, 1)

				gotTags := listedPaks[0].Proto().GetAclTags()
				slices.Sort(gotTags)
				assert.Equal(t, expectedTags, gotTags)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			db, err := newSQLiteTestDB()
			require.NoError(t, err)

			tt.test(t, db)
		})
	}
}

func TestCannotDeleteAssignedPreAuthKey(t *testing.T) {
	db, err := newSQLiteTestDB()
	require.NoError(t, err)
	user, err := db.CreateUser(types.User{Name: "test8"})
	assert.NoError(t, err)

	key, err := db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, []string{"tag:good"})
	assert.NoError(t, err)

	node := types.Node{
		ID:             0,
		Hostname:       "testest",
		UserID:         user.ID,
		RegisterMethod: util.RegisterMethodAuthKey,
		AuthKeyID:      ptr.To(key.ID),
	}
	db.DB.Save(&node)

	err = db.DB.Delete(&types.PreAuthKey{ID: key.ID}).Error
	require.ErrorContains(t, err, "constraint failed: FOREIGN KEY constraint failed")
}

func TestPreAuthKeyAuthentication(t *testing.T) {
	db, err := newSQLiteTestDB()
	require.NoError(t, err)
	user := db.CreateUserForTest("test-user")

	tests := []struct {
		name           string
		setupKey       func() string // Returns key string to test
		wantFindErr    bool          // Error when finding the key
		wantValidateErr bool         // Error when validating the key
		validateResult func(*testing.T, *types.PreAuthKey)
	}{
		{
			name: "legacy_key_plaintext",
			setupKey: func() string {
				// Insert legacy key directly using GORM (simulate existing production key)
				// Note: We use raw SQL to bypass GORM's handling and set prefix to empty string
				// which simulates how legacy keys exist in production databases
				legacyKey := "abc123def456ghi789jkl012mno345pqr678stu901vwx234yz"
				now := time.Now()

				// Use raw SQL to insert with empty prefix to avoid UNIQUE constraint
				err := db.DB.Exec(`
					INSERT INTO pre_auth_keys (key, user_id, reusable, ephemeral, used, created_at)
					VALUES (?, ?, ?, ?, ?, ?)
				`, legacyKey, user.ID, true, false, false, now).Error
				require.NoError(t, err)
				return legacyKey
			},
			wantFindErr: false,
			wantValidateErr: false,
			validateResult: func(t *testing.T, pak *types.PreAuthKey) {
				assert.Equal(t, user.ID, pak.UserID)
				assert.NotEmpty(t, pak.Key) // Legacy keys have Key populated
				assert.Empty(t, pak.Prefix) // Legacy keys have empty Prefix
				assert.Nil(t, pak.Hash)     // Legacy keys have nil Hash
			},
		},
		{
			name: "new_key_bcrypt",
			setupKey: func() string {
				// Create new key via API
				keyStr, err := db.CreatePreAuthKey(
					types.UserID(user.ID),
					true, false, nil, []string{"tag:test"},
				)
				require.NoError(t, err)
				return keyStr.Key
			},
			wantFindErr: false,
			wantValidateErr: false,
			validateResult: func(t *testing.T, pak *types.PreAuthKey) {
				assert.Equal(t, user.ID, pak.UserID)
				assert.Empty(t, pak.Key)       // New keys have empty Key
				assert.NotEmpty(t, pak.Prefix) // New keys have Prefix
				assert.NotNil(t, pak.Hash)     // New keys have Hash
				assert.Len(t, pak.Prefix, 12)  // Prefix is 12 chars
			},
		},
		{
			name: "new_key_format_validation",
			setupKey: func() string {
				keyStr, err := db.CreatePreAuthKey(
					types.UserID(user.ID),
					true, false, nil, nil,
				)
				require.NoError(t, err)

				// Verify format: hskey-auth-{12-char-prefix}-{64-char-hash}
				// Use fixed-length parsing since prefix/hash can contain dashes (base64 URL-safe)
				assert.True(t, strings.HasPrefix(keyStr.Key, "hskey-auth-"))

				// Extract prefix and hash using fixed-length parsing like the real code does
				_, prefixAndHash, found := strings.Cut(keyStr.Key, "hskey-auth-")
				assert.True(t, found)
				assert.True(t, len(prefixAndHash) >= 12+1+64) // prefix + '-' + hash minimum

				prefix := prefixAndHash[:12]
				assert.Len(t, prefix, 12) // Prefix is 12 chars
				assert.Equal(t, byte('-'), prefixAndHash[12]) // Separator
				hash := prefixAndHash[13:]
				assert.Len(t, hash, 64) // Hash is 64 chars

				return keyStr.Key
			},
			wantFindErr: false,
			wantValidateErr: false,
		},
		{
			name: "invalid_bcrypt_hash",
			setupKey: func() string {
				// Create valid key
				key, err := db.CreatePreAuthKey(
					types.UserID(user.ID),
					true, false, nil, nil,
				)
				require.NoError(t, err)
				keyStr := key.Key

				// Return key with tampered hash using fixed-length parsing
				_, prefixAndHash, _ := strings.Cut(keyStr, "hskey-auth-")
				prefix := prefixAndHash[:12]
				return "hskey-auth-" + prefix + "-" + "wrong_hash_here_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "empty_key",
			setupKey: func() string {
				return ""
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "key_too_short",
			setupKey: func() string {
				return "hskey-auth-short"
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "missing_separator",
			setupKey: func() string {
				return "hskey-auth-ABCDEFGHIJKLabcdefghijklmnopqrstuvwxyz1234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ"
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "hash_too_short",
			setupKey: func() string {
				return "hskey-auth-ABCDEFGHIJKL-short"
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "prefix_with_invalid_chars",
			setupKey: func() string {
				return "hskey-auth-ABC$EF@HIJKL-" + strings.Repeat("a", 64)
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "hash_with_invalid_chars",
			setupKey: func() string {
				return "hskey-auth-ABCDEFGHIJKL-" + "invalid$chars" + strings.Repeat("a", 54)
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "prefix_not_found_in_db",
			setupKey: func() string {
				// Create a validly formatted key but with a prefix that doesn't exist
				return "hskey-auth-NotInDB12345-" + strings.Repeat("a", 64)
			},
			wantFindErr: true,
			wantValidateErr: false,
		},
		{
			name: "expired_legacy_key",
			setupKey: func() string {
				legacyKey := "expired_legacy_key_123456789012345678901234"
				now := time.Now()
				expiration := time.Now().Add(-1 * time.Hour) // Expired 1 hour ago

				// Use raw SQL to avoid UNIQUE constraint on empty prefix
				err := db.DB.Exec(`
					INSERT INTO pre_auth_keys (key, user_id, reusable, ephemeral, used, created_at, expiration)
					VALUES (?, ?, ?, ?, ?, ?, ?)
				`, legacyKey, user.ID, true, false, false, now, expiration).Error
				require.NoError(t, err)
				return legacyKey
			},
			wantFindErr: false,
			wantValidateErr: true,
		},
		{
			name: "used_single_use_legacy_key",
			setupKey: func() string {
				legacyKey := "used_legacy_key_123456789012345678901234567"
				now := time.Now()

				// Use raw SQL to avoid UNIQUE constraint on empty prefix
				err := db.DB.Exec(`
					INSERT INTO pre_auth_keys (key, user_id, reusable, ephemeral, used, created_at)
					VALUES (?, ?, ?, ?, ?, ?)
				`, legacyKey, user.ID, false, false, true, now).Error
				require.NoError(t, err)
				return legacyKey
			},
			wantFindErr: false,
			wantValidateErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			keyStr := tt.setupKey()

			pak, err := db.GetPreAuthKey(keyStr)

			if tt.wantFindErr {
				assert.Error(t, err)
				return
			}

			require.NoError(t, err)
			require.NotNil(t, pak)

			// Check validation if needed
			if tt.wantValidateErr {
				err := pak.Validate()
				assert.Error(t, err)
				return
			}

			if tt.validateResult != nil {
				tt.validateResult(t, pak)
			}
		})
	}
}