refactor: Update getReportMessage function to handle unsupported nested context types

pkg/analyzer/analyzer.go

@@ -7,6 +7,7 @@ import (
 	"go/ast"
 	"go/printer"
 	"go/token"
+	"go/types"
 
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/inspect"
@@ -28,6 +29,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	nodeFilter := []ast.Node{
 		(*ast.ForStmt)(nil),
 		(*ast.RangeStmt)(nil),
+		(*ast.FuncLit)(nil),
 	}
 
 	inspctr.Preorder(nodeFilter, func(node ast.Node) {
@@ -65,7 +67,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 
 		pass.Report(analysis.Diagnostic{
 			Pos:            assignStmt.Pos(),
-			Message:        "nested context in loop",
+			Message:        getReportMessage(node),
 			SuggestedFixes: fixes,
 		})
 
@@ -74,6 +76,19 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	return nil, nil
 }
 
+func getReportMessage(node ast.Node) string {
+	switch node.(type) {
+	case *ast.ForStmt:
+		return "nested context in loop"
+	case *ast.RangeStmt:
+		return "nested context in loop"
+	case *ast.FuncLit:
+		return "nested context in function literal"
+	default:
+		return "unsupported nested context type"
+	}
+}
+
 func getBody(node ast.Node) (*ast.BlockStmt, error) {
 	forStmt, ok := node.(*ast.ForStmt)
 	if ok {
@@ -85,6 +100,11 @@ func getBody(node ast.Node) (*ast.BlockStmt, error) {
 		return rangeStmt.Body, nil
 	}
 
+	funcLit, ok := node.(*ast.FuncLit)
+	if ok {
+		return funcLit.Body, nil
+	}
+
 	return nil, errUnknown
 }
 
@@ -155,7 +175,7 @@ func findNestedContext(pass *analysis.Pass, block *ast.BlockStmt, stmts []ast.St
 		// allow assignment to non-pointer children of values defined within the loop
 		if lhs := getRootIdent(pass, assignStmt.Lhs[0]); lhs != nil {
 			if obj := pass.TypesInfo.ObjectOf(lhs); obj != nil {
-				if obj.Pos() >= block.Pos() && obj.Pos() < block.End() {
+				if checkObjectScopeWithinBlock(obj.Parent(), block) {
 					continue // definition is within the loop
 				}
 			}
@@ -167,6 +187,18 @@ func findNestedContext(pass *analysis.Pass, block *ast.BlockStmt, stmts []ast.St
 	return nil
 }
 
+func checkObjectScopeWithinBlock(scope *types.Scope, block *ast.BlockStmt) bool {
+	if scope == nil {
+		return false
+	}
+
+	if scope.Pos() >= block.Pos() && scope.End() <= block.End() {
+		return true
+	}
+
+	return false
+}
+
 func getRootIdent(pass *analysis.Pass, node ast.Node) *ast.Ident {
 	for {
 		switch n := node.(type) {

testdata/src/example.go

@@ -59,6 +59,12 @@ func example() {
 
 		break
 	}
+
+	// detects contexts wrapped in function literals (this is risky as function literals can be called multiple times)
+	_ = func() {
+		ctx = wrapContext(ctx) // want "nested context in function literal"
+	}
+
 }
 
 func wrapContext(ctx context.Context) context.Context {
@@ -180,3 +186,31 @@ func inVariousNestedBlocks(ctx context.Context) {
 		break
 	}
 }
+
+// this middleware could run on every request, bloating the request parameter level context and causing a memory leak
+func badMiddleware(ctx context.Context) func() error {
+	return func() error {
+		ctx = wrapContext(ctx) // want "nested context in function literal"
+		return doSomethingWithCtx(ctx)
+	}
+}
+
+// this middleware is fine, as it doesn't modify the context of parent function
+func okMiddleware(ctx context.Context) func() error {
+	return func() error {
+		ctx := wrapContext(ctx)
+		return doSomethingWithCtx(ctx)
+	}
+}
+
+// this middleware is fine, as it only modifies the context passed to it
+func okMiddleware2(ctx context.Context) func(ctx context.Context) error {
+	return func(ctx context.Context) error {
+		ctx = wrapContext(ctx)
+		return doSomethingWithCtx(ctx)
+	}
+}
+
+func doSomethingWithCtx(ctx context.Context) error {
+	return nil
+}