diff --git a/internal/master/hub/hub.go b/internal/master/hub/hub.go new file mode 100644 index 0000000..ce6968f --- /dev/null +++ b/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 +} diff --git a/internal/master/hub/hub_test.go b/internal/master/hub/hub_test.go new file mode 100644 index 0000000..1df2ec7 --- /dev/null +++ b/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) +}