455 lines
15 KiB
Go
455 lines
15 KiB
Go
package hub
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
"github.com/coder/websocket/wsjson"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/shcizo/package-updater/internal/agenttoken"
|
|
"github.com/shcizo/package-updater/internal/master/events"
|
|
"github.com/shcizo/package-updater/internal/master/model"
|
|
"github.com/shcizo/package-updater/internal/protocol"
|
|
)
|
|
|
|
type fakeStore struct {
|
|
mu sync.Mutex
|
|
agents map[string]model.Agent // by token hash
|
|
hellos []protocol.Hello
|
|
snapshots [][]protocol.Service
|
|
upserts []string // "clusterID|key"
|
|
removes []string // "clusterID|key"
|
|
seen []int64
|
|
lookups int
|
|
// revokeFromLookup, when > 0, makes AgentByTokenHash return the agent as
|
|
// revoked from that (1-based) lookup onwards.
|
|
revokeFromLookup int
|
|
|
|
// Job calls, recorded; the *Result fields script the answers.
|
|
lostCalls [][]int64
|
|
lostResult []int64
|
|
lostErr error
|
|
lostGate chan struct{} // when set, FailLostJobs blocks until it is closed
|
|
progress []string // "agent|job|chunk"
|
|
progressOK bool
|
|
finishes []model.JobFinish
|
|
finishAgents []int64
|
|
finishOK bool
|
|
finishErr error
|
|
}
|
|
|
|
func newFakeStore(agents ...model.Agent) *fakeStore {
|
|
f := &fakeStore{agents: map[string]model.Agent{}}
|
|
for _, a := range agents {
|
|
f.agents[agenttoken.Hash("pua_"+a.Name)] = a
|
|
}
|
|
return f
|
|
}
|
|
|
|
func (f *fakeStore) AgentByTokenHash(_ context.Context, hash string) (model.Agent, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.lookups++
|
|
a, ok := f.agents[hash]
|
|
if !ok {
|
|
return model.Agent{}, model.ErrNotFound
|
|
}
|
|
if f.revokeFromLookup > 0 && f.lookups >= f.revokeFromLookup {
|
|
a.RevokedAt = ptr(time.Now())
|
|
}
|
|
return a, nil
|
|
}
|
|
|
|
func (f *fakeStore) RecordHello(_ context.Context, _ int64, h protocol.Hello, _ time.Time) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.hellos = append(f.hellos, h)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeStore) ApplySnapshot(_ context.Context, _ int64, _ protocol.Hello, s []protocol.Service, _ time.Time) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.snapshots = append(f.snapshots, s)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeStore) UpsertService(_ context.Context, _ int64, cluster string, s protocol.Service, _ time.Time) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.upserts = append(f.upserts, cluster+"|"+s.Key())
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeStore) RemoveService(_ context.Context, _ int64, cluster, key string, _ time.Time) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.removes = append(f.removes, cluster+"|"+key)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeStore) MarkSeen(_ context.Context, id int64, _ time.Time) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.seen = append(f.seen, id)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeStore) FailLostJobs(_ context.Context, _ int64, known []int64, _ time.Time) ([]int64, error) {
|
|
if f.lostGate != nil {
|
|
<-f.lostGate
|
|
}
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.lostCalls = append(f.lostCalls, known)
|
|
return f.lostResult, f.lostErr
|
|
}
|
|
|
|
func (f *fakeStore) JobProgress(_ context.Context, agentID, jobID int64, chunk string, _ time.Time) (bool, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.progress = append(f.progress, fmt.Sprintf("%d|%d|%s", agentID, jobID, chunk))
|
|
return f.progressOK, nil
|
|
}
|
|
|
|
func (f *fakeStore) FinishJob(_ context.Context, agentID, _ int64, fin model.JobFinish) (bool, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.finishAgents = append(f.finishAgents, agentID)
|
|
f.finishes = append(f.finishes, fin)
|
|
return f.finishOK, f.finishErr
|
|
}
|
|
|
|
func (f *fakeStore) count(get func(*fakeStore) int) func() int {
|
|
return func() int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return get(f)
|
|
}
|
|
}
|
|
|
|
type fakePub struct {
|
|
mu sync.Mutex
|
|
topics []events.Topic
|
|
}
|
|
|
|
func (p *fakePub) Publish(t events.Topic) {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
p.topics = append(p.topics, t)
|
|
}
|
|
|
|
func (p *fakePub) has(t events.Topic) bool {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
for _, x := range p.topics {
|
|
if x == t {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// reset forgets every topic published so far.
|
|
func (p *fakePub) reset() {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
p.topics = nil
|
|
}
|
|
|
|
var (
|
|
node2 = model.Agent{ID: 2, Name: "node-2"}
|
|
mgr = model.Agent{ID: 1, Name: "swarm-mgr"}
|
|
revoked = model.Agent{ID: 3, Name: "old", RevokedAt: ptr(time.Now())}
|
|
)
|
|
|
|
func ptr[T any](v T) *T { return &v }
|
|
|
|
func newHub(t *testing.T, store *fakeStore) (*Hub, *fakePub, *httptest.Server) {
|
|
t.Helper()
|
|
pub := &fakePub{}
|
|
h := New(store, pub, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
|
srv := httptest.NewServer(h)
|
|
t.Cleanup(srv.Close)
|
|
return h, pub, srv
|
|
}
|
|
|
|
func dial(t *testing.T, srv *httptest.Server, token string) (*websocket.Conn, *http.Response, error) {
|
|
t.Helper()
|
|
header := http.Header{}
|
|
if token != "" {
|
|
header.Set("Authorization", "Bearer "+token)
|
|
}
|
|
return websocket.Dial(context.Background(), "ws"+strings.TrimPrefix(srv.URL, "http"),
|
|
&websocket.DialOptions{HTTPHeader: header})
|
|
}
|
|
|
|
func send(t *testing.T, c *websocket.Conn, typ string, payload any) {
|
|
t.Helper()
|
|
env, err := protocol.Encode(typ, payload)
|
|
require.NoError(t, err)
|
|
require.NoError(t, wsjson.Write(context.Background(), c, env))
|
|
}
|
|
|
|
func TestHub_RejectsBadTokens(t *testing.T) {
|
|
_, _, srv := newHub(t, newFakeStore(node2, revoked))
|
|
for name, token := range map[string]string{"missing": "", "unknown": "pua_nobody", "revoked": "pua_old"} {
|
|
t.Run(name, func(t *testing.T) {
|
|
_, resp, err := dial(t, srv, token)
|
|
require.Error(t, err)
|
|
require.NotNil(t, resp)
|
|
assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestHub_HelloAndSnapshotReachStore(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
h, pub, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
|
|
send(t, c, protocol.TypeSnapshot, protocol.Snapshot{Services: []protocol.Service{
|
|
{Kind: protocol.KindContainer, ContainerName: "a", Image: "a:1", Status: "running"},
|
|
}})
|
|
|
|
require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.snapshots) })() == 1 }, time.Second, 5*time.Millisecond)
|
|
assert.True(t, h.Online(node2.ID))
|
|
assert.True(t, pub.has(events.Agents))
|
|
assert.True(t, pub.has(events.Services))
|
|
}
|
|
|
|
func TestHub_LargeSnapshotArrivesIntact(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
services := make([]protocol.Service, 500)
|
|
for i := range services {
|
|
services[i] = protocol.Service{
|
|
Kind: protocol.KindContainer, ContainerName: strings.Repeat("x", 20) + string(rune('a'+i%26)) + strings.Repeat("y", i%7),
|
|
Image: "registry.example.com/some/long/image/name:v1.2.3@sha256:" + strings.Repeat("0", 64), Status: "running",
|
|
}
|
|
}
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
|
|
send(t, c, protocol.TypeSnapshot, protocol.Snapshot{Services: services})
|
|
|
|
require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.snapshots) })() == 1 }, time.Second, 5*time.Millisecond)
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
assert.Len(t, store.snapshots[0], 500)
|
|
}
|
|
|
|
func TestHub_UpsertAndRemoveCarryClusterID(t *testing.T) {
|
|
store := newFakeStore(mgr)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_swarm-mgr")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
gitea := protocol.Service{Kind: protocol.KindSwarm, SwarmService: "gitea_gitea", Image: "gitea/gitea:1", Status: "running"}
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true, SwarmClusterID: "c1"})
|
|
send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: gitea})
|
|
send(t, c, protocol.TypeRemove, protocol.Remove{Key: gitea.Key()})
|
|
|
|
require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.removes) })() == 1 }, time.Second, 5*time.Millisecond)
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
assert.Equal(t, []string{"c1|swarm|gitea_gitea"}, store.upserts)
|
|
assert.Equal(t, []string{"c1|swarm|gitea_gitea"}, store.removes)
|
|
}
|
|
|
|
func TestHub_IgnoresSwarmUpsertFromWorker(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
|
|
send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: protocol.Service{Kind: protocol.KindSwarm, SwarmService: "x"}})
|
|
send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: protocol.Service{Kind: protocol.KindContainer, ContainerName: "y"}})
|
|
|
|
require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.upserts) })() == 1 }, time.Second, 5*time.Millisecond)
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
assert.Equal(t, []string{"|container|y"}, store.upserts)
|
|
}
|
|
|
|
// Live swarm messages are gated exactly like snapshots: only a swarm manager
|
|
// with a cluster ID may upsert or remove swarm services.
|
|
func TestHub_SwarmLiveMessagesAreGatedOnManagerWithClusterID(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
hello protocol.Hello
|
|
wantReaches bool
|
|
}{
|
|
{"non-manager", protocol.Hello{Hostname: "swarm-mgr"}, false},
|
|
{"non-manager with cluster id", protocol.Hello{Hostname: "swarm-mgr", SwarmClusterID: "c1"}, false},
|
|
{"manager with empty cluster id", protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true}, false},
|
|
{"manager with cluster id", protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true, SwarmClusterID: "c1"}, true},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
store := newFakeStore(mgr)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_swarm-mgr")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
svc := protocol.Service{Kind: protocol.KindSwarm, SwarmService: "gitea_gitea"}
|
|
sentinel := protocol.Service{Kind: protocol.KindContainer, ContainerName: "sentinel"}
|
|
send(t, c, protocol.TypeHello, tt.hello)
|
|
send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: svc})
|
|
send(t, c, protocol.TypeRemove, protocol.Remove{Key: svc.Key()})
|
|
// Messages are handled in order, so once the sentinel's remove has
|
|
// arrived, the swarm messages before it have been handled too.
|
|
send(t, c, protocol.TypeRemove, protocol.Remove{Key: sentinel.Key()})
|
|
|
|
sentinelRemoved := func() bool {
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
for _, r := range store.removes {
|
|
if strings.HasSuffix(r, "|container|sentinel") {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
require.Eventually(t, sentinelRemoved, time.Second, 5*time.Millisecond)
|
|
|
|
store.mu.Lock()
|
|
defer store.mu.Unlock()
|
|
if tt.wantReaches {
|
|
assert.Equal(t, []string{"c1|swarm|gitea_gitea"}, store.upserts)
|
|
assert.Contains(t, store.removes, "c1|swarm|gitea_gitea")
|
|
} else {
|
|
assert.Empty(t, store.upserts)
|
|
assert.Equal(t, []string{tt.hello.SwarmClusterID + "|container|sentinel"}, store.removes)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestHub_ClosesConnectionThatSkipsHello(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
send(t, c, protocol.TypeSnapshot, protocol.Snapshot{})
|
|
_, _, err = c.Read(context.Background())
|
|
assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
|
|
}
|
|
|
|
// lockedBuffer is a bytes.Buffer safe for the hub's handler goroutines to
|
|
// write while the test reads it.
|
|
type lockedBuffer struct {
|
|
mu sync.Mutex
|
|
buf bytes.Buffer
|
|
}
|
|
|
|
func (b *lockedBuffer) Write(p []byte) (int, error) {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
return b.buf.Write(p)
|
|
}
|
|
|
|
func (b *lockedBuffer) String() string {
|
|
b.mu.Lock()
|
|
defer b.mu.Unlock()
|
|
return b.buf.String()
|
|
}
|
|
|
|
func TestHub_NewConnectionReplacesOld(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
logs := &lockedBuffer{}
|
|
h := New(store, &fakePub{}, slog.New(slog.NewTextHandler(logs, nil)))
|
|
srv := httptest.NewServer(h)
|
|
t.Cleanup(srv.Close)
|
|
first, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer first.CloseNow()
|
|
send(t, first, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
|
|
require.Eventually(t, func() bool { return h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
|
|
|
|
second, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer second.CloseNow()
|
|
|
|
_, _, err = first.Read(context.Background())
|
|
assert.Error(t, err, "the replaced connection must be closed")
|
|
assert.True(t, h.Online(node2.ID), "the agent stays online on its new connection")
|
|
assert.Contains(t, logs.String(), "msg=agent_replaced")
|
|
assert.Contains(t, logs.String(), "agent=node-2")
|
|
assert.NotContains(t, logs.String(), "pua_", "tokens are never logged")
|
|
}
|
|
|
|
func TestHub_DisconnectCutsLiveConnection(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
h, pub, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
|
|
require.Eventually(t, func() bool { return h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
|
|
// Online turns true before the hello is handled, so wait for the hello's
|
|
// own Agents signal, then forget it; the assertion below can then only be
|
|
// satisfied by the disconnect.
|
|
require.Eventually(t, func() bool { return pub.has(events.Agents) }, time.Second, 5*time.Millisecond)
|
|
pub.reset()
|
|
|
|
h.Disconnect(node2.ID)
|
|
|
|
_, _, err = c.Read(context.Background())
|
|
assert.Error(t, err)
|
|
require.Eventually(t, func() bool { return !h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
|
|
require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.seen) })() == 1 }, time.Second, 5*time.Millisecond)
|
|
require.Eventually(t, func() bool { return pub.has(events.Agents) }, time.Second, 5*time.Millisecond)
|
|
}
|
|
|
|
func TestHub_RevokeBetweenLookupAndRegisterCutsConnection(t *testing.T) {
|
|
store := newFakeStore(node2)
|
|
store.revokeFromLookup = 2 // first lookup passes, the re-check sees the revocation
|
|
h, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_node-2")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
_, _, err = c.Read(context.Background())
|
|
assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
|
|
require.Eventually(t, func() bool { return !h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
|
|
}
|
|
|
|
func TestHub_ClosesConnectionOnSecondHello(t *testing.T) {
|
|
store := newFakeStore(mgr)
|
|
_, _, srv := newHub(t, store)
|
|
c, _, err := dial(t, srv, "pua_swarm-mgr")
|
|
require.NoError(t, err)
|
|
defer c.CloseNow()
|
|
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "swarm-mgr"})
|
|
send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true, SwarmClusterID: "c1"})
|
|
|
|
_, _, err = c.Read(context.Background())
|
|
assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
|
|
assert.Equal(t, 1, store.count(func(f *fakeStore) int { return len(f.hellos) })())
|
|
}
|