feat(master): agent hub that authenticates agents and stores their inventory
This commit is contained in:
@@ -0,0 +1,352 @@
|
||||
package hub
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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
|
||||
}
|
||||
|
||||
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()
|
||||
a, ok := f.agents[hash]
|
||||
if !ok {
|
||||
return model.Agent{}, model.ErrNotFound
|
||||
}
|
||||
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) 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))
|
||||
}
|
||||
|
||||
func TestHub_NewConnectionReplacesOld(t *testing.T) {
|
||||
store := newFakeStore(node2)
|
||||
h, _, srv := newHub(t, store)
|
||||
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")
|
||||
}
|
||||
|
||||
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)
|
||||
// The hello itself published Agents; forget that so the assertion below
|
||||
// can only be satisfied by the disconnect.
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user