From 3e9e29f41c6048145785a204b7a2b33ace9bdec0 Mon Sep 17 00:00:00 2001
From: Michael Urman <MichaelUrman@users.noreply.github.com>
Date: Wed, 29 May 2024 09:47:58 -0500
Subject: [PATCH] Test and fix cases with nested contexts

As long as the context is rooted in a non-pointer value that has a new
copy in the loop, it is as safe to copy that value as it is to copy the
context. So only report such cases when they are indirected and thus
shared.
---
 pkg/analyzer/analyzer.go | 27 +++++++++++++++++++++++++++
 testdata/src/example.go  | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/pkg/analyzer/analyzer.go b/pkg/analyzer/analyzer.go
index 7bb8fd4..40ce3b6 100644
--- a/pkg/analyzer/analyzer.go
+++ b/pkg/analyzer/analyzer.go
@@ -7,6 +7,7 @@ import (
 	"go/ast"
 	"go/printer"
 	"go/token"
+
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/inspect"
 	"golang.org/x/tools/go/ast/inspector"
@@ -54,6 +55,14 @@ func run(pass *analysis.Pass) (interface{}, error) {
 				break
 			}
 
+			if lhs := getRootIdent(pass, assignStmt.Lhs[0]); lhs != nil {
+				if obj := pass.TypesInfo.ObjectOf(lhs); obj != nil {
+					if obj.Pos() >= body.Pos() && obj.Pos() < body.End() {
+						continue
+					}
+				}
+			}
+
 			suggestedStmt := ast.AssignStmt{
 				Lhs:    assignStmt.Lhs,
 				TokPos: assignStmt.TokPos,
@@ -103,6 +112,24 @@ func getBody(node ast.Node) (*ast.BlockStmt, error) {
 	return nil, errUnknown
 }
 
+func getRootIdent(pass *analysis.Pass, node ast.Node) *ast.Ident {
+	for {
+		switch n := node.(type) {
+		case *ast.Ident:
+			return n
+		case *ast.IndexExpr:
+			node = n.X
+		case *ast.SelectorExpr:
+			if sel, ok := pass.TypesInfo.Selections[n]; ok && sel.Indirect() {
+				return nil // indirected (pointer) roots don't imply a (safe) copy
+			}
+			node = n.X
+		default:
+			return nil
+		}
+	}
+}
+
 // render returns the pretty-print of the given node
 func render(fset *token.FileSet, x interface{}) (string, error) {
 	var buf bytes.Buffer
diff --git a/testdata/src/example.go b/testdata/src/example.go
index 69777f7..a7e83a8 100644
--- a/testdata/src/example.go
+++ b/testdata/src/example.go
@@ -30,3 +30,38 @@ func example() {
 func wrapContext(ctx context.Context) context.Context {
 	return context.WithoutCancel(ctx)
 }
+
+// storing contexts in a struct isn't recommended, but local copies of a non-pointer struct should act like local copies of a context.
+func inStructs(ctx context.Context) {
+	for i := 0; i < 10; i++ {
+		c := struct{ Ctx context.Context }{ctx}
+		c.Ctx = context.WithValue(c.Ctx, "key", i)
+		c.Ctx = context.WithValue(c.Ctx, "other", "val")
+	}
+
+	for i := 0; i < 10; i++ {
+		c := []struct{ Ctx context.Context }{{ctx}}
+		c[0].Ctx = context.WithValue(c[0].Ctx, "key", i)
+		c[0].Ctx = context.WithValue(c[0].Ctx, "other", "val")
+	}
+
+	c := struct{ Ctx context.Context }{ctx}
+	for i := 0; i < 10; i++ {
+		c := c
+		c.Ctx = context.WithValue(c.Ctx, "key", i)
+		c.Ctx = context.WithValue(c.Ctx, "other", "val")
+	}
+
+	pc := &struct{ Ctx context.Context }{ctx}
+	for i := 0; i < 10; i++ {
+		c := pc
+		c.Ctx = context.WithValue(c.Ctx, "key", i) // want "nested context in loop"
+		c.Ctx = context.WithValue(c.Ctx, "other", "val")
+	}
+
+	for i := 0; i < 10; i++ {
+		c := []*struct{ Ctx context.Context }{{ctx}}
+		c[0].Ctx = context.WithValue(c[0].Ctx, "key", i) // want "nested context in loop"
+		c[0].Ctx = context.WithValue(c[0].Ctx, "other", "val")
+	}
+}