From 5bb5d201e4cecd1186e8faee7e55072f62a4f18d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Proulx?= <76956526+fproulx-boostsecurity@users.noreply.github.com> Date: Thu, 6 Aug 2026 12:09:14 -0400 Subject: [PATCH] fix: make GraphHub publication revision-safe --- .claude/e2e/exploit_test.go | 6 + .claude/e2e/graph.go | 224 +++++++++++++++++++++++ internal/kitchen/browser_assets/graph.js | 68 +++++-- internal/kitchen/graph.go | 2 +- internal/kitchen/graph_hub.go | 214 +++++++++++----------- internal/kitchen/graph_hub_test.go | 194 +++++++++++++++++--- internal/kitchen/graph_ws.go | 24 ++- internal/kitchen/graph_ws_test.go | 33 +++- internal/pantry/graph.go | 8 + 9 files changed, 613 insertions(+), 160 deletions(-) create mode 100644 .claude/e2e/graph.go diff --git a/.claude/e2e/exploit_test.go b/.claude/e2e/exploit_test.go index f77ae0d..dbd68f6 100644 --- a/.claude/e2e/exploit_test.go +++ b/.claude/e2e/exploit_test.go @@ -18,15 +18,21 @@ func TestPublicExploitSmoke(t *testing.T) { root := findProjectRoot() require.NotEmpty(t, root, "could not find project root") + kitchenURL := getEnvOrFile("KITCHEN_URL", e2eEnvPath) + require.NotEmpty(t, kitchenURL, "KITCHEN_URL required") + authToken := getEnvOrFile("AUTH_TOKEN", e2eEnvPath) + require.NotEmpty(t, authToken, "AUTH_TOKEN required") require.NoError(t, resetE2EWorkspace()) require.NoError(t, waitForKitchenHealth(root, 90*time.Second)) + graphProbes := connectGraphSmokeProbes(t, kitchenURL, authToken) require.NoError(t, writeConfig(token, targetRepo)) require.NoError(t, restartCounter("e2e-smoke")) tmux := newTmuxController(tmuxSessionName) waitForReconPhase(t, tmux) + verifyGraphPublicationAfterAnalysis(t, graphProbes) requireContent(t, tmux, "whooli", 15*time.Second, "Attack tree should show target org") requireContent(t, tmux, "xyz", 15*time.Second, "Attack tree should show target repo") diff --git a/.claude/e2e/graph.go b/.claude/e2e/graph.go new file mode 100644 index 0000000..ee8dbf5 --- /dev/null +++ b/.claude/e2e/graph.go @@ -0,0 +1,224 @@ +// Copyright (C) 2026 boostsecurity.io +// SPDX-License-Identifier: AGPL-3.0-or-later + +//go:build e2e + +package e2e + +import ( + "context" + "encoding/json" + "fmt" + "net/url" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const graphProbeTimeout = 30 * time.Second + +type graphWireMessage struct { + Type string `json:"type"` + Data json.RawMessage `json:"data"` +} + +type graphWireSnapshot struct { + Revision uint64 `json:"revision"` + Mode string `json:"mode"` + TotalNodes int `json:"total_nodes"` + TotalEdges int `json:"total_edges"` + Nodes []json.RawMessage `json:"nodes"` + Edges []json.RawMessage `json:"edges"` +} + +type graphWireDelta struct { + BaseRevision uint64 `json:"base_revision"` + Revision uint64 `json:"revision"` +} + +type graphWireFence struct { + Revision uint64 `json:"revision"` +} + +type graphProbe struct { + connection *websocket.Conn + mode string + revision uint64 +} + +type graphSmokeProbes struct { + full *graphProbe + filtered *graphProbe + auto *graphProbe +} + +func connectGraphSmokeProbes(t *testing.T, kitchenURL, authToken string) *graphSmokeProbes { + t.Helper() + + return &graphSmokeProbes{ + full: connectGraphProbe(t, kitchenURL, authToken, "full"), + filtered: connectGraphProbe(t, kitchenURL, authToken, "filtered"), + auto: connectGraphProbe(t, kitchenURL, authToken, "auto"), + } +} + +func connectGraphProbe(t *testing.T, kitchenURL, authToken, mode string) *graphProbe { + t.Helper() + + endpoint, err := graphWebSocketURL(kitchenURL, authToken, mode) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), graphProbeTimeout) + defer cancel() + connection, _, err := websocket.Dial(ctx, endpoint, nil) + require.NoError(t, err) + connection.SetReadLimit(-1) + t.Cleanup(func() { _ = connection.Close(websocket.StatusNormalClosure, "") }) + + probe := &graphProbe{connection: connection, mode: mode} + message := probe.read(t) + require.Equal(t, "snapshot", message.Type, "%s mode must start from a complete snapshot", mode) + snapshot := decodeGraphData[graphWireSnapshot](t, message) + require.Equal(t, uint64(0), snapshot.Revision, "%s mode should connect to the purged Pantry", mode) + require.Zero(t, snapshot.TotalNodes) + require.Zero(t, snapshot.TotalEdges) + probe.revision = snapshot.Revision + return probe +} + +func graphWebSocketURL(kitchenURL, authToken, mode string) (string, error) { + endpoint, err := url.Parse(kitchenURL) + if err != nil { + return "", fmt.Errorf("parse Kitchen URL: %w", err) + } + switch endpoint.Scheme { + case "http": + endpoint.Scheme = "ws" + case "https": + endpoint.Scheme = "wss" + default: + return "", fmt.Errorf("unsupported Kitchen URL scheme %q", endpoint.Scheme) + } + endpoint.Path = "/graph/ws" + query := endpoint.Query() + query.Set("token", authToken) + query.Set("mode", mode) + endpoint.RawQuery = query.Encode() + return endpoint.String(), nil +} + +func verifyGraphPublicationAfterAnalysis(t *testing.T, probes *graphSmokeProbes) { + t.Helper() + + fullSnapshot := probes.full.refreshAfterFence(t) + require.Equal(t, "full", fullSnapshot.Mode) + require.NotZero(t, fullSnapshot.TotalNodes) + require.Len(t, fullSnapshot.Nodes, fullSnapshot.TotalNodes) + require.Len(t, fullSnapshot.Edges, fullSnapshot.TotalEdges) + + filteredSnapshot := probes.filtered.waitForSnapshot(t, fullSnapshot.Revision) + require.Equal(t, "filtered", filteredSnapshot.Mode) + require.NotZero(t, filteredSnapshot.TotalNodes) + assert.LessOrEqual(t, len(filteredSnapshot.Nodes), filteredSnapshot.TotalNodes) + assert.LessOrEqual(t, len(filteredSnapshot.Edges), filteredSnapshot.TotalEdges) + + autoSnapshot := probes.auto.waitForSnapshot(t, fullSnapshot.Revision) + require.NotZero(t, autoSnapshot.TotalNodes) + assert.Contains(t, []string{"full", "filtered"}, autoSnapshot.Mode) + assert.LessOrEqual(t, len(autoSnapshot.Nodes), autoSnapshot.TotalNodes) + assert.LessOrEqual(t, len(autoSnapshot.Edges), autoSnapshot.TotalEdges) +} + +func (p *graphProbe) refreshAfterFence(t *testing.T) graphWireSnapshot { + t.Helper() + + minimumRevision := p.revision + for { + message := p.read(t) + switch message.Type { + case "delta": + delta := decodeGraphData[graphWireDelta](t, message) + require.Equal(t, p.revision, delta.BaseRevision, "full-mode deltas must be contiguous before a fence") + require.Greater(t, delta.Revision, p.revision) + p.revision = delta.Revision + case "snapshot_required": + fence := decodeGraphData[graphWireFence](t, message) + require.Greater(t, fence.Revision, p.revision) + minimumRevision = fence.Revision + p.requestSnapshot(t) + return p.waitForFullSnapshot(t, minimumRevision) + default: + t.Fatalf("full mode received unexpected graph message %q before committed-state fence", message.Type) + } + } +} + +func (p *graphProbe) waitForFullSnapshot(t *testing.T, minimumRevision uint64) graphWireSnapshot { + t.Helper() + + for { + message := p.read(t) + switch message.Type { + case "snapshot": + snapshot := decodeGraphData[graphWireSnapshot](t, message) + require.GreaterOrEqual(t, snapshot.Revision, minimumRevision) + p.revision = snapshot.Revision + return snapshot + case "delta": + delta := decodeGraphData[graphWireDelta](t, message) + minimumRevision = max(minimumRevision, delta.Revision) + case "snapshot_required": + fence := decodeGraphData[graphWireFence](t, message) + minimumRevision = max(minimumRevision, fence.Revision) + default: + t.Fatalf("full mode received unexpected graph message %q while refreshing", message.Type) + } + } +} + +func (p *graphProbe) waitForSnapshot(t *testing.T, minimumRevision uint64) graphWireSnapshot { + t.Helper() + + for { + message := p.read(t) + require.Equal(t, "snapshot", message.Type, "%s mode must receive complete snapshots while refreshing", p.mode) + snapshot := decodeGraphData[graphWireSnapshot](t, message) + require.GreaterOrEqual(t, snapshot.Revision, p.revision) + p.revision = snapshot.Revision + if snapshot.Revision >= minimumRevision { + return snapshot + } + } +} + +func (p *graphProbe) requestSnapshot(t *testing.T) { + t.Helper() + + ctx, cancel := context.WithTimeout(context.Background(), graphProbeTimeout) + defer cancel() + require.NoError(t, wsjson.Write(ctx, p.connection, map[string]any{ + "type": "snapshot_request", + "data": map[string]string{"mode": p.mode}, + })) +} + +func (p *graphProbe) read(t *testing.T) graphWireMessage { + t.Helper() + + ctx, cancel := context.WithTimeout(context.Background(), graphProbeTimeout) + defer cancel() + var message graphWireMessage + require.NoError(t, wsjson.Read(ctx, p.connection, &message), "%s graph stream did not publish in time", p.mode) + return message +} + +func decodeGraphData[T any](t *testing.T, message graphWireMessage) T { + t.Helper() + + var data T + require.NoError(t, json.Unmarshal(message.Data, &data)) + return data +} diff --git a/internal/kitchen/browser_assets/graph.js b/internal/kitchen/browser_assets/graph.js index 515bd8c..2eb802d 100644 --- a/internal/kitchen/browser_assets/graph.js +++ b/internal/kitchen/browser_assets/graph.js @@ -31,12 +31,14 @@ SPDX-License-Identifier: AGPL-3.0-or-later let cy; let ws; - let graphVersion = 0; + let graphRevision = 0; let reconnectAttempts = 0; let requestedGraphMode = readGraphMode(); let resolvedGraphMode = 'full'; let graphMeta = { totalNodes: 0, totalEdges: 0, largeGraph: false, filterDescription: '' }; - let filteredRefreshTimer = null; + let minimumSnapshotRevision = 0; + let snapshotRequestPending = false; + let awaitingInitialSnapshot = true; function initCytoscape() { cy = cytoscape({ @@ -152,6 +154,9 @@ SPDX-License-Identifier: AGPL-3.0-or-later wsUrl += '?' + wsParams.toString(); updateConnectionStatus('connecting'); + awaitingInitialSnapshot = true; + snapshotRequestPending = false; + minimumSnapshotRevision = Math.max(minimumSnapshotRevision, graphRevision); ws = new WebSocket(wsUrl); ws.onopen = function() { @@ -165,6 +170,7 @@ SPDX-License-Identifier: AGPL-3.0-or-later }; ws.onclose = function() { + snapshotRequestPending = false; updateConnectionStatus('disconnected'); scheduleReconnect(); }; @@ -188,13 +194,24 @@ SPDX-License-Identifier: AGPL-3.0-or-later case 'delta': handleDelta(msg.data); break; + case 'snapshot_required': + requireSnapshot(msg.data.revision); + break; case 'pong': break; } } function handleSnapshot(data) { - graphVersion = data.version; + snapshotRequestPending = false; + if (!Number.isSafeInteger(data.revision) || data.revision < minimumSnapshotRevision || data.revision < graphRevision) { + requestRequiredSnapshot(); + return; + } + + graphRevision = data.revision; + minimumSnapshotRevision = graphRevision; + awaitingInitialSnapshot = false; resolvedGraphMode = data.mode || 'full'; graphMeta = { totalNodes: data.total_nodes || 0, @@ -238,12 +255,21 @@ SPDX-License-Identifier: AGPL-3.0-or-later } function handleDelta(data) { - if (resolvedGraphMode !== 'full' || prefersFilteredSnapshots()) { - scheduleFilteredRefresh(); + if (!Number.isSafeInteger(data.base_revision) || !Number.isSafeInteger(data.revision)) { + requireSnapshot(graphRevision + 1); + return; + } + if (data.revision <= graphRevision) { + return; + } + if (awaitingInitialSnapshot || requestedGraphMode !== 'full' || resolvedGraphMode !== 'full') { + requireSnapshot(data.revision); + return; + } + if (minimumSnapshotRevision > graphRevision || data.base_revision !== graphRevision) { + requireSnapshot(data.revision); return; } - - graphVersion = data.version; (data.added_nodes || []).forEach(node => { cy.add({ @@ -297,6 +323,8 @@ SPDX-License-Identifier: AGPL-3.0-or-later if (!graphMeta.largeGraph && (data.added_nodes || []).length > 2) { runLayout(); } + graphRevision = data.revision; + minimumSnapshotRevision = graphRevision; updateStats(); } @@ -344,10 +372,10 @@ SPDX-License-Identifier: AGPL-3.0-or-later const edges = cy.edges().length; if (resolvedGraphMode === 'filtered') { document.getElementById('stats').textContent = - 'Filtered: ' + nodes + '/' + graphMeta.totalNodes + ' nodes, ' + edges + '/' + graphMeta.totalEdges + ' edges (v' + graphVersion + ')'; + 'Filtered: ' + nodes + '/' + graphMeta.totalNodes + ' nodes, ' + edges + '/' + graphMeta.totalEdges + ' edges (r' + graphRevision + ')'; return; } - document.getElementById('stats').textContent = 'Full: ' + nodes + ' nodes, ' + edges + ' edges (v' + graphVersion + ')'; + document.getElementById('stats').textContent = 'Full: ' + nodes + ' nodes, ' + edges + ' edges (r' + graphRevision + ')'; } function updateConnectionStatus(status) { @@ -387,26 +415,28 @@ SPDX-License-Identifier: AGPL-3.0-or-later function requestSnapshot(mode) { if (!ws || ws.readyState !== WebSocket.OPEN) { - return; + return false; } ws.send(JSON.stringify({ type: 'snapshot_request', data: { mode: mode } })); + snapshotRequestPending = true; + return true; } - function scheduleFilteredRefresh() { - if (filteredRefreshTimer) { - return; + function requireSnapshot(revision) { + if (Number.isSafeInteger(revision)) { + minimumSnapshotRevision = Math.max(minimumSnapshotRevision, revision); } - filteredRefreshTimer = window.setTimeout(function() { - filteredRefreshTimer = null; - requestSnapshot(requestedGraphMode); - }, 250); + requestRequiredSnapshot(); } - function prefersFilteredSnapshots() { - return requestedGraphMode === 'filtered' || (requestedGraphMode === 'auto' && graphMeta.largeGraph); + function requestRequiredSnapshot() { + if (snapshotRequestPending) { + return; + } + requestSnapshot(requestedGraphMode); } function setGraphMode(mode) { diff --git a/internal/kitchen/graph.go b/internal/kitchen/graph.go index 7446e57..a610c29 100644 --- a/internal/kitchen/graph.go +++ b/internal/kitchen/graph.go @@ -34,7 +34,7 @@ func (h *Handler) handleGraph(w http.ResponseWriter, _ *http.Request) { func (h *Handler) handleGraphData(w http.ResponseWriter, r *http.Request) { writeGraphSecurityHeaders(w) p := h.Pantry() - snapshot := buildGraphSnapshot(p, p.Revision(), r.URL.Query().Get("mode")) + snapshot := buildGraphSnapshot(p, r.URL.Query().Get("mode")) data := graphData{ Mode: snapshot.Mode, LargeGraph: snapshot.LargeGraph, diff --git a/internal/kitchen/graph_hub.go b/internal/kitchen/graph_hub.go index a0c0177..b55da13 100644 --- a/internal/kitchen/graph_hub.go +++ b/internal/kitchen/graph_hub.go @@ -8,7 +8,6 @@ import ( "log/slog" "net/http" "sync" - "time" "github.com/coder/websocket" "github.com/coder/websocket/wsjson" @@ -16,32 +15,22 @@ import ( "github.com/boostsecurityio/smokedmeat/internal/pantry" ) -const ( - graphBatchWindow = 100 * time.Millisecond - graphSendBuffer = 256 - graphStaleVersions = 100 -) +const graphSendBuffer = 256 // GraphHub manages WebSocket connections for real-time graph updates. -// Implements pantry.Observer to receive change notifications. type GraphHub struct { mu sync.RWMutex clients map[*GraphClient]bool pantry *pantry.Pantry - - // Delta batching - deltaMu sync.Mutex - pendingDelta *GraphDelta - batchTimer *time.Timer } // GraphClient represents a connected graph visualization client. type GraphClient struct { - conn *websocket.Conn - send chan GraphMessage - hub *GraphHub - version uint64 - mode string + conn *websocket.Conn + send chan GraphMessage + hub *GraphHub + cancel context.CancelFunc + mode string } // NewGraphHub creates a new graph hub. @@ -63,113 +52,129 @@ func (h *GraphHub) HandleWebSocket(w http.ResponseWriter, r *http.Request) { return } + ctx, cancel := context.WithCancel(r.Context()) client := &GraphClient{ - conn: conn, - send: make(chan GraphMessage, graphSendBuffer), - hub: h, - mode: normalizeGraphMode(r.URL.Query().Get("mode")), + conn: conn, + send: make(chan GraphMessage, graphSendBuffer), + hub: h, + cancel: cancel, + mode: normalizeGraphMode(r.URL.Query().Get("mode")), + } + if !h.register(client) { + cancel() + _ = conn.Close(websocket.StatusPolicyViolation, "state queue unavailable") + return } - - h.register(client) - defer h.unregister(client) - - // Send initial snapshot - snapshot := h.buildSnapshot(client.mode) - client.version = snapshot.Version - client.send <- GraphMessage{Type: "snapshot", Data: snapshot} - - ctx, cancel := context.WithCancel(r.Context()) defer cancel() + defer h.unregister(client) go client.writePump(ctx) client.readPump(ctx) } -func (h *GraphHub) register(client *GraphClient) { +func (h *GraphHub) register(client *GraphClient) bool { h.mu.Lock() defer h.mu.Unlock() + h.clients[client] = true + snapshot := h.buildSnapshot(client.mode) + if !h.enqueueLocked(client, GraphMessage{Type: "snapshot", Data: snapshot}) { + return false + } slog.Debug("graph client connected", "total", len(h.clients)) + return true } func (h *GraphHub) unregister(client *GraphClient) { h.mu.Lock() defer h.mu.Unlock() - if _, ok := h.clients[client]; ok { - delete(h.clients, client) - close(client.send) - } + h.disconnectLocked(client) slog.Debug("graph client disconnected", "total", len(h.clients)) } -func (h *GraphHub) buildSnapshot(mode string) GraphSnapshot { - return buildGraphSnapshot(h.pantry, h.pantry.Revision(), mode) +func (h *GraphHub) disconnectLocked(client *GraphClient) { + if _, ok := h.clients[client]; !ok { + return + } + delete(h.clients, client) + close(client.send) + if client.cancel != nil { + client.cancel() + } } -// broadcast sends a message to all connected clients. -func (h *GraphHub) broadcast(msg GraphMessage) { - h.mu.RLock() - defer h.mu.RUnlock() - - for client := range h.clients { - select { - case client.send <- msg: - default: - slog.Warn("graph client buffer full, dropping message") - } +func (h *GraphHub) enqueueLocked(client *GraphClient, message GraphMessage) bool { + select { + case client.send <- message: + return true + default: + slog.Warn("graph client state queue full, disconnecting") + h.disconnectLocked(client) + return false } } -// flushDelta sends accumulated changes to all clients. -func (h *GraphHub) flushDelta() { - h.deltaMu.Lock() - defer h.deltaMu.Unlock() - - delta := h.pendingDelta - h.pendingDelta = nil - h.batchTimer = nil - - if delta == nil { - return +func (h *GraphHub) enqueue(client *GraphClient, message GraphMessage) bool { + h.mu.Lock() + defer h.mu.Unlock() + if _, ok := h.clients[client]; !ok { + return false } - - // Keep the publication fence held so a replacement snapshot cannot overtake this older delta. - h.broadcast(GraphMessage{Type: "delta", Data: delta}) + return h.enqueueLocked(client, message) } -// scheduleDeltaFlush ensures a delta is flushed after the batch window. -func (h *GraphHub) scheduleDeltaFlush() { - if h.batchTimer == nil { - h.batchTimer = time.AfterFunc(graphBatchWindow, h.flushDelta) - } +func (h *GraphHub) buildSnapshot(mode string) GraphSnapshot { + return buildGraphSnapshot(h.pantry, mode) } func (h *GraphHub) OnPantryChange(change pantry.ChangeSet) { - if change.Kind == pantry.ChangeCommittedState { - h.deltaMu.Lock() - if h.batchTimer != nil { - h.batchTimer.Stop() + h.mu.Lock() + defer h.mu.Unlock() + + var view *pantry.Pantry + var delta GraphDelta + hasDelta := false + snapshots := make(map[string]GraphSnapshot, 2) + for client := range h.clients { + if client.mode == graphModeFull { + message := GraphMessage{ + Type: "snapshot_required", + Data: GraphSnapshotRequired{Revision: change.Revision}, + } + if change.Kind == pantry.ChangeGranular { + if !hasDelta { + delta = graphDeltaFromChange(change) + hasDelta = true + } + message = GraphMessage{Type: "delta", Data: delta} + } + h.enqueueLocked(client, message) + continue } - h.batchTimer = nil - h.pendingDelta = nil - h.deltaMu.Unlock() - h.broadcastSnapshots() - return - } - h.deltaMu.Lock() - defer h.deltaMu.Unlock() + if view == nil { + view = h.pantry.Clone() + } + snapshot, ok := snapshots[client.mode] + if !ok { + snapshot = buildGraphSnapshotFromView(view, client.mode) + snapshots[client.mode] = snapshot + } + h.enqueueLocked(client, GraphMessage{Type: "snapshot", Data: snapshot}) + } +} - if h.pendingDelta == nil { - h.pendingDelta = &GraphDelta{} +func graphDeltaFromChange(change pantry.ChangeSet) GraphDelta { + delta := GraphDelta{ + BaseRevision: change.BaseRevision, + Revision: change.Revision, } - h.pendingDelta.Version = change.Revision for _, asset := range change.Granular.AddedAssets { - h.pendingDelta.AddedNodes = append(h.pendingDelta.AddedNodes, AssetToGraphNode(asset)) + delta.AddedNodes = append(delta.AddedNodes, AssetToGraphNode(asset)) } for _, update := range change.Granular.UpdatedAssets { node := AssetToGraphNode(update.After) - h.pendingDelta.UpdatedNodes = append(h.pendingDelta.UpdatedNodes, NodeUpdate{ + delta.UpdatedNodes = append(delta.UpdatedNodes, NodeUpdate{ ID: update.After.ID, OldState: string(update.Before.State), NewState: string(update.After.State), @@ -179,30 +184,28 @@ func (h *GraphHub) OnPantryChange(change pantry.ChangeSet) { }) } for _, edge := range change.Granular.AddedRelationships { - h.pendingDelta.AddedEdges = append(h.pendingDelta.AddedEdges, GraphEdge{ + delta.AddedEdges = append(delta.AddedEdges, GraphEdge{ Source: edge.From, Target: edge.To, Type: string(edge.Relationship.Type), }) } - h.pendingDelta.RemovedNodes = append(h.pendingDelta.RemovedNodes, change.Granular.RemovedAssetIDs...) + delta.RemovedNodes = append(delta.RemovedNodes, change.Granular.RemovedAssetIDs...) for _, edge := range change.Granular.RemovedRelationships { - h.pendingDelta.RemovedEdges = append(h.pendingDelta.RemovedEdges, EdgeRef{Source: edge.From, Target: edge.To}) + delta.RemovedEdges = append(delta.RemovedEdges, EdgeRef{Source: edge.From, Target: edge.To}) } - h.scheduleDeltaFlush() + return delta } -func (h *GraphHub) broadcastSnapshots() { - h.mu.RLock() - defer h.mu.RUnlock() - for client := range h.clients { - snapshot := h.buildSnapshot(client.mode) - select { - case client.send <- GraphMessage{Type: "snapshot", Data: snapshot}: - default: - slog.Warn("graph client buffer full, dropping snapshot") - } +func (h *GraphHub) sendSnapshot(client *GraphClient, mode string) { + h.mu.Lock() + defer h.mu.Unlock() + if _, ok := h.clients[client]; !ok { + return } + client.mode = normalizeGraphMode(mode) + snapshot := h.buildSnapshot(client.mode) + h.enqueueLocked(client, GraphMessage{Type: "snapshot", Data: snapshot}) } // ClientCount returns the number of connected graph clients. @@ -217,9 +220,8 @@ func (c *GraphClient) readPump(ctx context.Context) { for { var msg GraphMessage - err := wsjson.Read(ctx, c.conn, &msg) - if err != nil { - if websocket.CloseStatus(err) != websocket.StatusNormalClosure { + if err := wsjson.Read(ctx, c.conn, &msg); err != nil { + if websocket.CloseStatus(err) != websocket.StatusNormalClosure && ctx.Err() == nil { slog.Debug("graph websocket read error", "error", err) } return @@ -227,12 +229,9 @@ func (c *GraphClient) readPump(ctx context.Context) { switch msg.Type { case "ping": - c.send <- GraphMessage{Type: "pong"} + c.hub.enqueue(c, GraphMessage{Type: "pong"}) case "snapshot_request": - c.mode = graphModeFromData(msg.Data, c.mode) - snapshot := c.hub.buildSnapshot(c.mode) - c.version = snapshot.Version - c.send <- GraphMessage{Type: "snapshot", Data: snapshot} + c.hub.sendSnapshot(c, graphModeFromData(msg.Data, c.mode)) } } } @@ -248,6 +247,7 @@ func (c *GraphClient) writePump(ctx context.Context) { } if err := wsjson.Write(ctx, c.conn, msg); err != nil { slog.Debug("graph websocket write error", "error", err) + c.hub.unregister(c) return } } diff --git a/internal/kitchen/graph_hub_test.go b/internal/kitchen/graph_hub_test.go index 71d19ca..c299e2f 100644 --- a/internal/kitchen/graph_hub_test.go +++ b/internal/kitchen/graph_hub_test.go @@ -5,6 +5,7 @@ package kitchen import ( "context" + "fmt" "testing" "time" @@ -18,54 +19,199 @@ type graphHubSnapshotStore struct{} func (graphHubSnapshotStore) Replace([]byte) error { return nil } -func TestGraphHubTranslatesCommittedGranularChangeSet(t *testing.T) { +func TestGraphHubPublishesContiguousFullModeDeltas(t *testing.T) { live := pantry.New() hub := NewGraphHub(live) - client := &GraphClient{send: make(chan GraphMessage, 2), hub: hub, mode: graphModeFull} - hub.register(client) - t.Cleanup(func() { hub.unregister(client) }) + client := registerGraphClient(t, hub, graphModeFull, 4) + state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) + + for _, name := range []string{"acme", "globex"} { + require.NoError(t, state.Update(context.Background(), func(candidate *pantry.Pantry) error { + return candidate.AddAsset(pantry.NewOrganization(name, "github")) + })) + } + + first := receiveGraphMessage(t, client) + assert.Equal(t, "delta", first.Type) + firstDelta, ok := first.Data.(GraphDelta) + require.True(t, ok) + assert.Equal(t, uint64(0), firstDelta.BaseRevision) + assert.Equal(t, uint64(1), firstDelta.Revision) + require.Len(t, firstDelta.AddedNodes, 1) + assert.Equal(t, "github:org:acme", firstDelta.AddedNodes[0].ID) + + second := receiveGraphMessage(t, client) + assert.Equal(t, "delta", second.Type) + secondDelta, ok := second.Data.(GraphDelta) + require.True(t, ok) + assert.Equal(t, uint64(1), secondDelta.BaseRevision) + assert.Equal(t, uint64(2), secondDelta.Revision) + require.Len(t, secondDelta.AddedNodes, 1) + assert.Equal(t, "github:org:globex", secondDelta.AddedNodes[0].ID) +} + +func TestGraphHubPublishesSnapshotForProjectedModesAfterEveryChange(t *testing.T) { + for _, mode := range []string{graphModeFiltered, graphModeAuto} { + t.Run(mode, func(t *testing.T) { + live := pantry.New() + hub := NewGraphHub(live) + client := registerGraphClient(t, hub, mode, 2) + state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) + + require.NoError(t, state.Update(context.Background(), func(candidate *pantry.Pantry) error { + return candidate.AddAsset(pantry.NewOrganization("acme", "github")) + })) + + message := receiveGraphMessage(t, client) + assert.Equal(t, "snapshot", message.Type) + snapshot, ok := message.Data.(GraphSnapshot) + require.True(t, ok) + assert.Equal(t, uint64(1), snapshot.Revision) + assert.Equal(t, 1, snapshot.TotalNodes) + }) + } +} + +func TestGraphHubPublishesSnapshotRequiredForCommittedState(t *testing.T) { + live := pantry.New() + hub := NewGraphHub(live) + client := registerGraphClient(t, hub, graphModeFull, 3) state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) require.NoError(t, state.Update(context.Background(), func(candidate *pantry.Pantry) error { return candidate.AddAsset(pantry.NewOrganization("acme", "github")) })) - hub.flushDelta() + require.NoError(t, state.Replace(context.Background(), func(candidate *pantry.Pantry) error { + return candidate.AddAsset(pantry.NewOrganization("globex", "github")) + })) - message := <-client.send - assert.Equal(t, "delta", message.Type) - delta, ok := message.Data.(*GraphDelta) + deltaMessage := receiveGraphMessage(t, client) + assert.Equal(t, "delta", deltaMessage.Type) + + fenceMessage := receiveGraphMessage(t, client) + assert.Equal(t, "snapshot_required", fenceMessage.Type) + fence, ok := fenceMessage.Data.(GraphSnapshotRequired) require.True(t, ok) - assert.Equal(t, uint64(1), delta.Version) - require.Len(t, delta.AddedNodes, 1) - assert.Equal(t, "github:org:acme", delta.AddedNodes[0].ID) + assert.Equal(t, uint64(2), fence.Revision) } -func TestGraphHubCommittedStateMarkerSupersedesPendingDelta(t *testing.T) { +func TestGraphHubPublishesSnapshotForProjectedModeAfterCommittedState(t *testing.T) { live := pantry.New() hub := NewGraphHub(live) - client := &GraphClient{send: make(chan GraphMessage, 2), hub: hub, mode: graphModeFull} - hub.register(client) - t.Cleanup(func() { hub.unregister(client) }) + client := registerGraphClient(t, hub, graphModeFiltered, 2) + state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) + + require.NoError(t, state.Replace(context.Background(), func(candidate *pantry.Pantry) error { + return candidate.AddAsset(pantry.NewOrganization("acme", "github")) + })) + + message := receiveGraphMessage(t, client) + assert.Equal(t, "snapshot", message.Type) + snapshot, ok := message.Data.(GraphSnapshot) + require.True(t, ok) + assert.Equal(t, uint64(1), snapshot.Revision) + assert.Equal(t, 1, snapshot.TotalNodes) +} + +func TestGraphHubDisconnectsSaturatedClientWithoutAffectingHealthyClients(t *testing.T) { + live := pantry.New() + hub := NewGraphHub(live) + healthy := registerGraphClient(t, hub, graphModeFull, 2) + saturated := &GraphClient{send: make(chan GraphMessage, 1), hub: hub, mode: graphModeFull} + require.True(t, hub.register(saturated)) + t.Cleanup(func() { hub.unregister(saturated) }) state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) require.NoError(t, state.Update(context.Background(), func(candidate *pantry.Pantry) error { return candidate.AddAsset(pantry.NewOrganization("acme", "github")) })) - require.NoError(t, state.Replace(context.Background(), func(candidate *pantry.Pantry) error { - return candidate.AddAsset(pantry.NewOrganization("globex", "github")) + + assert.Equal(t, 1, hub.ClientCount()) + message := receiveGraphMessage(t, healthy) + assert.Equal(t, "delta", message.Type) + + initial, ok := <-saturated.send + require.True(t, ok) + assert.Equal(t, "snapshot", initial.Type) + _, ok = <-saturated.send + assert.False(t, ok) +} + +func TestGraphHubReconnectStartsWithCurrentSnapshot(t *testing.T) { + live := pantry.New() + hub := NewGraphHub(live) + state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) + + require.NoError(t, state.Update(context.Background(), func(candidate *pantry.Pantry) error { + return candidate.AddAsset(pantry.NewOrganization("acme", "github")) })) - message := <-client.send + client := &GraphClient{send: make(chan GraphMessage, 1), hub: hub, mode: graphModeFull} + require.True(t, hub.register(client)) + t.Cleanup(func() { hub.unregister(client) }) + + message := receiveGraphMessage(t, client) assert.Equal(t, "snapshot", message.Type) snapshot, ok := message.Data.(GraphSnapshot) require.True(t, ok) - assert.Equal(t, uint64(2), snapshot.Version) - assert.Equal(t, 2, snapshot.TotalNodes) + assert.Equal(t, uint64(1), snapshot.Revision) + assert.Equal(t, 1, snapshot.TotalNodes) +} + +func TestGraphHubSnapshotRevisionMatchesConcurrentState(t *testing.T) { + live := pantry.New() + hub := NewGraphHub(live) + state := pantry.NewCommittedState(live, graphHubSnapshotStore{}) + + const writes = 200 + done := make(chan struct{}) + var writerErr error + go func() { + defer close(done) + for i := range writes { + if err := state.Update(context.Background(), func(candidate *pantry.Pantry) error { + return candidate.AddAsset(pantry.NewOrganization(fmt.Sprintf("org-%03d", i), "github")) + }); err != nil { + writerErr = err + return + } + } + }() + + for { + snapshot := hub.buildSnapshot(graphModeFull) + assert.Equal(t, snapshot.Revision, uint64(snapshot.TotalNodes)) + select { + case <-done: + require.NoError(t, writerErr) + final := hub.buildSnapshot(graphModeFull) + assert.Equal(t, uint64(writes), final.Revision) + assert.Equal(t, writes, final.TotalNodes) + return + default: + } + } +} + +func registerGraphClient(t *testing.T, hub *GraphHub, mode string, capacity int) *GraphClient { + t.Helper() + client := &GraphClient{send: make(chan GraphMessage, capacity), hub: hub, mode: mode} + require.True(t, hub.register(client)) + t.Cleanup(func() { hub.unregister(client) }) + + initial := receiveGraphMessage(t, client) + require.Equal(t, "snapshot", initial.Type) + return client +} - time.Sleep(2 * graphBatchWindow) +func receiveGraphMessage(t *testing.T, client *GraphClient) GraphMessage { + t.Helper() select { - case unexpected := <-client.send: - t.Fatalf("unexpected state after committed snapshot: %#v", unexpected) - default: + case message, ok := <-client.send: + require.True(t, ok) + return message + case <-time.After(time.Second): + t.Fatal("timed out waiting for graph message") + return GraphMessage{} } } diff --git a/internal/kitchen/graph_ws.go b/internal/kitchen/graph_ws.go index a4b62b0..ed25969 100644 --- a/internal/kitchen/graph_ws.go +++ b/internal/kitchen/graph_ws.go @@ -25,13 +25,13 @@ const ( // GraphMessage is the envelope for all graph WebSocket messages. type GraphMessage struct { - Type string `json:"type"` // "snapshot", "delta", "ping", "pong" + Type string `json:"type"` Data any `json:"data,omitempty"` } -// GraphSnapshot is the initial full graph state sent on connect. +// GraphSnapshot is one complete graph projection at a committed revision. type GraphSnapshot struct { - Version uint64 `json:"version"` + Revision uint64 `json:"revision"` Mode string `json:"mode"` LargeGraph bool `json:"large_graph"` TotalNodes int `json:"total_nodes"` @@ -43,7 +43,8 @@ type GraphSnapshot struct { // GraphDelta contains incremental changes to the graph. type GraphDelta struct { - Version uint64 `json:"version"` + BaseRevision uint64 `json:"base_revision"` + Revision uint64 `json:"revision"` AddedNodes []GraphNode `json:"added_nodes,omitempty"` UpdatedNodes []NodeUpdate `json:"updated_nodes,omitempty"` AddedEdges []GraphEdge `json:"added_edges,omitempty"` @@ -51,6 +52,11 @@ type GraphDelta struct { RemovedEdges []EdgeRef `json:"removed_edges,omitempty"` } +// GraphSnapshotRequired fences deltas until the client loads a qualifying snapshot. +type GraphSnapshotRequired struct { + Revision uint64 `json:"revision"` +} + // EdgeRef identifies an edge by its source and target. type EdgeRef struct { Source string `json:"source"` @@ -189,10 +195,14 @@ func buildGraphSelection(p *pantry.Pantry, requestedMode string) graphSelection } } -func buildGraphSnapshot(p *pantry.Pantry, revision uint64, requestedMode string) GraphSnapshot { - selection := buildGraphSelection(p, requestedMode) +func buildGraphSnapshot(p *pantry.Pantry, requestedMode string) GraphSnapshot { + return buildGraphSnapshotFromView(p.Clone(), requestedMode) +} + +func buildGraphSnapshotFromView(view *pantry.Pantry, requestedMode string) GraphSnapshot { + selection := buildGraphSelection(view, requestedMode) return GraphSnapshot{ - Version: revision, + Revision: view.Revision(), Mode: selection.mode, LargeGraph: selection.largeGraph, TotalNodes: selection.totalNodes, diff --git a/internal/kitchen/graph_ws_test.go b/internal/kitchen/graph_ws_test.go index eeac66c..37d20a5 100644 --- a/internal/kitchen/graph_ws_test.go +++ b/internal/kitchen/graph_ws_test.go @@ -4,6 +4,7 @@ package kitchen import ( + "encoding/json" "fmt" "net/http" "net/http/httptest" @@ -149,6 +150,29 @@ func TestGraphScriptUsesTooltipProperties(t *testing.T) { assert.False(t, strings.Contains(body, "String(v)")) } +func TestGraphDeltaUsesRevisionFenceFields(t *testing.T) { + payload, err := json.Marshal(GraphDelta{BaseRevision: 4, Revision: 5}) + require.NoError(t, err) + + assert.JSONEq(t, `{"base_revision":4,"revision":5}`, string(payload)) + assert.NotContains(t, string(payload), "version") +} + +func TestGraphSnapshotUsesCommittedRevision(t *testing.T) { + payload, err := json.Marshal(GraphSnapshot{Revision: 5}) + require.NoError(t, err) + + assert.Contains(t, string(payload), `"revision":5`) + assert.NotContains(t, string(payload), "version") +} + +func TestGraphSnapshotRequiredCarriesMinimumRevision(t *testing.T) { + payload, err := json.Marshal(GraphSnapshotRequired{Revision: 5}) + require.NoError(t, err) + + assert.JSONEq(t, `{"revision":5}`, string(payload)) +} + func TestHandleGraphSetsBrowserSecurityHeaders(t *testing.T) { h := NewHandlerWithPublisher(nil, nil) rec := httptest.NewRecorder() @@ -256,8 +280,13 @@ func TestGraphTemplateAndScriptContainGraphModeControls(t *testing.T) { assert.Contains(t, page, "data-mode=\"filtered\"") assert.Contains(t, page, "data-mode=\"full\"") assert.Contains(t, script, "setGraphMode") - assert.Contains(t, script, "prefersFilteredSnapshots") - assert.Contains(t, script, "resolvedGraphMode !== 'full' || prefersFilteredSnapshots()") + assert.Contains(t, script, "case 'snapshot_required':") + assert.Contains(t, script, "data.base_revision !== graphRevision") + assert.Contains(t, script, "data.revision <= graphRevision") + assert.Contains(t, script, "requireSnapshot(data.revision)") + assert.Contains(t, script, "if (snapshotRequestPending)") + assert.Contains(t, script, "data.revision < minimumSnapshotRevision") + assert.NotContains(t, script, "scheduleFilteredRefresh") } func TestAssetToGraphNode_RepositorySSHAccessLabel(t *testing.T) { diff --git a/internal/pantry/graph.go b/internal/pantry/graph.go index 3e5153c..5710417 100644 --- a/internal/pantry/graph.go +++ b/internal/pantry/graph.go @@ -532,6 +532,14 @@ func (p *Pantry) Snapshot() Snapshot { return snapshot } +// Clone returns an independently readable Pantry captured from one coherent state. +func (p *Pantry) Clone() *Pantry { + if p == nil { + return New() + } + return pantryFromSnapshot(p.Snapshot()) +} + func pantryFromSnapshot(snapshot Snapshot) *Pantry { p := New() p.revision = snapshot.Revision