fix(master): recheck revocation after register, reject duplicate hello

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