mirror of
https://github.com/juanfont/headscale.git
synced 2025-11-20 17:56:02 -05:00
457 lines
13 KiB
Go
457 lines
13 KiB
Go
package db
|
|
|
|
import (
|
|
"fmt"
|
|
"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) {
|
|
t.Helper()
|
|
|
|
_, err := db.CreatePreAuthKey(12345, true, false, nil, nil)
|
|
assert.Error(t, err)
|
|
},
|
|
},
|
|
{
|
|
name: "success_create_and_list",
|
|
test: func(t *testing.T, db *HSDatabase) {
|
|
t.Helper()
|
|
|
|
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) {
|
|
t.Helper()
|
|
|
|
_, 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) {
|
|
t.Helper()
|
|
|
|
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) {
|
|
t.Helper()
|
|
|
|
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"})
|
|
require.NoError(t, err)
|
|
|
|
key, err := db.CreatePreAuthKey(types.UserID(user.ID), false, false, nil, []string{"tag:good"})
|
|
require.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) {
|
|
t.Helper()
|
|
|
|
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) {
|
|
t.Helper()
|
|
|
|
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.GreaterOrEqual(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)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMultipleLegacyKeysAllowed(t *testing.T) {
|
|
db, err := newSQLiteTestDB()
|
|
require.NoError(t, err)
|
|
|
|
user, err := db.CreateUser(types.User{Name: "test-legacy"})
|
|
require.NoError(t, err)
|
|
|
|
// Create multiple legacy keys by directly inserting with empty prefix
|
|
// This simulates the migration scenario where existing databases have multiple
|
|
// plaintext keys without prefix/hash fields
|
|
now := time.Now()
|
|
|
|
for i := range 5 {
|
|
legacyKey := fmt.Sprintf("legacy_key_%d_%s", i, strings.Repeat("x", 40))
|
|
|
|
err := db.DB.Exec(`
|
|
INSERT INTO pre_auth_keys (key, prefix, hash, user_id, reusable, ephemeral, used, created_at)
|
|
VALUES (?, '', NULL, ?, ?, ?, ?, ?)
|
|
`, legacyKey, user.ID, true, false, false, now).Error
|
|
require.NoError(t, err, "should allow multiple legacy keys with empty prefix")
|
|
}
|
|
|
|
// Verify all legacy keys can be retrieved
|
|
var legacyKeys []types.PreAuthKey
|
|
|
|
err = db.DB.Where("prefix = '' OR prefix IS NULL").Find(&legacyKeys).Error
|
|
require.NoError(t, err)
|
|
assert.Len(t, legacyKeys, 5, "should have created 5 legacy keys")
|
|
|
|
// Now create new bcrypt-based keys - these should have unique prefixes
|
|
key1, err := db.CreatePreAuthKey(types.UserID(user.ID), true, false, nil, nil)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, key1.Key)
|
|
|
|
key2, err := db.CreatePreAuthKey(types.UserID(user.ID), true, false, nil, nil)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, key2.Key)
|
|
|
|
// Verify the new keys have different prefixes
|
|
pak1, err := db.GetPreAuthKey(key1.Key)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, pak1.Prefix)
|
|
|
|
pak2, err := db.GetPreAuthKey(key2.Key)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, pak2.Prefix)
|
|
|
|
assert.NotEqual(t, pak1.Prefix, pak2.Prefix, "new keys should have unique prefixes")
|
|
|
|
// Verify we cannot manually insert duplicate non-empty prefixes
|
|
duplicatePrefix := "test_prefix1"
|
|
hash1 := []byte("hash1")
|
|
hash2 := []byte("hash2")
|
|
|
|
// First insert should succeed
|
|
err = db.DB.Exec(`
|
|
INSERT INTO pre_auth_keys (key, prefix, hash, user_id, reusable, ephemeral, used, created_at)
|
|
VALUES ('', ?, ?, ?, ?, ?, ?, ?)
|
|
`, duplicatePrefix, hash1, user.ID, true, false, false, now).Error
|
|
require.NoError(t, err, "first key with prefix should succeed")
|
|
|
|
// Second insert with same prefix should fail
|
|
err = db.DB.Exec(`
|
|
INSERT INTO pre_auth_keys (key, prefix, hash, user_id, reusable, ephemeral, used, created_at)
|
|
VALUES ('', ?, ?, ?, ?, ?, ?, ?)
|
|
`, duplicatePrefix, hash2, user.ID, true, false, false, now).Error
|
|
require.Error(t, err, "duplicate non-empty prefix should be rejected")
|
|
assert.Contains(t, err.Error(), "UNIQUE constraint failed", "should fail with UNIQUE constraint error")
|
|
}
|