mirror of
https://github.com/juanfont/headscale.git
synced 2026-08-14 11:20:44 +09:00
Kristoffer Dalby
560b6d81ad
With the v1 API served by Huma and the CLI on the HTTP client, nothing uses the gRPC service or its generated code. Delete proto/, the generated gen/go and gen/openapiv2, grpcv1.go, the buf config, the proto .Proto() helpers and gRPC config, and the proto build tooling from the flake and CI.
503 lines
14 KiB
Go
503 lines
14 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"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
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(new(types.UserID(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(user.TypedID(), true, false, nil, nil)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, key.Key)
|
|
|
|
// List keys for the user
|
|
keys, err := db.ListPreAuthKeys()
|
|
require.NoError(t, err)
|
|
assert.Len(t, keys, 1)
|
|
|
|
// Verify User association is populated
|
|
assert.Equal(t, user.ID, keys[0].User.ID)
|
|
},
|
|
},
|
|
}
|
|
|
|
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(user.TypedID(), 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(user.TypedID(), false, false, nil, tagsWithDuplicate)
|
|
require.NoError(t, err)
|
|
|
|
listedPaks, err := db.ListPreAuthKeys()
|
|
require.NoError(t, err)
|
|
require.Len(t, listedPaks, 1)
|
|
|
|
gotTags := slices.Clone(listedPaks[0].Tags)
|
|
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(user.TypedID(), false, false, nil, []string{"tag:good"})
|
|
require.NoError(t, err)
|
|
|
|
node := types.Node{
|
|
ID: 0,
|
|
Hostname: "testest",
|
|
UserID: &user.ID,
|
|
RegisterMethod: util.RegisterMethodAuthKey,
|
|
AuthKeyID: new(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(
|
|
user.TypedID(),
|
|
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(
|
|
user.TypedID(),
|
|
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(
|
|
user.TypedID(),
|
|
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(user.TypedID(), true, false, nil, nil)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, key1.Key)
|
|
|
|
key2, err := db.CreatePreAuthKey(user.TypedID(), 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")
|
|
}
|
|
|
|
// TestUsePreAuthKeyAtomicCAS verifies that UsePreAuthKey is an atomic
|
|
// compare-and-set: a second call against an already-used key reports
|
|
// PAKError("authkey already used") rather than silently succeeding.
|
|
func TestUsePreAuthKeyAtomicCAS(t *testing.T) {
|
|
db, err := newSQLiteTestDB()
|
|
require.NoError(t, err)
|
|
|
|
user, err := db.CreateUser(types.User{Name: "atomic-cas"})
|
|
require.NoError(t, err)
|
|
|
|
pakNew, err := db.CreatePreAuthKey(user.TypedID(), false /* reusable */, false, nil, nil)
|
|
require.NoError(t, err)
|
|
|
|
pak, err := db.GetPreAuthKey(pakNew.Key)
|
|
require.NoError(t, err)
|
|
require.False(t, pak.Reusable, "test sanity: key must be single-use")
|
|
|
|
// First Use should commit cleanly.
|
|
err = db.Write(func(tx *gorm.DB) error {
|
|
return UsePreAuthKey(tx, pak)
|
|
})
|
|
require.NoError(t, err, "first UsePreAuthKey should succeed")
|
|
|
|
// Reload from disk to drop the in-memory Used=true the first call
|
|
// set on the struct, simulating a second concurrent transaction
|
|
// that loaded the same row before the first one committed.
|
|
stale, err := db.GetPreAuthKey(pakNew.Key)
|
|
require.NoError(t, err)
|
|
|
|
stale.Used = false
|
|
|
|
err = db.Write(func(tx *gorm.DB) error {
|
|
return UsePreAuthKey(tx, stale)
|
|
})
|
|
require.Error(t, err, "second UsePreAuthKey on the same single-use key must fail")
|
|
|
|
var pakErr types.PAKError
|
|
require.ErrorAs(t, err, &pakErr,
|
|
"second UsePreAuthKey error must be a PAKError, got: %v", err)
|
|
assert.Equal(t, "authkey already used", pakErr.Error())
|
|
}
|
|
|
|
// TestGetPreAuthKeyUnknownMapsToRecordNotFound ensures an unknown (or deleted)
|
|
// pre-auth key resolves to a record-not-found error, which the registration
|
|
// handler maps to a 401 rather than a raw server error.
|
|
func TestGetPreAuthKeyUnknownMapsToRecordNotFound(t *testing.T) {
|
|
db, err := newSQLiteTestDB()
|
|
require.NoError(t, err)
|
|
|
|
_, err = db.GetPreAuthKey("nonexistent-key")
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, gorm.ErrRecordNotFound,
|
|
"unknown pre-auth key must map to record-not-found (handled as 401)")
|
|
}
|