Merge branch 'main' into fix/errgroup-conditional-reassignment

mpyw · mpyw · commit 0e6eafc7453e · 2025-12-26T12:35:42.000+09:00
Resolved conflict in internal/probe/assignment.go: - Kept FuncLitsAssignedTo from main - Removed duplicate FuncLitAssignmentsOfIdent (already added in main via PR #24)
diff --git a/internal/checkers/goroutine.go b/internal/checkers/goroutine.go
@@ -61,11 +61,11 @@ func (*Goroutine) checkFromAST(cctx *probe.Context, stmt *ast.GoStmt) bool {
 	}
 
 	if ident, ok := call.Fun.(*ast.Ident); ok {
-		funcLit := cctx.FuncLitOfIdent(ident)
-		if funcLit == nil {
+		assigns := cctx.FuncLitAssignmentsOfIdent(ident)
+		if len(assigns) == 0 {
 			return true
 		}
-		return cctx.FuncLitCapturesContext(funcLit)
+		return cctx.FuncLitsAllCaptureContext(assigns)
 	}
 
 	if sel, ok := call.Fun.(*ast.SelectorExpr); ok {
@@ -166,16 +166,35 @@ func (c *GoroutineDerive) checkFromSSA(cctx *probe.Context, lit *ast.FuncLit) (*
 }
 
 func (c *GoroutineDerive) checkIdent(cctx *probe.Context, ident *ast.Ident) bool {
-	funcLit := cctx.FuncLitOfIdent(ident)
-	if funcLit == nil {
+	assigns := cctx.FuncLitAssignmentsOfIdent(ident)
+	if len(assigns) == 0 {
 		return true
 	}
 
-	if cctx.FuncLitHasContextParam(funcLit) {
-		return true
+	// 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
+	for i := startIdx; i < len(assigns); i++ {
+		lit := assigns[i].Lit
+		if !cctx.FuncLitHasContextParam(lit) && !c.derivers.SatisfiesAnyGroup(cctx.Pass, lit.Body) {
+			return false
+		}
 	}
 
-	return c.derivers.SatisfiesAnyGroup(cctx.Pass, funcLit.Body)
+	return true
 }
 
 func (c *GoroutineDerive) checkHigherOrder(cctx *probe.Context, innerCall *ast.CallExpr) bool {
@@ -192,14 +211,34 @@ func (c *GoroutineDerive) checkHigherOrder(cctx *probe.Context, innerCall *ast.C
 	}
 
 	if ident, ok := fun.(*ast.Ident); ok {
-		funcLit := cctx.FuncLitOfIdent(ident)
-		if funcLit == nil {
+		assigns := cctx.FuncLitAssignmentsOfIdent(ident)
+		if len(assigns) == 0 {
 			return true
 		}
-		if cctx.FuncLitHasContextParam(funcLit) {
-			return true
+
+		// 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
 		}
-		return c.factoryReturnsCallingFunc(cctx, funcLit)
+
+		// Check all assignments from startIdx onwards
+		for i := startIdx; i < len(assigns); i++ {
+			lit := assigns[i].Lit
+			if !cctx.FuncLitHasContextParam(lit) && !c.factoryReturnsCallingFunc(cctx, lit) {
+				return false
+			}
+		}
+		return true
 	}
 
 	return true
@@ -254,13 +293,32 @@ func (c *GoroutineDerive) returnedValueCalls(cctx *probe.Context, result ast.Exp
 		return false
 	}
 
-	innerFuncLit := cctx.FuncLitOfIdent(ident)
-	if innerFuncLit == nil {
+	assigns := cctx.FuncLitAssignmentsOfIdent(ident)
+	if len(assigns) == 0 {
 		return false
 	}
 
-	if cctx.FuncLitHasContextParam(innerFuncLit) {
-		return true
+	// 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
+		}
 	}
-	return c.derivers.SatisfiesAnyGroup(cctx.Pass, innerFuncLit.Body)
+
+	// Determine the starting point for checks
+	startIdx := 0
+	if lastUnconditionalIdx >= 0 {
+		startIdx = lastUnconditionalIdx
+	}
+
+	// Check all assignments from startIdx onwards
+	for i := startIdx; i < len(assigns); i++ {
+		lit := assigns[i].Lit
+		if !cctx.FuncLitHasContextParam(lit) && !c.derivers.SatisfiesAnyGroup(cctx.Pass, lit.Body) {
+			return false
+		}
+	}
+	return true
 }
diff --git a/internal/probe/assignment.go b/internal/probe/assignment.go
@@ -24,6 +24,26 @@ func (c *Context) FuncLitOfIdent(ident *ast.Ident) *ast.FuncLit {
 	return c.FuncLitAssignedTo(v, token.NoPos)
 }
 
+// FuncLitsOfIdent returns ALL func literals assigned to the identifier's variable.
+// This is needed for conditional reassignment patterns where different branches
+// assign different closures to the same variable.
+func (c *Context) FuncLitsOfIdent(ident *ast.Ident) []*ast.FuncLit {
+	v := c.VarOf(ident)
+	if v == nil {
+		return nil
+	}
+	return c.FuncLitsAssignedTo(v, token.NoPos)
+}
+
+// FuncLitAssignmentsOfIdent returns ALL func literal assignments with conditionality info.
+func (c *Context) FuncLitAssignmentsOfIdent(ident *ast.Ident) []FuncLitAssignment {
+	v := c.VarOf(ident)
+	if v == nil {
+		return nil
+	}
+	return c.FuncLitAssignmentsTo(v, token.NoPos)
+}
+
 // FuncLitAssignedTo searches for the func literal assigned to the variable.
 // If beforePos is token.NoPos, returns the LAST assignment found.
 // If beforePos is set, returns the last assignment BEFORE that position.
@@ -51,13 +71,32 @@ func (c *Context) FuncLitAssignedTo(v *types.Var, beforePos token.Pos) *ast.Func
 	return result
 }
 
-// FuncLitAssignmentsOfIdent returns ALL func literal assignments with conditionality info.
-func (c *Context) FuncLitAssignmentsOfIdent(ident *ast.Ident) []FuncLitAssignment {
-	v := c.VarOf(ident)
-	if v == nil {
+// FuncLitsAssignedTo searches for ALL func literals assigned to the variable.
+// If beforePos is token.NoPos, returns ALL assignments found.
+// If beforePos is set, returns all assignments BEFORE that position.
+// This is needed for conditional reassignment patterns.
+func (c *Context) FuncLitsAssignedTo(v *types.Var, beforePos token.Pos) []*ast.FuncLit {
+	f := c.FileOf(v.Pos())
+	if f == nil {
 		return nil
 	}
-	return c.FuncLitAssignmentsTo(v, token.NoPos)
+
+	var results []*ast.FuncLit
+	ast.Inspect(f, func(n ast.Node) bool {
+		assign, ok := n.(*ast.AssignStmt)
+		if !ok {
+			return true
+		}
+		if beforePos != token.NoPos && assign.Pos() >= beforePos {
+			return true
+		}
+		if fl := c.funcLitInAssignment(assign, v); fl != nil {
+			results = append(results, fl)
+		}
+		return true
+	})
+
+	return results
 }
 
 // FuncLitAssignmentsTo searches for ALL func literal assignments with conditionality info.
diff --git a/internal/probe/capture.go b/internal/probe/capture.go
@@ -53,6 +53,41 @@ func (c *Context) FuncLitCapturesContext(lit *ast.FuncLit) bool {
 	return c.FuncLitHasContextParam(lit) || c.FuncLitUsesContext(lit)
 }
 
+// FuncLitsAllCaptureContext checks if func literals properly capture context.
+// Uses conditionality information to determine the correct check:
+// - Find the last unconditional assignment
+// - Check all assignments from that point onwards (including conditional ones)
+// - ALL must capture context for the check to pass
+func (c *Context) FuncLitsAllCaptureContext(assigns []FuncLitAssignment) bool {
+	if len(assigns) == 0 {
+		return true
+	}
+
+	// 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 capture context (because conditional assignments may override)
+	for i := startIdx; i < len(assigns); i++ {
+		if !c.FuncLitCapturesContext(assigns[i].Lit) {
+			return false
+		}
+	}
+	return true
+}
+
 // FuncLitUsesContext checks if a function literal references any context variable.
 // Does NOT descend into nested func literals.
 func (c *Context) FuncLitUsesContext(lit *ast.FuncLit) bool {
diff --git a/testdata/metatest/tests/conditionalReassignmentAllUseCtx.json b/testdata/metatest/tests/conditionalReassignmentAllUseCtx.json
@@ -0,0 +1,16 @@
+{
+  "title": "Conditional reassignment - all paths use ctx",
+  "targets": [
+    "goroutine"
+  ],
+  "variants": {
+    "bad": null,
+    "good": {
+      "description": "All assignments use context, so all paths propagate it.",
+      "functions": {
+        "goroutine": "goodConditionalReassignAllUseCtx"
+      }
+    }
+  },
+  "level": "evil"
+}
diff --git a/testdata/metatest/tests/conditionalReassignmentFirstNoCtx.json b/testdata/metatest/tests/conditionalReassignmentFirstNoCtx.json
@@ -0,0 +1,16 @@
+{
+  "title": "Conditional reassignment - first doesn't use ctx, conditional does",
+  "targets": [
+    "goroutine"
+  ],
+  "variants": {
+    "bad": {
+      "description": "First assignment ignores context, conditional reassignment uses it.",
+      "functions": {
+        "goroutine": "badConditionalReassignFirstNoCtx"
+      }
+    },
+    "good": null
+  },
+  "level": "evil"
+}
diff --git a/testdata/metatest/tests/conditionalReassignmentFirstUsesCtx.json b/testdata/metatest/tests/conditionalReassignmentFirstUsesCtx.json
@@ -0,0 +1,16 @@
+{
+  "title": "Conditional reassignment - first uses ctx, conditional doesn't",
+  "targets": [
+    "goroutine"
+  ],
+  "variants": {
+    "bad": {
+      "description": "First assignment uses context, but conditional reassignment doesn't.",
+      "functions": {
+        "goroutine": "badConditionalReassignFirstUsesCtx"
+      }
+    },
+    "good": null
+  },
+  "level": "evil"
+}
diff --git a/testdata/src/goroutine/evil.go b/testdata/src/goroutine/evil.go
@@ -711,3 +711,47 @@ func packageFactoryWithCtxParam(ctx context.Context) func() {
 func packageFactoryWithoutCtx() func() {
 	return func() { fmt.Println("no ctx") }
 }
+
+// ===== CONDITIONAL REASSIGNMENT PATTERNS =====
+// These test the analyzer's handling of variable reassignment in conditionals.
+// All paths must propagate context - if ANY assignment doesn't use ctx, warn.
+
+//vt:helper
+var conditionFlag bool
+
+// [BAD]: Conditional reassignment - first uses ctx, conditional doesn't
+//
+// First assignment uses context, but conditional reassignment doesn't.
+// If condition is true, goroutine won't propagate context.
+// All paths must be safe, so this should warn.
+func badConditionalReassignFirstUsesCtx(ctx context.Context) {
+	fn := func() { _ = ctx } // First assignment uses ctx
+	if conditionFlag {
+		fn = func() { fmt.Println("no ctx") } // Conditional reassignment without ctx
+	}
+	go fn() // want `goroutine does not propagate context "ctx"`
+}
+
+// [BAD]: Conditional reassignment - first doesn't use ctx, conditional does
+//
+// First assignment ignores context, conditional reassignment uses it.
+// If condition is false, goroutine won't propagate context.
+// All paths must be safe, so this should warn.
+func badConditionalReassignFirstNoCtx(ctx context.Context) {
+	fn := func() { fmt.Println("no ctx") } // First assignment without ctx
+	if conditionFlag {
+		fn = func() { _ = ctx } // Conditional reassignment with ctx
+	}
+	go fn() // want `goroutine does not propagate context "ctx"`
+}
+
+// [GOOD]: Conditional reassignment - all paths use ctx
+//
+// All assignments use context, so all paths propagate it.
+func goodConditionalReassignAllUseCtx(ctx context.Context) {
+	fn := func() { _ = ctx } // First assignment uses ctx
+	if conditionFlag {
+		fn = func() { _ = ctx } // Conditional reassignment also uses ctx
+	}
+	go fn() // All paths safe
+}
