From 049324b2f4d249e43a15bb89aef02008c93d9dd1 Mon Sep 17 00:00:00 2001
From: Samuel Enocsson <samuel.enocsson@gmail.com>
Date: Mon, 12 Oct 2026 14:52:24 +0200
Subject: [PATCH] fix(master): recheck revocation after register, reject
 duplicate hello

---
 internal/master/hub/hub.go      | 12 ++++++++++
 internal/master/hub/hub_test.go | 42 +++++++++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/internal/master/hub/hub.go b/internal/master/hub/hub.go
index ce6968f..b0c920e 100644
--- a/internal/master/hub/hub.go
+++ b/internal/master/hub/hub.go
@@ -81,6 +81,15 @@ func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	h.register(agent.ID, c)
 	defer h.unregister(agent.ID, c)
 
+	// Revoking commits revoked_at before calling Disconnect. Re-checking after
+	// register closes the window where a revoke lands between the first lookup
+	// and register, when Disconnect would find nothing to cut.
+	again, err := h.store.AgentByTokenHash(r.Context(), agenttoken.Hash(token))
+	if err != nil || again.ID != agent.ID || again.RevokedAt != nil {
+		_ = c.Close(websocket.StatusPolicyViolation, "revoked")
+		return
+	}
+
 	s := &session{hub: h, agent: agent}
 	err = s.run(r.Context(), c)
 	h.log.Info("agent_disconnected", "agent", agent.Name, "reason", err)
@@ -165,6 +174,9 @@ func (s *session) handle(ctx context.Context, env protocol.Envelope) error {
 	h := s.hub
 	switch env.Type {
 	case protocol.TypeHello:
+		if s.hello != nil {
+			return errors.New("duplicate hello")
+		}
 		v, err := protocol.Decode[protocol.Hello](env)
 		if err != nil {
 			return err
diff --git a/internal/master/hub/hub_test.go b/internal/master/hub/hub_test.go
index 1df2ec7..6911f24 100644
--- a/internal/master/hub/hub_test.go
+++ b/internal/master/hub/hub_test.go
@@ -30,6 +30,10 @@ type fakeStore struct {
 	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
 }
 
 func newFakeStore(agents ...model.Agent) *fakeStore {
@@ -43,10 +47,14 @@ func newFakeStore(agents ...model.Agent) *fakeStore {
 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
 }
 
@@ -338,8 +346,10 @@ func TestHub_DisconnectCutsLiveConnection(t *testing.T) {
 	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.
+	// 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)
@@ -350,3 +360,31 @@ func TestHub_DisconnectCutsLiveConnection(t *testing.T) {
 	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) })())
+}