ctxhttp-analyzer/analyzer.go

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/printer"
	"go/token"
	"go/types"
	"strings"

	"golang.org/x/tools/go/analysis"
)

var Analyzer = &analysis.Analyzer{
	Name: "ctxhttp",
	Doc:  "detects HTTP calls that should use context and suggests fixes",
	Run:  run,
}

func run(pass *analysis.Pass) (interface{}, error) {
	for _, file := range pass.Files {
		// Track which functions have already had context injected,
		// so we don't duplicate the injection for multiple HTTP calls
		// in the same function.
		injected := map[*ast.FuncDecl]bool{}

		// Track whether we've already added the "context" import for this file.
		importAdded := false

		ast.Inspect(file, func(n ast.Node) bool {
			call, ok := n.(*ast.CallExpr)
			if !ok {
				return true
			}

			kind := detectHTTPCall(pass, call)
			if kind == "" {
				return true
			}

			enc := enclosingFunc(file, call)
			ctxExpr := contextExpr(enc, pass)

			// Should we inject a ctx := ... statement?
			inject := !injected[enc] && needsContextInjection(enc, pass)

			// Should we add the "context" import?
			addImport := !importAdded && !hasContextImport(file)

			var fixes []analysis.SuggestedFix
			switch kind {
			case "http.NewRequest":
				fixes = fixNewRequest(pass, file, call, ctxExpr, enc, inject, addImport)
			case "http.Get", "http.Head":
				fixes = fixSimpleHTTPFunc(pass, file, call, kind, ctxExpr, enc, inject, addImport, "http.DefaultClient")
			case "http.Post":
				fixes = fixHTTPPost(pass, file, call, ctxExpr, enc, inject, addImport, "http.DefaultClient")
			case "http.PostForm":
				fixes = fixHTTPPostForm(pass, file, call, ctxExpr, enc, inject, addImport, "http.DefaultClient")
			case "client.Get", "client.Head":
				receiver := nodeText(pass, call.Fun.(*ast.SelectorExpr).X)
				fixes = fixSimpleHTTPFunc(pass, file, call, kind, ctxExpr, enc, inject, addImport, receiver)
			case "client.Post":
				receiver := nodeText(pass, call.Fun.(*ast.SelectorExpr).X)
				fixes = fixHTTPPost(pass, file, call, ctxExpr, enc, inject, addImport, receiver)
			case "client.PostForm":
				receiver := nodeText(pass, call.Fun.(*ast.SelectorExpr).X)
				fixes = fixHTTPPostForm(pass, file, call, ctxExpr, enc, inject, addImport, receiver)
			}

			if inject {
				injected[enc] = true
			}
			if addImport {
				importAdded = true
			}

			pass.Report(analysis.Diagnostic{
				Pos:            call.Pos(),
				End:            call.End(),
				Message:        kind + " should use context",
				SuggestedFixes: fixes,
			})

			return true
		})
	}
	return nil, nil
}

// detectHTTPCall returns the kind of HTTP call, or "" if not relevant.
func detectHTTPCall(pass *analysis.Pass, call *ast.CallExpr) string {
	sel, ok := call.Fun.(*ast.SelectorExpr)
	if !ok {
		return ""
	}

	obj := pass.TypesInfo.ObjectOf(sel.Sel)
	if obj == nil {
		return ""
	}

	fn, ok := obj.(*types.Func)
	if !ok {
		return ""
	}

	switch fn.FullName() {
	case "net/http.NewRequest":
		return "http.NewRequest"
	case "net/http.Get":
		return "http.Get"
	case "net/http.Head":
		return "http.Head"
	case "net/http.Post":
		return "http.Post"
	case "net/http.PostForm":
		return "http.PostForm"
	case "(*net/http.Client).Get":
		return "client.Get"
	case "(*net/http.Client).Head":
		return "client.Head"
	case "(*net/http.Client).Post":
		return "client.Post"
	case "(*net/http.Client).PostForm":
		return "client.PostForm"
	}
	return ""
}

// enclosingFunc finds the FuncDecl containing a node.
func enclosingFunc(file *ast.File, target ast.Node) *ast.FuncDecl {
	var result *ast.FuncDecl
	ast.Inspect(file, func(n ast.Node) bool {
		fd, ok := n.(*ast.FuncDecl)
		if !ok {
			return true
		}
		if fd.Body != nil && fd.Pos() <= target.Pos() && target.End() <= fd.End() {
			result = fd
		}
		return true
	})
	return result
}

// enclosingStmt finds the statement in a block that directly contains the call.
// Returns the BlockStmt and the index of the statement within it.
func enclosingStmt(file *ast.File, call *ast.CallExpr) (*ast.BlockStmt, int) {
	var resultBlock *ast.BlockStmt
	var resultIdx int

	ast.Inspect(file, func(n ast.Node) bool {
		block, ok := n.(*ast.BlockStmt)
		if !ok {
			return true
		}
		for i, stmt := range block.List {
			if stmt.Pos() <= call.Pos() && call.End() <= stmt.End() {
				resultBlock = block
				resultIdx = i
			}
		}
		return true
	})
	return resultBlock, resultIdx
}

// contextExpr determines what expression provides context in the enclosing function.
func contextExpr(fd *ast.FuncDecl, pass *analysis.Pass) string {
	if fd == nil {
		return "context.TODO()"
	}
	if name := existingContextParam(fd, pass); name != "" {
		return name
	}
	return "ctx"
}

// existingContextParam returns the name of an existing context.Context parameter, or "".
func existingContextParam(fd *ast.FuncDecl, pass *analysis.Pass) string {
	if fd.Type.Params == nil {
		return ""
	}
	for _, field := range fd.Type.Params.List {
		typ := pass.TypesInfo.TypeOf(field.Type)
		if typ != nil && typ.String() == "context.Context" {
			if len(field.Names) > 0 {
				return field.Names[0].Name
			}
			return "ctx"
		}
	}
	return ""
}

// needsContextInjection returns true if the enclosing function needs a ctx := ... statement.
func needsContextInjection(fd *ast.FuncDecl, pass *analysis.Pass) bool {
	if fd == nil {
		return false
	}
	return existingContextParam(fd, pass) == ""
}

// contextInjectionText returns the statement to inject for context creation.
func contextInjectionText(fd *ast.FuncDecl, pass *analysis.Pass) string {
	if fd == nil {
		return ""
	}

	name := fd.Name.Name

	switch {
	case name == "main" || name == "init":
		return "ctx := context.Background()\n\t"
	case strings.HasPrefix(name, "Test") || strings.HasPrefix(name, "Benchmark"):
		tName := testParamName(fd)
		return "ctx := " + tName + ".Context()\n\t"
	case isHTTPHandler(fd, pass):
		reqName := httpRequestParamName(fd, pass)
		return "ctx := " + reqName + ".Context()\n\t"
	default:
		return "ctx := context.TODO()\n\t"
	}
}

// testParamName returns the name of the *testing.T or *testing.B parameter.
func testParamName(fd *ast.FuncDecl) string {
	if fd.Type.Params == nil {
		return "t"
	}
	for _, field := range fd.Type.Params.List {
		if len(field.Names) > 0 {
			return field.Names[0].Name
		}
	}
	return "t"
}

// isHTTPHandler checks if a function has the (http.ResponseWriter, *http.Request) signature.
func isHTTPHandler(fd *ast.FuncDecl, pass *analysis.Pass) bool {
	if fd.Type.Params == nil {
		return false
	}
	params := flattenParams(fd.Type.Params.List)
	if len(params) < 2 {
		return false
	}

	p0Type := pass.TypesInfo.TypeOf(params[len(params)-2].Type)
	p1Type := pass.TypesInfo.TypeOf(params[len(params)-1].Type)
	if p0Type == nil || p1Type == nil {
		return false
	}

	return p0Type.String() == "net/http.ResponseWriter" &&
		p1Type.String() == "*net/http.Request"
}

// httpRequestParamName returns the name of the *http.Request parameter.
func httpRequestParamName(fd *ast.FuncDecl, pass *analysis.Pass) string {
	for _, field := range flattenParams(fd.Type.Params.List) {
		typ := pass.TypesInfo.TypeOf(field.Type)
		if typ != nil && typ.String() == "*net/http.Request" {
			if len(field.Names) > 0 {
				return field.Names[0].Name
			}
		}
	}
	return "r"
}

// flattenParams expands grouped params (a, b int) into individual fields.
func flattenParams(fields []*ast.Field) []*ast.Field {
	var result []*ast.Field
	for _, f := range fields {
		if len(f.Names) <= 1 {
			result = append(result, f)
		} else {
			for _, name := range f.Names {
				result = append(result, &ast.Field{
					Names: []*ast.Ident{name},
					Type:  f.Type,
				})
			}
		}
	}
	return result
}

// fixNewRequest generates a SuggestedFix for http.NewRequest → http.NewRequestWithContext.
func fixNewRequest(pass *analysis.Pass, file *ast.File, call *ast.CallExpr, ctxExpr string, enc *ast.FuncDecl, inject bool, addImport bool) []analysis.SuggestedFix {
	sel := call.Fun.(*ast.SelectorExpr)

	var edits []analysis.TextEdit

	// Replace "NewRequest" with "NewRequestWithContext".
	edits = append(edits, analysis.TextEdit{
		Pos:     sel.Sel.Pos(),
		End:     sel.Sel.End(),
		NewText: []byte("NewRequestWithContext"),
	})

	// Insert ctx as first argument.
	if len(call.Args) > 0 {
		edits = append(edits, analysis.TextEdit{
			Pos:     call.Args[0].Pos(),
			End:     call.Args[0].Pos(),
			NewText: []byte(ctxExpr + ", "),
		})
	}

	if inject {
		edits = append(edits, contextInjectionEdit(enc, pass))
	}
	if addImport {
		edits = append(edits, contextImportEdit(file)...)
	}

	return []analysis.SuggestedFix{{
		Message:   "use http.NewRequestWithContext",
		TextEdits: edits,
	}}
}

// fixSimpleHTTPFunc rewrites http.Get(url) or http.Head(url) (and client equivalents).
//
// Before:
//
//	resp, err := http.Get(url)
//
// After:
//
//	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
//	if err != nil {
//		return ..., err  // or handle based on context
//	}
//	resp, err := http.DefaultClient.Do(req)
func fixSimpleHTTPFunc(pass *analysis.Pass, file *ast.File, call *ast.CallExpr, kind string, ctxExpr string, enc *ast.FuncDecl, inject bool, addImport bool, client string) []analysis.SuggestedFix {
	method := "GET"
	if strings.HasSuffix(kind, "Head") {
		method = "HEAD"
	}

	urlArg := nodeText(pass, call.Args[0])

	// Find what the result is assigned to.
	block, idx := enclosingStmt(file, call)
	if block == nil {
		return nil
	}
	stmt := block.List[idx]

	var lhs string
	assign, isAssign := stmt.(*ast.AssignStmt)
	if isAssign && len(assign.Lhs) == 2 {
		lhs = nodeText(pass, assign.Lhs[0])
	}
	if lhs == "" {
		lhs = "_"
	}

	tok := ":="
	if isAssign {
		tok = assign.Tok.String()
	}

	replacement := fmt.Sprintf(
		"req, err %s http.NewRequestWithContext(%s, %q, %s, nil)\n\tif err != nil {\n\t\treturn err\n\t}\n\t%s, err %s %s.Do(req)",
		tok, ctxExpr, method, urlArg,
		lhs, tok, client,
	)

	var edits []analysis.TextEdit
	edits = append(edits, analysis.TextEdit{
		Pos:     stmt.Pos(),
		End:     stmt.End(),
		NewText: []byte(replacement),
	})

	if inject {
		edits = append(edits, contextInjectionEdit(enc, pass))
	}
	if addImport {
		edits = append(edits, contextImportEdit(file)...)
	}

	return []analysis.SuggestedFix{{
		Message:   fmt.Sprintf("use http.NewRequestWithContext with %s.Do", client),
		TextEdits: edits,
	}}
}

// fixHTTPPost rewrites http.Post(url, contentType, body) and client.Post().
func fixHTTPPost(pass *analysis.Pass, file *ast.File, call *ast.CallExpr, ctxExpr string, enc *ast.FuncDecl, inject bool, addImport bool, client string) []analysis.SuggestedFix {
	urlArg := nodeText(pass, call.Args[0])
	ctypeArg := nodeText(pass, call.Args[1])
	bodyArg := nodeText(pass, call.Args[2])

	block, idx := enclosingStmt(file, call)
	if block == nil {
		return nil
	}
	stmt := block.List[idx]

	var lhs string
	assign, isAssign := stmt.(*ast.AssignStmt)
	if isAssign && len(assign.Lhs) == 2 {
		lhs = nodeText(pass, assign.Lhs[0])
	}
	if lhs == "" {
		lhs = "_"
	}

	tok := ":="
	if isAssign {
		tok = assign.Tok.String()
	}

	replacement := fmt.Sprintf(
		"req, err %s http.NewRequestWithContext(%s, \"POST\", %s, %s)\n\tif err != nil {\n\t\treturn err\n\t}\n\treq.Header.Set(\"Content-Type\", %s)\n\t%s, err %s %s.Do(req)",
		tok, ctxExpr, urlArg, bodyArg,
		ctypeArg,
		lhs, tok, client,
	)

	var edits []analysis.TextEdit
	edits = append(edits, analysis.TextEdit{
		Pos:     stmt.Pos(),
		End:     stmt.End(),
		NewText: []byte(replacement),
	})

	if inject {
		edits = append(edits, contextInjectionEdit(enc, pass))
	}
	if addImport {
		edits = append(edits, contextImportEdit(file)...)
	}

	return []analysis.SuggestedFix{{
		Message:   fmt.Sprintf("use http.NewRequestWithContext with %s.Do", client),
		TextEdits: edits,
	}}
}

// fixHTTPPostForm rewrites http.PostForm(url, data) and client.PostForm().
func fixHTTPPostForm(pass *analysis.Pass, file *ast.File, call *ast.CallExpr, ctxExpr string, enc *ast.FuncDecl, inject bool, addImport bool, client string) []analysis.SuggestedFix {
	urlArg := nodeText(pass, call.Args[0])
	dataArg := nodeText(pass, call.Args[1])

	block, idx := enclosingStmt(file, call)
	if block == nil {
		return nil
	}
	stmt := block.List[idx]

	var lhs string
	assign, isAssign := stmt.(*ast.AssignStmt)
	if isAssign && len(assign.Lhs) == 2 {
		lhs = nodeText(pass, assign.Lhs[0])
	}
	if lhs == "" {
		lhs = "_"
	}

	tok := ":="
	if isAssign {
		tok = assign.Tok.String()
	}

	replacement := fmt.Sprintf(
		"req, err %s http.NewRequestWithContext(%s, \"POST\", %s, strings.NewReader(%s.Encode()))\n\tif err != nil {\n\t\treturn err\n\t}\n\treq.Header.Set(\"Content-Type\", \"application/x-www-form-urlencoded\")\n\t%s, err %s %s.Do(req)",
		tok, ctxExpr, urlArg, dataArg,
		lhs, tok, client,
	)

	var edits []analysis.TextEdit
	edits = append(edits, analysis.TextEdit{
		Pos:     stmt.Pos(),
		End:     stmt.End(),
		NewText: []byte(replacement),
	})

	if inject {
		edits = append(edits, contextInjectionEdit(enc, pass))
	}
	if addImport {
		edits = append(edits, contextImportEdit(file)...)
	}

	return []analysis.SuggestedFix{{
		Message:   fmt.Sprintf("use http.NewRequestWithContext with %s.Do", client),
		TextEdits: edits,
	}}
}

// nodeText returns the source text of an AST node.
func nodeText(pass *analysis.Pass, node ast.Node) string {
	var buf bytes.Buffer
	if err := printer.Fprint(&buf, pass.Fset, node); err != nil {
		return ""
	}
	return buf.String()
}

// contextInjectionEdit creates a TextEdit that inserts a context statement at the
// beginning of a function body.
func contextInjectionEdit(fd *ast.FuncDecl, pass *analysis.Pass) analysis.TextEdit {
	text := contextInjectionText(fd, pass)
	insertPos := fd.Body.Lbrace + 1
	return analysis.TextEdit{
		Pos:     insertPos,
		End:     insertPos,
		NewText: []byte("\n\t" + text),
	}
}

// hasContextImport checks if the file already imports "context".
func hasContextImport(file *ast.File) bool {
	for _, imp := range file.Imports {
		if imp.Path.Value == `"context"` {
			return true
		}
	}
	return false
}

// contextImportEdit creates TextEdits to add "context" to the import block.
func contextImportEdit(file *ast.File) []analysis.TextEdit {
	// Find the import declaration.
	for _, decl := range file.Decls {
		gd, ok := decl.(*ast.GenDecl)
		if !ok || gd.Tok != token.IMPORT {
			continue
		}
		if gd.Lparen == 0 {
			// Single import without parens: import "net/http"
			// Insert before it.
			return []analysis.TextEdit{{
				Pos:     gd.Pos(),
				End:     gd.Pos(),
				NewText: []byte("import \"context\"\n"),
			}}
		}
		// Grouped import: insert "context" right after the opening paren.
		return []analysis.TextEdit{{
			Pos:     gd.Lparen + 1,
			End:     gd.Lparen + 1,
			NewText: []byte("\n\t\"context\""),
		}}
	}
	// No import block at all — add one after the package declaration.
	return []analysis.TextEdit{{
		Pos:     file.Name.End(),
		End:     file.Name.End(),
		NewText: []byte("\n\nimport \"context\"\n"),
	}}
}