Merge pull request #23 from mpyw/fix/errgroup-conditional-reassignment

fix: track conditional reassignments in errgroup/spawner checkers

Author: mpyw

internal/checkers/spawner.go

@@ -150,11 +150,11 @@ func (c *SpawnCallbackChecker) checkArgFromAST(cctx *probe.Context, arg ast.Expr
 	}
 
 	if ident, ok := arg.(*ast.Ident); ok {
-		funcLit := cctx.FuncLitOfIdent(ident)
-		if funcLit == nil {
+		assigns := cctx.FuncLitAssignmentsOfIdent(ident)
+		if len(assigns) == 0 {
 			return true
 		}
-		return c.checkFuncLitAST(cctx, funcLit)
+		return c.checkFuncLitAssignments(cctx, assigns)
 	}
 
 	if call, ok := arg.(*ast.CallExpr); ok {
@@ -172,6 +172,34 @@ func (c *SpawnCallbackChecker) checkArgFromAST(cctx *probe.Context, arg ast.Expr
 	return true
 }
 
+// checkFuncLitAssignments checks all func literal assignments from last unconditional onwards.
+// ALL must pass for the check to succeed.
+func (c *SpawnCallbackChecker) checkFuncLitAssignments(cctx *probe.Context, assigns []probe.FuncLitAssignment) bool {
+	// Find the index of the last unconditional assignment
+	lastUnconditionalIdx := -1
+	for i := len(assigns) - 1; i >= 0; i-- {
+		if !assigns[i].Conditional {
+			lastUnconditionalIdx = i
+			break
+		}
+	}
+
+	// Determine the starting point for checks
+	startIdx := 0
+	if lastUnconditionalIdx >= 0 {
+		startIdx = lastUnconditionalIdx
+	}
+
+	// Check all assignments from startIdx onwards
+	// ALL must pass (because conditional assignments may override)
+	for i := startIdx; i < len(assigns); i++ {
+		if !c.checkFuncLitAST(cctx, assigns[i].Lit) {
+			return false
+		}
+	}
+	return true
+}
+
 // checkFuncLitAST checks a func literal using AST-based analysis.
 func (c *SpawnCallbackChecker) checkFuncLitAST(cctx *probe.Context, lit *ast.FuncLit) bool {
 	// Check context capture
@@ -333,11 +361,11 @@ func (c *SpawnerChecker) checkFuncArg(cctx *probe.Context, arg ast.Expr) bool {
 	}
 
 	if ident, ok := arg.(*ast.Ident); ok {
-		funcLit := cctx.FuncLitOfIdent(ident)
-		if funcLit == nil {
+		assigns := cctx.FuncLitAssignmentsOfIdent(ident)
+		if len(assigns) == 0 {
 			return true
 		}
-		return c.checkFuncLitAST(cctx, funcLit)
+		return c.checkFuncLitAssignments(cctx, assigns)
 	}
 
 	if call, ok := arg.(*ast.CallExpr); ok {
@@ -347,6 +375,34 @@ func (c *SpawnerChecker) checkFuncArg(cctx *probe.Context, arg ast.Expr) bool {
 	return true
 }
 
+// checkFuncLitAssignments checks all func literal assignments from last unconditional onwards.
+// ALL must pass for the check to succeed.
+func (c *SpawnerChecker) checkFuncLitAssignments(cctx *probe.Context, assigns []probe.FuncLitAssignment) bool {
+	// Find the index of the last unconditional assignment
+	lastUnconditionalIdx := -1
+	for i := len(assigns) - 1; i >= 0; i-- {
+		if !assigns[i].Conditional {
+			lastUnconditionalIdx = i
+			break
+		}
+	}
+
+	// Determine the starting point for checks
+	startIdx := 0
+	if lastUnconditionalIdx >= 0 {
+		startIdx = lastUnconditionalIdx
+	}
+
+	// Check all assignments from startIdx onwards
+	// ALL must pass (because conditional assignments may override)
+	for i := startIdx; i < len(assigns); i++ {
+		if !c.checkFuncLitAST(cctx, assigns[i].Lit) {
+			return false
+		}
+	}
+	return true
+}
+
 // checkFuncLitSSA checks a func literal using SSA analysis for SpawnerChecker.
 func (c *SpawnerChecker) checkFuncLitSSA(cctx *probe.Context, lit *ast.FuncLit) (bool, bool) {
 	if cctx.SSAProg == nil || cctx.Tracer == nil {

testdata/metatest/tests/errgroupConditionalReassignmentAllUseCtx.json

@@ -0,0 +1,16 @@
+{
+  "title": "Conditional reassignment - all paths use ctx",
+  "targets": [
+    "errgroup"
+  ],
+  "variants": {
+    "bad": null,
+    "good": {
+      "description": "All assignments use context, so all paths propagate it.",
+      "functions": {
+        "errgroup": "goodConditionalReassignAllUseCtx"
+      }
+    }
+  },
+  "level": "evil"
+}

testdata/metatest/tests/errgroupConditionalReassignmentFirstNoCtx.json

@@ -0,0 +1,16 @@
+{
+  "title": "Conditional reassignment - first doesn't use ctx, conditional does",
+  "targets": [
+    "errgroup"
+  ],
+  "variants": {
+    "bad": {
+      "description": "First assignment ignores context, conditional reassignment uses it.",
+      "functions": {
+        "errgroup": "badConditionalReassignFirstNoCtx"
+      }
+    },
+    "good": null
+  },
+  "level": "evil"
+}

testdata/metatest/tests/errgroupConditionalReassignmentFirstUsesCtx.json

@@ -0,0 +1,16 @@
+{
+  "title": "Conditional reassignment - first uses ctx, conditional doesn't",
+  "targets": [
+    "errgroup"
+  ],
+  "variants": {
+    "bad": {
+      "description": "First assignment uses context, but conditional reassignment doesn't.",
+      "functions": {
+        "errgroup": "badConditionalReassignFirstUsesCtx"
+      }
+    },
+    "good": null
+  },
+  "level": "evil"
+}

testdata/src/errgroup/evil.go

@@ -843,3 +843,78 @@ func badNestedFactoryWithoutCtx(ctx context.Context) {
 	g.Go(makeFactory()()) // want `errgroup.Group.Go\(\) closure should use context "ctx"`
 	_ = g.Wait()
 }
+
+// ===== CONDITIONAL REASSIGNMENT PATTERNS =====
+// These patterns test conditional reassignment of function variables.
+// ALL assignments from last unconditional onwards must capture context.
+
+var conditionFlag bool
+
+// [BAD]: Conditional reassignment - first uses ctx, conditional doesn't
+//
+// First assignment uses context, but conditional reassignment doesn't.
+//
+// See also:
+//   goroutine: badConditionalReassignFirstUsesCtx
+//   waitgroup: badConditionalReassignFirstUsesCtx
+func badConditionalReassignFirstUsesCtx(ctx context.Context) {
+	g := new(errgroup.Group)
+	fn := func() error {
+		_ = ctx
+		return nil
+	}
+	if conditionFlag {
+		fn = func() error {
+			fmt.Println("no ctx")
+			return nil
+		}
+	}
+	g.Go(fn) // want `errgroup.Group.Go\(\) closure should use context "ctx"`
+	_ = g.Wait()
+}
+
+// [BAD]: Conditional reassignment - first doesn't use ctx, conditional does
+//
+// First assignment ignores context, conditional reassignment uses it.
+//
+// See also:
+//   goroutine: badConditionalReassignFirstNoCtx
+//   waitgroup: badConditionalReassignFirstNoCtx
+func badConditionalReassignFirstNoCtx(ctx context.Context) {
+	g := new(errgroup.Group)
+	fn := func() error {
+		fmt.Println("no ctx")
+		return nil
+	}
+	if conditionFlag {
+		fn = func() error {
+			_ = ctx
+			return nil
+		}
+	}
+	g.Go(fn) // want `errgroup.Group.Go\(\) closure should use context "ctx"`
+	_ = g.Wait()
+}
+
+// [GOOD]: Conditional reassignment - all paths use ctx
+//
+// All assignments use context, so all paths propagate it.
+//
+// See also:
+//   goroutine: goodConditionalReassignAllUseCtx
+//   waitgroup: goodConditionalReassignAllUseCtx
+func goodConditionalReassignAllUseCtx(ctx context.Context) {
+	g := new(errgroup.Group)
+	fn := func() error {
+		_ = ctx
+		return nil
+	}
+	if conditionFlag {
+		fn = func() error {
+			_ = ctx
+			return nil
+		}
+	}
+	g.Go(fn) // OK - all assignments use ctx
+	_ = g.Wait()
+}