fix(graphql): memoize fragment depth to bound depth-limit walk

Author: JiahuiWho

internal/serve/graphql/depth_limit.go

@@ -50,7 +50,7 @@ func (d DepthLimit) MutateOperationContext(_ context.Context, opCtx *graphql.Ope
 		return nil
 	}
 
-	depth := selectionSetDepth(op.SelectionSet, opCtx.Doc.Fragments, map[string]bool{}, 0)
+	depth := selectionSetDepth(op.SelectionSet, opCtx.Doc.Fragments)
 	if depth > maxDepth {
 		err := gqlerror.Errorf("operation has depth %d, which exceeds the limit of %d", depth, maxDepth)
 		err.Extensions = map[string]interface{}{"code": errQueryTooDeep}
@@ -59,35 +59,57 @@ func (d DepthLimit) MutateOperationContext(_ context.Context, opCtx *graphql.Ope
 	return nil
 }
 
-// selectionSetDepth returns the maximum nesting depth reachable from set, starting at depth.
+// selectionSetDepth returns the maximum nesting depth reachable from set.
 // Every Field adds one level (whether or not it has children — a leaf scalar still counts as the
 // level it was selected at); InlineFragments are transparent (they don't add a level, matching
 // how they read in the query); FragmentSpreads are resolved against fragments and otherwise
 // treated the same as an InlineFragment.
+func selectionSetDepth(set ast.SelectionSet, fragments ast.FragmentDefinitionList) int {
+	return selectionSetDepthMemo(set, fragments, map[string]bool{}, map[string]int{}, 0)
+}
+
+// selectionSetDepthMemo is the recursive core of selectionSetDepth, threading two maps through the
+// walk: visitedFragments and fragmentDepths.
 //
 // visitedFragments tracks fragment names currently on the recursion stack, guarding against a
 // fragment that (directly or transitively) spreads itself: rather than recursing forever, a
 // repeated name is simply skipped.
-func selectionSetDepth(set ast.SelectionSet, fragments ast.FragmentDefinitionList, visitedFragments map[string]bool, depth int) int {
+//
+// A fragment spread adds no level of its own, so the depth a fragment contributes below its spread
+// point is independent of where it is spread (it's a pure additive offset onto the current depth).
+// fragmentDepths caches each fragment's depth, computed once relative to depth 0, so on any later
+// encounter we simply add the current depth. Without this cache, a fragment reachable through
+// multiple spreads is re-expanded on every path to it, so a document whose fragments fan out into
+// one another can force work that grows exponentially with the number of fragments even though the
+// document itself is small and its actual depth is shallow. Memoizing bounds the walk to
+// O(number of fragments + selections).
+func selectionSetDepthMemo(set ast.SelectionSet, fragments ast.FragmentDefinitionList, visitedFragments map[string]bool, fragmentDepths map[string]int, depth int) int {
 	maxDepth := depth
 	for _, sel := range set {
 		var childDepth int
 		switch sel := sel.(type) {
 		case *ast.Field:
-			childDepth = selectionSetDepth(sel.SelectionSet, fragments, visitedFragments, depth+1)
+			childDepth = selectionSetDepthMemo(sel.SelectionSet, fragments, visitedFragments, fragmentDepths, depth+1)
 		case *ast.InlineFragment:
-			childDepth = selectionSetDepth(sel.SelectionSet, fragments, visitedFragments, depth)
+			childDepth = selectionSetDepthMemo(sel.SelectionSet, fragments, visitedFragments, fragmentDepths, depth)
 		case *ast.FragmentSpread:
 			if visitedFragments[sel.Name] {
 				continue
 			}
-			def := fragments.ForName(sel.Name)
-			if def == nil {
-				continue
+			// Reuse the fragment's cached relative depth if we've expanded it before; otherwise
+			// expand it once (from depth 0, guarded against cycles) and cache the result.
+			relDepth, cached := fragmentDepths[sel.Name]
+			if !cached {
+				def := fragments.ForName(sel.Name)
+				if def == nil {
+					continue
+				}
+				visitedFragments[sel.Name] = true
+				relDepth = selectionSetDepthMemo(def.SelectionSet, fragments, visitedFragments, fragmentDepths, 0)
+				delete(visitedFragments, sel.Name)
+				fragmentDepths[sel.Name] = relDepth
 			}
-			visitedFragments[sel.Name] = true
-			childDepth = selectionSetDepth(def.SelectionSet, fragments, visitedFragments, depth)
-			delete(visitedFragments, sel.Name)
+			childDepth = depth + relDepth
 		default:
 			continue
 		}

internal/serve/graphql/depth_limit_test.go

@@ -2,6 +2,7 @@ package graphql
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -19,7 +20,7 @@ func field(name string, sub ast.SelectionSet) *ast.Field {
 func TestSelectionSetDepth(t *testing.T) {
 	t.Run("flat selection set is depth 1", func(t *testing.T) {
 		set := ast.SelectionSet{field("hash", nil), field("ledgerNumber", nil)}
-		assert.Equal(t, 1, selectionSetDepth(set, nil, map[string]bool{}, 0))
+		assert.Equal(t, 1, selectionSetDepth(set, nil))
 	})
 
 	t.Run("nested fields add one level each", func(t *testing.T) {
@@ -35,7 +36,7 @@ func TestSelectionSetDepth(t *testing.T) {
 				}),
 			}),
 		}
-		assert.Equal(t, 5, selectionSetDepth(set, nil, map[string]bool{}, 0))
+		assert.Equal(t, 5, selectionSetDepth(set, nil))
 	})
 
 	t.Run("inline fragment does not add a level", func(t *testing.T) {
@@ -49,7 +50,7 @@ func TestSelectionSetDepth(t *testing.T) {
 			}),
 		}
 		// node=1, inline fragment transparent, balance=2
-		assert.Equal(t, 2, selectionSetDepth(set, nil, map[string]bool{}, 0))
+		assert.Equal(t, 2, selectionSetDepth(set, nil))
 	})
 
 	t.Run("fragment spread is resolved against the document's fragments and adds no level itself", func(t *testing.T) {
@@ -70,7 +71,7 @@ func TestSelectionSetDepth(t *testing.T) {
 			}),
 		}
 		// transactionByHash=1, spread transparent, operations=2, id=3
-		assert.Equal(t, 3, selectionSetDepth(set, fragments, map[string]bool{}, 0))
+		assert.Equal(t, 3, selectionSetDepth(set, fragments))
 	})
 
 	t.Run("unresolvable fragment spread is skipped without panicking", func(t *testing.T) {
@@ -79,7 +80,7 @@ func TestSelectionSetDepth(t *testing.T) {
 				&ast.FragmentSpread{Name: "DoesNotExist"},
 			}),
 		}
-		assert.Equal(t, 1, selectionSetDepth(set, ast.FragmentDefinitionList{}, map[string]bool{}, 0))
+		assert.Equal(t, 1, selectionSetDepth(set, ast.FragmentDefinitionList{}))
 	})
 
 	t.Run("fragment cycle does not hang", func(t *testing.T) {
@@ -91,7 +92,7 @@ func TestSelectionSetDepth(t *testing.T) {
 
 		done := make(chan int, 1)
 		go func() {
-			done <- selectionSetDepth(set, fragments, map[string]bool{}, 0)
+			done <- selectionSetDepth(set, fragments)
 		}()
 
 		select {
@@ -102,16 +103,55 @@ func TestSelectionSetDepth(t *testing.T) {
 		}
 	})
 
-	t.Run("a fragment spread more than once (not a cycle) is still counted each time", func(t *testing.T) {
+	t.Run("a fragment spread more than once (not a cycle) is counted correctly once", func(t *testing.T) {
 		fragments := ast.FragmentDefinitionList{
 			{Name: "Shared", SelectionSet: ast.SelectionSet{field("value", nil)}},
 		}
 		set := ast.SelectionSet{
 			field("a", ast.SelectionSet{&ast.FragmentSpread{Name: "Shared"}}),
 			field("b", ast.SelectionSet{&ast.FragmentSpread{Name: "Shared"}}),
 		}
-		// a=1/b=1, Shared transparent, value=2
-		assert.Equal(t, 2, selectionSetDepth(set, fragments, map[string]bool{}, 0))
+		// a=1/b=1, Shared transparent, value=2. Depth is a max, so spreading Shared twice yields
+		// the same depth as once — the memo cache changes cost, not the result.
+		assert.Equal(t, 2, selectionSetDepth(set, fragments))
+	})
+
+	// A document whose fragments fan out into one another must stay cheap to measure. Here each
+	// fragment spreads the next twice, so without per-fragment memoization the walk re-expands the
+	// chain on every path and does O(2^N) work; at N=64 that would never finish. With memoization
+	// it is linear and returns immediately, while still reporting the true (shallow) depth of 2.
+	t.Run("fragment fan-out does not blow up the walk", func(t *testing.T) {
+		const n = 64
+		fragments := make(ast.FragmentDefinitionList, 0, n)
+		for i := 0; i < n-1; i++ {
+			fragments = append(fragments, &ast.FragmentDefinition{
+				Name: fmt.Sprintf("f%d", i),
+				SelectionSet: ast.SelectionSet{
+					&ast.FragmentSpread{Name: fmt.Sprintf("f%d", i+1)},
+					&ast.FragmentSpread{Name: fmt.Sprintf("f%d", i+1)},
+				},
+			})
+		}
+		fragments = append(fragments, &ast.FragmentDefinition{
+			Name:         fmt.Sprintf("f%d", n-1),
+			SelectionSet: ast.SelectionSet{field("address", nil)},
+		})
+		set := ast.SelectionSet{
+			field("accountByAddress", ast.SelectionSet{&ast.FragmentSpread{Name: "f0"}}),
+		}
+
+		done := make(chan int, 1)
+		go func() {
+			done <- selectionSetDepth(set, fragments)
+		}()
+
+		select {
+		case depth := <-done:
+			// accountByAddress=1, fragment chain transparent, address=2.
+			assert.Equal(t, 2, depth)
+		case <-time.After(5 * time.Second):
+			t.Fatal("selectionSetDepth did not terminate on a binary fan-out fragment bomb; memoization is not working and the DoS is present")
+		}
 	})
 }