mirror of
https://github.com/juanfont/headscale.git
synced 2026-09-10 02:31:59 +09:00
Kristoffer Dalby
373c60efe3
Drops the Add/Done bookkeeping, which in batcher.go was spread across three functions.
98 lines
2.0 KiB
Go
98 lines
2.0 KiB
Go
package state
|
|
|
|
import (
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/juanfont/headscale/hscontrol/types"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func newTestStateForSSHCheck() *State {
|
|
return &State{
|
|
sshCheckAuth: make(map[sshCheckPair]time.Time),
|
|
}
|
|
}
|
|
|
|
func TestSSHCheckAuth(t *testing.T) {
|
|
s := newTestStateForSSHCheck()
|
|
|
|
src := types.NodeID(1)
|
|
dst := types.NodeID(2)
|
|
otherDst := types.NodeID(3)
|
|
otherSrc := types.NodeID(4)
|
|
|
|
// No record initially
|
|
_, ok := s.GetLastSSHAuth(src, dst)
|
|
require.False(t, ok)
|
|
|
|
// Record auth for (src, dst)
|
|
s.SetLastSSHAuth(src, dst)
|
|
|
|
// Same src+dst: found
|
|
authTime, ok := s.GetLastSSHAuth(src, dst)
|
|
require.True(t, ok)
|
|
assert.WithinDuration(t, time.Now(), authTime, time.Second)
|
|
|
|
// Same src, different dst: not found (auth is per-pair)
|
|
_, ok = s.GetLastSSHAuth(src, otherDst)
|
|
require.False(t, ok)
|
|
|
|
// Different src: not found
|
|
_, ok = s.GetLastSSHAuth(otherSrc, dst)
|
|
require.False(t, ok)
|
|
}
|
|
|
|
func TestSSHCheckAuthClear(t *testing.T) {
|
|
s := newTestStateForSSHCheck()
|
|
|
|
s.SetLastSSHAuth(types.NodeID(1), types.NodeID(2))
|
|
s.SetLastSSHAuth(types.NodeID(1), types.NodeID(3))
|
|
|
|
_, ok := s.GetLastSSHAuth(types.NodeID(1), types.NodeID(2))
|
|
require.True(t, ok)
|
|
|
|
_, ok = s.GetLastSSHAuth(types.NodeID(1), types.NodeID(3))
|
|
require.True(t, ok)
|
|
|
|
// Clear
|
|
s.ClearSSHCheckAuth()
|
|
|
|
_, ok = s.GetLastSSHAuth(types.NodeID(1), types.NodeID(2))
|
|
require.False(t, ok)
|
|
|
|
_, ok = s.GetLastSSHAuth(types.NodeID(1), types.NodeID(3))
|
|
require.False(t, ok)
|
|
}
|
|
|
|
func TestSSHCheckAuthConcurrent(t *testing.T) {
|
|
s := newTestStateForSSHCheck()
|
|
|
|
var wg sync.WaitGroup
|
|
|
|
for i := range 100 {
|
|
wg.Go(func() {
|
|
src := types.NodeID(uint64(i % 10)) //nolint:gosec
|
|
dst := types.NodeID(uint64(i%5 + 10)) //nolint:gosec
|
|
|
|
s.SetLastSSHAuth(src, dst)
|
|
s.GetLastSSHAuth(src, dst)
|
|
})
|
|
}
|
|
|
|
wg.Wait()
|
|
|
|
// Clear concurrently with reads
|
|
wg.Go(func() {
|
|
s.ClearSSHCheckAuth()
|
|
})
|
|
|
|
wg.Go(func() {
|
|
s.GetLastSSHAuth(types.NodeID(1), types.NodeID(2))
|
|
})
|
|
|
|
wg.Wait()
|
|
}
|