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

Author: mpyw

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
 }

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 {

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.

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 {

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"
+}

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"
+}

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"
+}