fix: prevent taint analysis hang on packages with many CHA call graph edges (#1608) (#1610)

Cap incoming call graph edges to 32 per function and add cross-query
parameter taint memoization to prevent combinatorial explosion when CHA
over-approximates interface method calls across transitive dependencies.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: ccojocar

taint/analyzer_internal_test.go

@@ -1,6 +1,7 @@
 package taint
 
 import (
+	"fmt"
 	"go/ast"
 	"go/constant"
 	"go/parser"
@@ -9,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/buildssa"
@@ -441,6 +443,109 @@ func f(w W)                               { w.Write([]byte("hello")) }
 	_ = analyzer.Analyze(prog, []*ssa.Function{fn})
 }
 
+// buildManySinkCallsFixture creates an SSA program with many interface implementations
+// and many sink-calling functions, producing a large CHA call graph. Used by both the
+// regression test and benchmark.
+func buildManySinkCallsFixture(tb testing.TB) (*ssa.Program, []*ssa.Function) {
+	tb.Helper()
+
+	src := `package p
+
+type W interface{ Write([]byte) (int, error) }
+`
+	// Generate 20 concrete implementations of W to inflate CHA edges.
+	for i := 0; i < 20; i++ {
+		src += fmt.Sprintf(`
+type Impl%d struct{}
+func (x *Impl%d) Write(p []byte) (int, error) { return len(p), nil }
+`, i, i)
+	}
+
+	// Generate 20 functions, each calling w.Write with a variable arg (potential sink).
+	for i := 0; i < 20; i++ {
+		src += fmt.Sprintf(`
+func caller%d(w W, data []byte) { w.Write(data) }
+`, i)
+	}
+
+	fset := token.NewFileSet()
+	parsed, err := parser.ParseFile(fset, "p.go", src, 0)
+	if err != nil {
+		tb.Fatalf("parse: %v", err)
+	}
+	info := &types.Info{
+		Types:      make(map[ast.Expr]types.TypeAndValue),
+		Defs:       make(map[*ast.Ident]types.Object),
+		Uses:       make(map[*ast.Ident]types.Object),
+		Implicits:  make(map[ast.Node]types.Object),
+		Scopes:     make(map[ast.Node]*types.Scope),
+		Selections: make(map[*ast.SelectorExpr]*types.Selection),
+	}
+	pkg, err := (&types.Config{}).Check("p", fset, []*ast.File{parsed}, info)
+	if err != nil {
+		tb.Fatalf("type-check: %v", err)
+	}
+	prog := ssa.NewProgram(fset, ssa.BuilderMode(0))
+	ssaPkg := prog.CreatePackage(pkg, []*ast.File{parsed}, info, true)
+	prog.Build()
+
+	// Collect all caller* functions as analysis targets.
+	var srcFuncs []*ssa.Function
+	for i := 0; i < 20; i++ {
+		fn := ssaPkg.Func(fmt.Sprintf("caller%d", i))
+		if fn == nil {
+			tb.Fatalf("SSA function caller%d not found", i)
+		}
+		srcFuncs = append(srcFuncs, fn)
+	}
+
+	return prog, srcFuncs
+}
+
+func TestTaintAnalysisPerformanceWithManySinkCalls(t *testing.T) {
+	t.Parallel()
+
+	// This test verifies that taint analysis completes in bounded time even when
+	// CHA produces a large call graph (many interface implementations × many sink calls).
+	// Before the maxCallerEdges cap and paramTaintCache, this scenario could hang.
+	prog, srcFuncs := buildManySinkCallsFixture(t)
+
+	analyzer := New(&Config{
+		Sinks: []Sink{
+			{Package: "p", Receiver: "W", Method: "Write", CheckArgs: []int{1}},
+		},
+	})
+
+	// Must complete within 10 seconds; without the fix this could hang indefinitely.
+	done := make(chan []Result, 1)
+	go func() {
+		done <- analyzer.Analyze(prog, srcFuncs)
+	}()
+
+	select {
+	case <-time.After(10 * time.Second):
+		t.Fatal("taint analysis did not complete within 10 seconds — possible hang regression")
+	case results := <-done:
+		_ = results
+	}
+}
+
+func BenchmarkTaintAnalysisManySinkCalls(b *testing.B) {
+	prog, srcFuncs := buildManySinkCallsFixture(b)
+
+	cfg := &Config{
+		Sinks: []Sink{
+			{Package: "p", Receiver: "W", Method: "Write", CheckArgs: []int{1}},
+		},
+	}
+
+	b.ResetTimer()
+	for b.Loop() {
+		analyzer := New(cfg)
+		analyzer.Analyze(prog, srcFuncs)
+	}
+}
+
 func TestResolveOriginalTypeMakeInterface(t *testing.T) {
 	t.Parallel()
 	// Build a minimal, self-contained SSA program (no external imports) that

taint/taint.go

@@ -23,6 +23,12 @@ import (
 // maxTaintDepth limits recursion depth to prevent stack overflow on large codebases
 const maxTaintDepth = 50
 
+// maxCallerEdges caps the number of incoming call graph edges examined per function
+// in isParameterTainted. CHA over-approximates call graphs (every interface method
+// call fans out to ALL implementations), so a function can have thousands of callers.
+// Real taint flows come from direct/nearby callers, not the 33rd+ CHA-generated edge.
+const maxCallerEdges = 32
+
 // isContextType checks if a type is context.Context.
 // context.Context is a control-flow mechanism (deadlines, cancellation, request-scoped values)
 // that does not carry user-controlled data relevant to taint sinks like XSS.
@@ -211,14 +217,21 @@ type Config struct {
 }
 
 // Analyzer performs taint analysis on SSA programs.
+// paramKey identifies a specific parameter of a function for memoization.
+type paramKey struct {
+	fn       *ssa.Function
+	paramIdx int
+}
+
 type Analyzer struct {
-	config     *Config
-	sources    map[string]Source   // keyed by full type string
-	funcSrcs   map[string]Source   // function sources keyed by "pkg.Func"
-	sinks      map[string]Sink     // keyed by full function string
-	sanitizers map[string]struct{} // keyed by full function string
-	callGraph  *callgraph.Graph
-	prog       *ssa.Program // set at Analyze time for ArgTypeGuards resolution
+	config          *Config
+	sources         map[string]Source   // keyed by full type string
+	funcSrcs        map[string]Source   // function sources keyed by "pkg.Func"
+	sinks           map[string]Sink     // keyed by full function string
+	sanitizers      map[string]struct{} // keyed by full function string
+	callGraph       *callgraph.Graph
+	prog            *ssa.Program      // set at Analyze time for ArgTypeGuards resolution
+	paramTaintCache map[paramKey]bool // caches true results from isParameterTainted
 }
 
 // SetCallGraph injects a precomputed call graph.
@@ -309,13 +322,17 @@ func (a *Analyzer) Analyze(prog *ssa.Program, srcFuncs []*ssa.Function) []Result
 		a.callGraph = cha.CallGraph(prog)
 	}
 
+	a.paramTaintCache = make(map[paramKey]bool)
+
 	var results []Result
 
 	// Find all sink calls in the program
 	for _, fn := range srcFuncs {
 		results = append(results, a.analyzeFunctionSinks(fn)...)
 	}
 
+	a.paramTaintCache = nil
+
 	return results
 }
 
@@ -824,11 +841,31 @@ func (a *Analyzer) isParameterTainted(param *ssa.Parameter, fn *ssa.Function, vi
 		return false
 	}
 
+	// Resolve paramIdx early so we can use it for cache lookups.
+	paramIdx := -1
+	for i, p := range fn.Params {
+		if p == param {
+			paramIdx = i
+			break
+		}
+	}
+
+	// Check memoization cache (only true results are cached).
+	if paramIdx >= 0 && a.paramTaintCache != nil {
+		key := paramKey{fn: fn, paramIdx: paramIdx}
+		if a.paramTaintCache[key] {
+			return true
+		}
+	}
+
 	// Check if parameter type is a source type.
 	// This is the ONLY place where type-based source matching should trigger
 	// automatic taint — because parameters represent data flowing IN from
 	// external callers we don't control.
 	if a.isSourceType(param.Type()) {
+		if paramIdx >= 0 && a.paramTaintCache != nil {
+			a.paramTaintCache[paramKey{fn: fn, paramIdx: paramIdx}] = true
+		}
 		return true
 	}
 
@@ -842,14 +879,6 @@ func (a *Analyzer) isParameterTainted(param *ssa.Parameter, fn *ssa.Function, vi
 		return false
 	}
 
-	paramIdx := -1
-	for i, p := range fn.Params {
-		if p == param {
-			paramIdx = i
-			break
-		}
-	}
-
 	if paramIdx < 0 {
 		return false
 	}
@@ -867,8 +896,14 @@ func (a *Analyzer) isParameterTainted(param *ssa.Parameter, fn *ssa.Function, vi
 		adjustedIdx = paramIdx
 	}
 
-	// Check each caller
+	// Check each caller, capping at maxCallerEdges to avoid combinatorial
+	// explosion from CHA over-approximation of interface method calls.
+	edgesChecked := 0
 	for _, inEdge := range node.In {
+		if edgesChecked >= maxCallerEdges {
+			break
+		}
+
 		site := inEdge.Site
 		if site == nil {
 			continue
@@ -877,7 +912,11 @@ func (a *Analyzer) isParameterTainted(param *ssa.Parameter, fn *ssa.Function, vi
 		callArgs := site.Common().Args
 
 		if adjustedIdx < len(callArgs) {
+			edgesChecked++
 			if a.isTainted(callArgs[adjustedIdx], inEdge.Caller.Func, visited, depth+1) {
+				if a.paramTaintCache != nil {
+					a.paramTaintCache[paramKey{fn: fn, paramIdx: paramIdx}] = true
+				}
 				return true
 			}
 		}