Add handled response directive

README.md

@@ -30,6 +30,38 @@ You can enable additional checks with the `-check-consumption` flag to also veri
 $ go vet -vettool=$(which bodyclose) -check-consumption github.com/timakin/go_api/...
 ```
 
+### Handled Response Directive
+
+Use `//bodyclose:handled` on a function that fully handles every returned
+`*http.Response` body before returning:
+
+```go
+// bodyclose:handled
+func doRequest(req *http.Request, handle func([]byte) error) (*http.Response, error) {
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	if err := handle(body); err != nil {
+		return nil, err
+	}
+
+	resp.Body = io.NopCloser(bytes.NewReader(body))
+	return resp, nil
+}
+```
+
+Calls to an annotated function are skipped by both the close and consumption
+checks, including calls from another package.
+
+The directive works only for statically resolved function calls.
+
 #### Supported Consumption Patterns
 
 When `-check-consumption` is enabled, the following patterns are recognized as valid body consumption:

passes/bodyclose/bodyclose.go

@@ -3,6 +3,8 @@ package bodyclose
 import (
 	"fmt"
 	"go/ast"
+	"go/parser"
+	"go/token"
 	"go/types"
 	"strconv"
 	"strings"
@@ -31,27 +33,30 @@ var checkConsumptionFlag bool
 const (
 	Doc = "checks whether HTTP response body is closed successfully"
 
-	nethttpPath = "net/http"
-	closeMethod = "Close"
+	nethttpPath              = "net/http"
+	closeMethod              = "Close"
+	responseHandledDirective = "bodyclose:handled"
 )
 
 type runner struct {
-	pass             *analysis.Pass
-	resObj           types.Object
-	resTyp           *types.Pointer
-	bodyObj          types.Object
-	closeMthd        *types.Func
-	skipFile         map[*ast.File]bool
-	checkConsumption bool
+	pass                      *analysis.Pass
+	resObj                    types.Object
+	resTyp                    *types.Pointer
+	bodyObj                   types.Object
+	closeMthd                 *types.Func
+	skipFile                  map[*ast.File]bool
+	responseHandledDirectives map[string]map[int]struct{}
+	checkConsumption          bool
 }
 
 // run executes an analysis for the pass
 func run(pass *analysis.Pass) (interface{}, error) {
 	r := runner{
-		pass:             pass,
-		checkConsumption: checkConsumptionFlag,
+		pass:                      pass,
+		skipFile:                  make(map[*ast.File]bool),
+		responseHandledDirectives: make(map[string]map[int]struct{}),
+		checkConsumption:          checkConsumptionFlag,
 	}
-	funcs := pass.ResultOf[buildssa.Analyzer].(*buildssa.SSA).SrcFuncs
 
 	r.resObj = analysisutil.LookupFromImports(pass.Pkg.Imports(), nethttpPath, "Response")
 	if r.resObj == nil {
@@ -91,7 +96,7 @@ func run(pass *analysis.Pass) (interface{}, error) {
 		}
 	}
 
-	r.skipFile = map[*ast.File]bool{}
+	funcs := pass.ResultOf[buildssa.Analyzer].(*buildssa.SSA).SrcFuncs
 FuncLoop:
 	for _, f := range funcs {
 		// skip if the function is just referenced
@@ -118,12 +123,64 @@ FuncLoop:
 	return nil, nil
 }
 
+func (r *runner) responseHandledByDirective(call *ssa.Call) bool {
+	callee := call.Call.StaticCallee()
+	if callee == nil {
+		return false
+	}
+	fn, ok := callee.Object().(*types.Func)
+	if !ok {
+		return false
+	}
+
+	pos := r.pass.Fset.PositionFor(fn.Pos(), false)
+	if !pos.IsValid() || pos.Filename == "" {
+		return false
+	}
+
+	lines, ok := r.responseHandledDirectives[pos.Filename]
+	if !ok {
+		lines = responseHandledDirectiveLines(pos.Filename)
+		r.responseHandledDirectives[pos.Filename] = lines
+	}
+	_, ok = lines[pos.Line]
+	return ok
+}
+
+func responseHandledDirectiveLines(filename string) map[int]struct{} {
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
+	if err != nil {
+		return nil
+	}
+
+	lines := make(map[int]struct{})
+	for _, decl := range file.Decls {
+		fn, ok := decl.(*ast.FuncDecl)
+		if !ok || fn.Doc == nil {
+			continue
+		}
+		for _, comment := range fn.Doc.List {
+			text := strings.TrimSpace(strings.TrimPrefix(comment.Text, "//"))
+			if text == responseHandledDirective {
+				lines[fset.Position(fn.Name.Pos()).Line] = struct{}{}
+				break
+			}
+		}
+	}
+	return lines
+}
+
 func (r *runner) isopen(b *ssa.BasicBlock, i int) bool {
 	call, ok := r.getReqCall(b.Instrs[i])
 	if !ok {
 		return false
 	}
 
+	if r.responseHandledByDirective(call) {
+		return false
+	}
+
 	if len(*call.Referrers()) == 0 {
 		return true
 	}

passes/bodyclose/testdata/src/a/handledresponse.go

@@ -0,0 +1,22 @@
+package a
+
+import (
+	"net/http"
+
+	"handledresponse"
+)
+
+//bodyclose:handled
+func handledResponse() (*http.Response, error) {
+	return http.Get("http://example.com/")
+}
+
+func openResponse() (*http.Response, error) {
+	return http.Get("http://example.com/")
+}
+
+func responseHandledDirectiveCallSites() {
+	_, _ = handledResponse()
+	_, _ = handledresponse.Get()
+	_, _ = openResponse() // want "response body must be closed"
+}

passes/bodyclose/testdata/src/consumption/consumption.go

@@ -200,6 +200,15 @@ func properResponseBodyConsumptionWithRequestBody(w http.ResponseWriter, r *http
 	io.ReadAll(resp.Body) // This is RESPONSE body consumption
 }
 
+// bodyclose:handled
+func handledResponse() (*http.Response, error) {
+	return http.Get("http://example.com")
+}
+
+func responseHandledDirectiveSkipsConsumption() {
+	_, _ = handledResponse()
+}
+
 // closeBeforeConsume is commented out because current implementation
 // doesn't detect execution order (documented limitation)
 //

passes/bodyclose/testdata/src/handledresponse/handledresponse.go

@@ -0,0 +1,8 @@
+package handledresponse
+
+import "net/http"
+
+// bodyclose:handled
+func Get() (*http.Response, error) {
+	return http.Get("http://example.com/")
+}