feat(master): agent hub that authenticates agents and stores their inventory

internal/master/hub/hub.go

@@ -0,0 +1,215 @@
+// Package hub accepts agent connections, authenticates them, and turns their
+// messages into stored inventory and live UI signals.
+package hub
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+
+	"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"
+)
+
+// readLimit caps one incoming message. The library default of 32 KiB is too
+// small for a full snapshot from a busy machine.
+const readLimit = 1 << 20
+
+// Store is the storage the hub writes to. *store.Store satisfies it.
+type Store interface {
+	AgentByTokenHash(ctx context.Context, hash string) (model.Agent, error)
+	RecordHello(ctx context.Context, agentID int64, h protocol.Hello, at time.Time) error
+	ApplySnapshot(ctx context.Context, agentID int64, h protocol.Hello, services []protocol.Service, at time.Time) error
+	UpsertService(ctx context.Context, agentID int64, clusterID string, svc protocol.Service, at time.Time) error
+	RemoveService(ctx context.Context, agentID int64, clusterID, agentKey string, at time.Time) error
+	MarkSeen(ctx context.Context, agentID int64, at time.Time) error
+}
+
+// Publisher receives "something changed" signals. *events.Broker satisfies it.
+type Publisher interface {
+	Publish(t events.Topic)
+}
+
+// Hub is the master's endpoint for agents.
+type Hub struct {
+	store Store
+	pub   Publisher
+	log   *slog.Logger
+	now   func() time.Time
+
+	mu    sync.Mutex
+	conns map[int64]*websocket.Conn
+}
+
+// New returns a Hub. Mount it at GET /agent on the TLS listener.
+func New(store Store, pub Publisher, log *slog.Logger) *Hub {
+	return &Hub{store: store, pub: pub, log: log, now: time.Now, conns: make(map[int64]*websocket.Conn)}
+}
+
+// ServeHTTP authenticates the agent's token and runs its connection. The token
+// is only ever hashed; it is never logged or stored.
+func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	token, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
+	if !ok || token == "" {
+		http.Error(w, "unauthorized", http.StatusUnauthorized)
+		return
+	}
+	agent, err := h.store.AgentByTokenHash(r.Context(), agenttoken.Hash(token))
+	if err != nil && !errors.Is(err, model.ErrNotFound) {
+		h.log.Error("agent_lookup_failed", "error", err)
+	}
+	if err != nil || agent.RevokedAt != nil {
+		http.Error(w, "unauthorized", http.StatusUnauthorized)
+		return
+	}
+
+	c, err := websocket.Accept(w, r, nil)
+	if err != nil {
+		return
+	}
+	c.SetReadLimit(readLimit)
+	h.register(agent.ID, c)
+	defer h.unregister(agent.ID, c)
+
+	s := &session{hub: h, agent: agent}
+	err = s.run(r.Context(), c)
+	h.log.Info("agent_disconnected", "agent", agent.Name, "reason", err)
+}
+
+// Online reports whether the agent has a live connection.
+func (h *Hub) Online(agentID int64) bool {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	_, ok := h.conns[agentID]
+	return ok
+}
+
+// Disconnect closes the agent's live connection, if any. Call it after
+// revoking the agent so the revocation takes effect immediately.
+func (h *Hub) Disconnect(agentID int64) {
+	h.mu.Lock()
+	c := h.conns[agentID]
+	h.mu.Unlock()
+	if c != nil {
+		c.CloseNow()
+	}
+}
+
+func (h *Hub) register(id int64, c *websocket.Conn) {
+	h.mu.Lock()
+	old := h.conns[id]
+	h.conns[id] = c
+	h.mu.Unlock()
+	if old != nil {
+		old.CloseNow()
+	}
+}
+
+func (h *Hub) unregister(id int64, c *websocket.Conn) {
+	h.mu.Lock()
+	current := h.conns[id] == c
+	if current {
+		delete(h.conns, id)
+	}
+	h.mu.Unlock()
+	if !current {
+		return // replaced by a newer connection; that one owns the state
+	}
+	if err := h.store.MarkSeen(context.Background(), id, h.now()); err != nil {
+		h.log.Error("mark_seen_failed", "agent_id", id, "error", err)
+	}
+	h.pub.Publish(events.Agents)
+}
+
+// session is one agent connection's state.
+type session struct {
+	hub   *Hub
+	agent model.Agent
+	hello *protocol.Hello
+}
+
+// swarmManager reports whether this agent may report Swarm services. The same
+// rule gates snapshots, live upserts and live removes.
+func (s *session) swarmManager() bool {
+	return s.hello.SwarmManager && s.hello.SwarmClusterID != ""
+}
+
+func (s *session) run(ctx context.Context, c *websocket.Conn) error {
+	for {
+		var env protocol.Envelope
+		if err := wsjson.Read(ctx, c, &env); err != nil {
+			return err
+		}
+		if s.hello == nil && env.Type != protocol.TypeHello {
+			_ = c.Close(websocket.StatusPolicyViolation, "hello required first")
+			return fmt.Errorf("%q before hello", env.Type)
+		}
+		if err := s.handle(ctx, env); err != nil {
+			_ = c.Close(websocket.StatusPolicyViolation, "bad message")
+			return err
+		}
+	}
+}
+
+func (s *session) handle(ctx context.Context, env protocol.Envelope) error {
+	h := s.hub
+	switch env.Type {
+	case protocol.TypeHello:
+		v, err := protocol.Decode[protocol.Hello](env)
+		if err != nil {
+			return err
+		}
+		if err := h.store.RecordHello(ctx, s.agent.ID, v, h.now()); err != nil {
+			return err
+		}
+		s.hello = &v
+		h.log.Info("agent_connected", "agent", s.agent.Name, "hostname", v.Hostname, "swarm_manager", v.SwarmManager)
+		h.pub.Publish(events.Agents)
+	case protocol.TypeSnapshot:
+		v, err := protocol.Decode[protocol.Snapshot](env)
+		if err != nil {
+			return err
+		}
+		if err := h.store.ApplySnapshot(ctx, s.agent.ID, *s.hello, v.Services, h.now()); err != nil {
+			return err
+		}
+		h.pub.Publish(events.Services)
+	case protocol.TypeUpsert:
+		v, err := protocol.Decode[protocol.Upsert](env)
+		if err != nil {
+			return err
+		}
+		if v.Service.Kind == protocol.KindSwarm && !s.swarmManager() {
+			return nil
+		}
+		if err := h.store.UpsertService(ctx, s.agent.ID, s.hello.SwarmClusterID, v.Service, h.now()); err != nil {
+			return err
+		}
+		h.pub.Publish(events.Services)
+	case protocol.TypeRemove:
+		v, err := protocol.Decode[protocol.Remove](env)
+		if err != nil {
+			return err
+		}
+		if strings.HasPrefix(v.Key, string(protocol.KindSwarm)+"|") && !s.swarmManager() {
+			return nil
+		}
+		if err := h.store.RemoveService(ctx, s.agent.ID, s.hello.SwarmClusterID, v.Key, h.now()); err != nil {
+			return err
+		}
+		h.pub.Publish(events.Services)
+	default:
+		return fmt.Errorf("unknown message type %q", env.Type)
+	}
+	return nil
+}

internal/master/hub/hub_test.go

@@ -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)
+}