mirror of
https://github.com/juanfont/headscale.git
synced 2026-10-30 18:24:06 +09:00
316 lines
9.4 KiB
Go
316 lines
9.4 KiB
Go
package state
|
|
|
|
import (
|
|
"bytes"
|
|
"runtime"
|
|
"runtime/pprof"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/hashicorp/golang-lru/v2/expirable"
|
|
"github.com/juanfont/headscale/hscontrol/types"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func newTestAuthCache(t *testing.T, maxEntries int, ttl time.Duration) *expirable.LRU[types.AuthID, *types.AuthRequest] {
|
|
t.Helper()
|
|
|
|
cache := expirable.NewLRU[types.AuthID, *types.AuthRequest](
|
|
maxEntries,
|
|
func(_ types.AuthID, rn *types.AuthRequest) {
|
|
rn.FinishAuth(types.AuthVerdict{Err: ErrRegistrationExpired})
|
|
},
|
|
ttl,
|
|
)
|
|
t.Cleanup(cache.Close)
|
|
|
|
return cache
|
|
}
|
|
|
|
func requireExpiredVerdict(t *testing.T, entry *types.AuthRequest) {
|
|
t.Helper()
|
|
|
|
select {
|
|
case verdict := <-entry.WaitForAuth():
|
|
require.False(t, verdict.Accept(), "evicted entry must not signal Accept")
|
|
require.ErrorIs(t,
|
|
verdict.Err, ErrRegistrationExpired,
|
|
"evicted entry must surface ErrRegistrationExpired, got: %v",
|
|
verdict.Err,
|
|
)
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("eviction callback did not wake the parked AuthRequest")
|
|
}
|
|
}
|
|
|
|
// TestAuthCacheBoundedLRU checks the underlying library's eviction callback.
|
|
// State admission prevents this eviction behavior for pending auth requests.
|
|
func TestAuthCacheBoundedLRU(t *testing.T) {
|
|
const maxEntries = 4
|
|
|
|
// Long TTL: eviction by size, not by time.
|
|
cache := newTestAuthCache(t, maxEntries, time.Hour)
|
|
|
|
entries := make([]*types.AuthRequest, 0, maxEntries+1)
|
|
ids := make([]types.AuthID, 0, maxEntries+1)
|
|
|
|
for range maxEntries + 1 {
|
|
id := types.MustAuthID()
|
|
entry := types.NewAuthRequest()
|
|
cache.Add(id, entry)
|
|
ids = append(ids, id)
|
|
entries = append(entries, entry)
|
|
}
|
|
|
|
// Cap should be respected.
|
|
assert.Equal(t, maxEntries, cache.Len(), "cache must not exceed the configured maxEntries")
|
|
|
|
// The oldest entry must have been evicted.
|
|
_, ok := cache.Get(ids[0])
|
|
assert.False(t, ok, "oldest entry must be evicted when maxEntries is exceeded")
|
|
|
|
requireExpiredVerdict(t, entries[0])
|
|
|
|
// All non-evicted entries must still be retrievable.
|
|
for i := 1; i <= maxEntries; i++ {
|
|
_, ok := cache.Get(ids[i])
|
|
assert.True(t, ok, "non-evicted entry %d should still be in the cache", i)
|
|
}
|
|
}
|
|
|
|
// TestAuthCacheTTLExpiry verifies the TTL sweeper evicts an entry that timed
|
|
// out and that the eviction callback wakes its waiter.
|
|
func TestAuthCacheTTLExpiry(t *testing.T) {
|
|
cache := newTestAuthCache(t, 4, 50*time.Millisecond)
|
|
|
|
id := types.MustAuthID()
|
|
entry := types.NewAuthRequest()
|
|
cache.Add(id, entry)
|
|
|
|
requireExpiredVerdict(t, entry)
|
|
|
|
_, ok := cache.Get(id)
|
|
assert.False(t, ok, "expired entry must be gone from the cache")
|
|
}
|
|
|
|
func TestAuthCacheAdmissionPreservesPending(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
cfg.Tuning.RegisterCacheMaxEntries = 1
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
id := types.MustAuthID()
|
|
pending := types.NewAuthRequest()
|
|
require.NoError(t, s.SetAuthCacheEntry(id, pending))
|
|
|
|
otherID := types.MustAuthID()
|
|
require.ErrorIs(t, s.SetAuthCacheEntry(otherID, types.NewAuthRequest()), ErrAuthCacheFull)
|
|
|
|
got, ok := s.GetAuthCacheEntry(id)
|
|
require.True(t, ok, "new work must not displace an admitted authentication")
|
|
require.Same(t, pending, got)
|
|
|
|
_, ok = s.GetAuthCacheEntry(otherID)
|
|
require.False(t, ok, "new work must be refused at capacity")
|
|
|
|
select {
|
|
case <-pending.WaitForAuth():
|
|
t.Fatal("capacity refusal must not finish an existing authentication")
|
|
default:
|
|
}
|
|
}
|
|
|
|
func TestAuthCacheAdmissionPreservesIdentity(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
id := types.MustAuthID()
|
|
pending := types.NewAuthRequest()
|
|
require.NoError(t, s.SetAuthCacheEntry(id, pending))
|
|
require.NoError(t, s.SetAuthCacheEntry(id, pending))
|
|
require.ErrorIs(t, s.SetAuthCacheEntry(id, types.NewAuthRequest()), ErrAuthCacheEntryExists)
|
|
|
|
got, ok := s.GetAuthCacheEntry(id)
|
|
require.True(t, ok)
|
|
require.Same(t, pending, got, "an existing auth ID must keep its original request and waiter")
|
|
}
|
|
|
|
func TestAuthCacheAdmissionConcurrent(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
cfg.Tuning.RegisterCacheMaxEntries = 2
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
id := types.MustAuthID()
|
|
pending := types.NewAuthRequest()
|
|
require.NoError(t, s.SetAuthCacheEntry(id, pending))
|
|
|
|
var wg sync.WaitGroup
|
|
|
|
results := make(chan error, 4)
|
|
|
|
for range 4 {
|
|
wg.Go(func() {
|
|
results <- s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest())
|
|
})
|
|
}
|
|
|
|
wg.Wait()
|
|
close(results)
|
|
|
|
var accepted int
|
|
|
|
for err := range results {
|
|
if err == nil {
|
|
accepted++
|
|
} else {
|
|
require.ErrorIs(t, err, ErrAuthCacheFull)
|
|
}
|
|
}
|
|
|
|
require.Equal(t, 1, accepted)
|
|
|
|
got, ok := s.GetAuthCacheEntry(id)
|
|
require.True(t, ok, "concurrent admission must preserve existing transactions")
|
|
require.Same(t, pending, got)
|
|
require.Equal(t, 2, s.authCache.Len())
|
|
}
|
|
|
|
func TestAuthCacheAdmissionReservesSSHCapacity(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
cfg.Tuning.RegisterCacheMaxEntries = 1
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
registrationID := types.MustAuthID()
|
|
require.NoError(t, s.SetAuthCacheEntry(registrationID, types.NewAuthRequest()))
|
|
|
|
sshID := types.MustAuthID()
|
|
require.NoError(t, s.SetAuthCacheEntry(sshID, types.NewSSHCheckAuthRequest(1, 2)),
|
|
"registration capacity must not consume the SSH reservation")
|
|
require.ErrorIs(t, s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest()), ErrAuthCacheFull)
|
|
_, ok := s.GetAuthCacheEntry(registrationID)
|
|
require.True(t, ok)
|
|
_, ok = s.GetAuthCacheEntry(sshID)
|
|
require.True(t, ok)
|
|
}
|
|
|
|
func TestAuthCacheAdmissionBoundsRegistrationPayload(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
id := types.MustAuthID()
|
|
request := types.NewRegisterAuthRequest(&types.RegistrationData{
|
|
Hostname: strings.Repeat("x", types.MaxRegistrationPayloadBytes),
|
|
})
|
|
require.ErrorIs(t, s.SetAuthCacheEntry(id, request), types.ErrRegistrationPayloadTooLarge)
|
|
_, ok := s.GetAuthCacheEntry(id)
|
|
require.False(t, ok)
|
|
require.Zero(t, s.authCache.Len())
|
|
require.Zero(t, s.authCacheRegularEntries.Load())
|
|
}
|
|
|
|
func TestAuthCacheAdmissionDoesNotReuseExpiredRequest(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
cfg.Tuning.RegisterCacheExpiration = 50 * time.Millisecond
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
id := types.MustAuthID()
|
|
request := types.NewAuthRequest()
|
|
require.NoError(t, s.SetAuthCacheEntry(id, request))
|
|
requireExpiredVerdict(t, request)
|
|
require.Error(t, s.SetAuthCacheEntry(id, request), "expired waiters cannot be reused as new pending requests")
|
|
}
|
|
|
|
func TestAuthCacheAdmissionReleasesCapacity(t *testing.T) {
|
|
for _, removal := range []string{"consumption", "expiry"} {
|
|
t.Run(removal, func(t *testing.T) {
|
|
cfg := persistTestConfig(t.TempDir() + "/headscale.db")
|
|
cfg.Tuning.RegisterCacheMaxEntries = 1
|
|
cfg.Tuning.RegisterCacheExpiration = 50 * time.Millisecond
|
|
s, err := NewState(cfg)
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, s.Close()) })
|
|
|
|
id := types.MustAuthID()
|
|
pending := types.NewAuthRequest()
|
|
require.NoError(t, s.SetAuthCacheEntry(id, pending))
|
|
|
|
if removal == "consumption" {
|
|
pending.FinishAuth(types.AuthVerdict{})
|
|
s.DeleteAuthCacheEntry(id)
|
|
} else {
|
|
requireExpiredVerdict(t, pending)
|
|
}
|
|
|
|
require.NoError(t, s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest()))
|
|
require.Equal(t, int64(1), s.authCacheRegularEntries.Load())
|
|
})
|
|
}
|
|
}
|
|
|
|
func goroutineID() string {
|
|
buf := make([]byte, 64)
|
|
n := runtime.Stack(buf, false)
|
|
|
|
// "goroutine 123 [running]:"
|
|
return strings.Fields(string(buf[:n]))[1]
|
|
}
|
|
|
|
// authCacheSweepers counts expirable LRU TTL sweeper goroutines started from
|
|
// goroutine gid, so caches other tests are still tearing down do not count.
|
|
func authCacheSweepers(t *testing.T, gid string) int {
|
|
t.Helper()
|
|
|
|
var buf bytes.Buffer
|
|
|
|
require.NoError(t, pprof.Lookup("goroutine").WriteTo(&buf, 2))
|
|
|
|
n := 0
|
|
|
|
for g := range strings.SplitSeq(buf.String(), "\n\n") {
|
|
if strings.Contains(g, "golang-lru/v2/expirable") &&
|
|
strings.Contains(g, " in goroutine "+gid+"\n") {
|
|
n++
|
|
}
|
|
}
|
|
|
|
return n
|
|
}
|
|
|
|
// TestStateCloseStopsAuthCacheSweeper verifies that closing a State stops
|
|
// the auth cache's TTL sweeper goroutine instead of leaking it, and that a
|
|
// request still pending at that point gets its expiry verdict rather than
|
|
// waiting forever for a sweep that never comes.
|
|
func TestStateCloseStopsAuthCacheSweeper(t *testing.T) {
|
|
gid := goroutineID()
|
|
|
|
s, err := NewState(persistTestConfig(t.TempDir() + "/headscale.db"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, 1, authCacheSweepers(t, gid), "NewState must start a sweeper")
|
|
|
|
pending := types.NewAuthRequest()
|
|
require.NoError(t, s.SetAuthCacheEntry(types.MustAuthID(), pending))
|
|
|
|
require.NoError(t, s.Close())
|
|
|
|
requireExpiredVerdict(t, pending)
|
|
require.ErrorIs(t, s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest()), ErrAuthCacheClosed)
|
|
|
|
assert.Eventually(t, func() bool {
|
|
return authCacheSweepers(t, gid) == 0
|
|
}, 5*time.Second, 10*time.Millisecond, "sweeper must stop on State.Close")
|
|
}
|