fix: context.Background and context.TODO properly mark the loop as non-fat

Author: Crocmagnon

.golangci.yml

@@ -6,6 +6,11 @@ linters:
   disable:
     - depguard
     - exhaustruct
+    - wsl
+    - noinlineerr
+  settings:
+    cyclop:
+      max-complexity: 20
 formatters:
   enable:
     - goimports

pkg/analyzer/analyzer.go

@@ -172,6 +172,9 @@ func getBody(node ast.Node) (*ast.BlockStmt, error) {
 }
 
 func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *ast.AssignStmt {
+	// Track which context variables have been reset to a known empty context
+	resetContexts := make(map[string]bool)
+
 	for _, stmt := range stmts {
 		// Recurse if necessary
 		stmtList := getStmtList(stmt)
@@ -198,8 +201,21 @@ func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *as
 			continue
 		}
 
-		// Ignore [context.Background] & [context.TODO].
+		// Get the variable name being assigned to
+		varName := getVarName(pass, assignStmt)
+
+		// If the assignment is to a known empty context, mark this variable as reset
 		if isContextFunction(assignStmt.Rhs[0], "Background", "TODO") {
+			if varName != "" {
+				resetContexts[varName] = true
+			}
+
+			continue
+		}
+
+		// If this variable was previously reset to a known empty context in this block,
+		// it's safe to modify it
+		if varName != "" && resetContexts[varName] {
 			continue
 		}
 
@@ -218,6 +234,21 @@ func findNestedContext(pass *analysis.Pass, node ast.Node, stmts []ast.Stmt) *as
 	return nil
 }
 
+func getVarName(pass *analysis.Pass, assignStmt *ast.AssignStmt) string {
+	varName := ""
+
+	if ident, ok := assignStmt.Lhs[0].(*ast.Ident); ok {
+		varName = ident.Name
+	} else if sel, ok := assignStmt.Lhs[0].(*ast.SelectorExpr); ok {
+		// For struct fields like tc.ctx
+		if rendered, err := render(pass.Fset, sel); err == nil {
+			varName = string(rendered)
+		}
+	}
+
+	return varName
+}
+
 func getStmtList(stmt ast.Stmt) []ast.Stmt {
 	switch typedStmt := stmt.(type) {
 	case *ast.BlockStmt:
@@ -240,7 +271,9 @@ func getStmtList(stmt ast.Stmt) []ast.Stmt {
 // render returns the pretty-print of the given node.
 func render(fset *token.FileSet, x interface{}) ([]byte, error) {
 	var buf bytes.Buffer
-	if err := printer.Fprint(&buf, fset, x); err != nil {
+
+	err := printer.Fprint(&buf, fset, x)
+	if err != nil {
 		return nil, fmt.Errorf("printing node: %w", err)
 	}
 

pkg/analyzer/analyzer_test.go

@@ -53,7 +53,7 @@ func TestAnalyzer(t *testing.T) {
 				}
 			}
 
-			analysistest.RunWithSuggestedFixes(t, analysistest.TestData(), anlzr, test.dir)
+			analysistest.Run(t, analysistest.TestData(), anlzr, test.dir)
 		})
 	}
 }
@@ -65,3 +65,11 @@ func TestAnalyzer_cgo(t *testing.T) {
 
 	analysistest.Run(t, analysistest.TestData(), a, "cgo")
 }
+
+func TestSuggestedFixes(t *testing.T) {
+	t.Parallel()
+
+	a := analyzer.NewAnalyzer()
+
+	analysistest.RunWithSuggestedFixes(t, analysistest.TestData(), a, "suggestedfix")
+}

pkg/analyzer/testdata/src/common/example.go

@@ -7,6 +7,7 @@ import (
 
 func example() {
 	ctx := context.Background()
+	ctx2 := context.Background()
 
 	for i := 0; i < 10; i++ {
 		ctx := context.WithValue(ctx, "key", i)
@@ -18,6 +19,15 @@ func example() {
 		ctx = context.WithValue(ctx, "other", "val")
 	}
 
+	for i := 0; i < 10; i++ {
+		ctx = context.Background()
+		// This is fine because the first action was to re-assign to a known empty context.
+		ctx = context.WithValue(ctx, "key", i)
+		ctx = context.WithValue(ctx, "other", "val")
+		// Not OK because this var has NOT been re-assigned to a known empty context.
+		ctx2 = context.WithValue(ctx2, "key", i) // want "nested context in loop"
+	}
+
 	for item := range []string{"one", "two", "three"} {
 		ctx = wrapContext(ctx) // want "nested context in loop"
 		ctx := context.WithValue(ctx, "key", item)
@@ -241,11 +251,29 @@ func testCasesInit(t *testing.T) {
 			ctx: context.WithValue(context.Background(), "key", "value"),
 		},
 	}
+
+	// This is fine because the first action is to re-assign to a known empty context.
 	for _, tc := range cases {
 		t.Run("some test", func(t *testing.T) {
 			if tc.ctx == nil {
 				tc.ctx = context.Background()
+				tc.ctx = context.WithValue(tc.ctx, "key", "value")
 			}
 		})
 	}
 }
+
+var (
+	globalCtx    context.Context
+	globalCancel context.CancelFunc
+)
+
+func BeforeSuite(_ func()) interface{} {
+	return nil
+}
+
+// This is fine because the first action is to re-assign to a known empty context.
+var _ = BeforeSuite(func() {
+	globalCtx = context.TODO()
+	globalCtx, globalCancel = context.WithCancel(globalCtx)
+})