cmd/vet: -lostcancel: check for discarded result of context.WithCancel

The cfg subpackage builds a control flow graph of ast.Nodes.
The lostcancel module checks this graph to find paths, from a call to
WithCancel to a return statement, on which the cancel variable is
not used.  (A trivial case is simply assigning the cancel result to
the blank identifier.)

In a sample of 50,000 source files, the new check found 2068 blank
assignments and 118 return-before-cancel errors.  I manually inspected
20 of the latter and didn't find a single false positive among them.

Change-Id: I84cd49445f9f8d04908b04881eb1496a96611205
Reviewed-on: https://go-review.googlesource.com/24150Reviewed-by: Robert Griesemer <gri@golang.org>
Reviewed-by: Rob Pike <r@golang.org>

src/cmd/vet/deadcode.go

@@ -29,6 +29,8 @@ type deadState struct {
 }
 
 // checkUnreachable checks a function body for dead code.
+//
+// TODO(adonovan): use the new cfg package, which is more precise.
 func checkUnreachable(f *File, node ast.Node) {
 	var body *ast.BlockStmt
 	switch n := node.(type) {

src/cmd/vet/doc.go

@@ -91,6 +91,15 @@ Flag: -tests
 Mistakes involving tests including functions with incorrect names or signatures
 and example tests that document identifiers not in the package.
 
+Failure to call the cancellation function returned by context.WithCancel.
+
+Flag: -lostcancel
+
+The cancellation function returned by context.WithCancel, WithTimeout,
+and WithDeadline must be called or the new context will remain live
+until its parent context is cancelled.
+(The background context is never cancelled.)
+
 Methods
 
 Flag: -methods

src/cmd/vet/internal/cfg/cfg.go

+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// This package constructs a simple control-flow graph (CFG) of the
+// statements and expressions within a single function.
+//
+// Use cfg.New to construct the CFG for a function body.
+//
+// The blocks of the CFG contain all the function's non-control
+// statements.  The CFG does not contain control statements such as If,
+// Switch, Select, and Branch, but does contain their subexpressions.
+// For example, this source code:
+//
+//	if x := f(); x != nil {
+//		T()
+//	} else {
+//		F()
+//	}
+//
+// produces this CFG:
+//
+//    1:  x := f()
+//        x != nil
+//        succs: 2, 3
+//    2:  T()
+//        succs: 4
+//    3:  F()
+//        succs: 4
+//    4:
+//
+// The CFG does contain Return statements; even implicit returns are
+// materialized (at the position of the function's closing brace).
+//
+// The CFG does not record conditions associated with conditional branch
+// edges, nor the short-circuit semantics of the && and || operators,
+// nor abnormal control flow caused by panic.  If you need this
+// information, use golang.org/x/tools/go/ssa instead.
+//
+package cfg
+
+// Although the vet tool has type information, it is often extremely
+// fragmentary, so for simplicity this package does not depend on
+// go/types.  Consequently control-flow conditions are ignored even
+// when constant, and "mayReturn" information must be provided by the
+// client.
+import (
+	"bytes"
+	"fmt"
+	"go/ast"
+	"go/format"
+	"go/token"
+)
+
+// A CFG represents the control-flow graph of a single function.
+//
+// The entry point is Blocks[0]; there may be multiple return blocks.
+type CFG struct {
+	Blocks []*Block // block[0] is entry; order otherwise undefined
+}
+
+// A Block represents a basic block: a list of statements and
+// expressions that are always evaluated sequentially.
+//
+// A block may have 0-2 successors: zero for a return block or a block
+// that calls a function such as panic that never returns; one for a
+// normal (jump) block; and two for a conditional (if) block.
+type Block struct {
+	Nodes []ast.Node // statements, expressions, and ValueSpecs
+	Succs []*Block   // successor nodes in the graph
+
+	comment     string    // for debugging
+	index       int32     // index within CFG.Blocks
+	unreachable bool      // is block of stmts following return/panic/for{}
+	succs2      [2]*Block // underlying array for Succs
+}
+
+// New returns a new control-flow graph for the specified function body,
+// which must be non-nil.
+//
+// The CFG builder calls mayReturn to determine whether a given function
+// call may return.  For example, calls to panic, os.Exit, and log.Fatal
+// do not return, so the builder can remove infeasible graph edges
+// following such calls.  The builder calls mayReturn only for a
+// CallExpr beneath an ExprStmt.
+func New(body *ast.BlockStmt, mayReturn func(*ast.CallExpr) bool) *CFG {
+	b := builder{
+		mayReturn: mayReturn,
+		cfg:       new(CFG),
+	}
+	b.current = b.newBlock("entry")
+	b.stmt(body)
+
+	// Does control fall off the end of the function's body?
+	// Make implicit return explicit.
+	if b.current != nil && !b.current.unreachable {
+		b.add(&ast.ReturnStmt{
+			Return: body.End() - 1,
+		})
+	}
+
+	return b.cfg
+}
+
+func (b *Block) String() string {
+	return fmt.Sprintf("block %d (%s)", b.index, b.comment)
+}
+
+// Return returns the return statement at the end of this block if present, nil otherwise.
+func (b *Block) Return() (ret *ast.ReturnStmt) {
+	if len(b.Nodes) > 0 {
+		ret, _ = b.Nodes[len(b.Nodes)-1].(*ast.ReturnStmt)
+	}
+	return
+}
+
+// Format formats the control-flow graph for ease of debugging.
+func (g *CFG) Format(fset *token.FileSet) string {
+	var buf bytes.Buffer
+	for _, b := range g.Blocks {
+		fmt.Fprintf(&buf, ".%d: # %s\n", b.index, b.comment)
+		for _, n := range b.Nodes {
+			fmt.Fprintf(&buf, "\t%s\n", formatNode(fset, n))
+		}
+		if len(b.Succs) > 0 {
+			fmt.Fprintf(&buf, "\tsuccs:")
+			for _, succ := range b.Succs {
+				fmt.Fprintf(&buf, " %d", succ.index)
+			}
+			buf.WriteByte('\n')
+		}
+		buf.WriteByte('\n')
+	}
+	return buf.String()
+}
+
+func formatNode(fset *token.FileSet, n ast.Node) string {
+	var buf bytes.Buffer
+	format.Node(&buf, fset, n)
+	// Indent secondary lines by a tab.
+	return string(bytes.Replace(buf.Bytes(), []byte("\n"), []byte("\n\t"), -1))
+}

src/cmd/vet/internal/cfg/cfg_test.go

+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package cfg
+
+import (
+	"bytes"
+	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"testing"
+)
+
+const src = `package main
+
+import "log"
+
+func f1() {
+	live()
+	return
+	dead()
+}
+
+func f2() {
+	for {
+		live()
+	}
+	dead()
+}
+
+func f3() {
+	if true { // even known values are ignored
+		return
+	}
+	for true { // even known values are ignored
+		live()
+	}
+	for {
+		live()
+	}
+	dead()
+}
+
+func f4(x int) {
+	switch x {
+	case 1:
+		live()
+		fallthrough
+	case 2:
+		live()
+		log.Fatal()
+	default:
+		panic("oops")
+	}
+	dead()
+}
+
+func f4(ch chan int) {
+	select {
+	case <-ch:
+		live()
+		return
+	default:
+		live()
+		panic("oops")
+	}
+	dead()
+}
+
+func f5(unknown bool) {
+	for {
+		if unknown {
+			break
+		}
+		continue
+		dead()
+	}
+	live()
+}
+
+func f6(unknown bool) {
+outer:
+	for {
+		for {
+			break outer
+			dead()
+		}
+		dead()
+	}
+	live()
+}
+
+func f7() {
+	for {
+		break nosuchlabel
+		dead()
+	}
+	dead()
+}
+
+func f8() {
+	select{}
+	dead()
+}
+
+func f9(ch chan int) {
+	select {
+	case <-ch:
+		return
+	}
+	dead()
+}
+
+func f10(ch chan int) {
+	select {
+	case <-ch:
+		return
+		dead()
+	default:
+	}
+	live()
+}
+`
+
+func TestDeadCode(t *testing.T) {
+	// We'll use dead code detection to verify the CFG.
+
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "dummy.go", src, parser.Mode(0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, decl := range f.Decls {
+		if decl, ok := decl.(*ast.FuncDecl); ok {
+			g := New(decl.Body, mayReturn)
+
+			// Mark blocks reachable from entry.
+			live := make(map[*Block]bool)
+			var visit func(*Block)
+			visit = func(b *Block) {
+				if !live[b] {
+					live[b] = true
+					for _, succ := range b.Succs {
+						visit(succ)
+					}
+				}
+			}
+			visit(g.Blocks[0])
+
+			// Print statements in unreachable blocks
+			// (in order determined by builder).
+			var buf bytes.Buffer
+			for _, b := range g.Blocks {
+				if !live[b] {
+					for _, n := range b.Nodes {
+						fmt.Fprintf(&buf, "\t%s\n", formatNode(fset, n))
+					}
+				}
+			}
+
+			// Check that the result contains "dead" at least once but not "live".
+			if !bytes.Contains(buf.Bytes(), []byte("dead")) ||
+				bytes.Contains(buf.Bytes(), []byte("live")) {
+				t.Errorf("unexpected dead statements in function %s:\n%s",
+					decl.Name.Name,
+					&buf)
+				t.Logf("control flow graph:\n%s", g.Format(fset))
+			}
+		}
+	}
+}
+
+// A trivial mayReturn predicate that looks only at syntax, not types.
+func mayReturn(call *ast.CallExpr) bool {
+	switch fun := call.Fun.(type) {
+	case *ast.Ident:
+		return fun.Name != "panic"
+	case *ast.SelectorExpr:
+		return fun.Sel.Name != "Fatal"
+	}
+	return true
+}

src/cmd/vet/lostcancel.go

+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"cmd/vet/internal/cfg"
+	"fmt"
+	"go/ast"
+	"go/types"
+	"strconv"
+)
+
+func init() {
+	register("lostcancel",
+		"check for failure to call cancellation function returned by context.WithCancel",
+		checkLostCancel,
+		funcDecl, funcLit)
+}
+
+const debugLostCancel = false
+
+var contextPackage = "context"
+
+// checkLostCancel reports a failure to the call the cancel function
+// returned by context.WithCancel, either because the variable was
+// assigned to the blank identifier, or because there exists a
+// control-flow path from the call to a return statement and that path
+// does not "use" the cancel function.  Any reference to the variable
+// counts as a use, even within a nested function literal.
+//
+// checkLostCancel analyzes a single named or literal function.
+func checkLostCancel(f *File, node ast.Node) {
+	// Fast path: bypass check if file doesn't use context.WithCancel.
+	if !hasImport(f.file, contextPackage) {
+		return
+	}
+
+	// Maps each cancel variable to its defining ValueSpec/AssignStmt.
+	cancelvars := make(map[*types.Var]ast.Node)
+
+	// Find the set of cancel vars to analyze.
+	stack := make([]ast.Node, 0, 32)
+	ast.Inspect(node, func(n ast.Node) bool {
+		switch n.(type) {
+		case *ast.FuncLit:
+			if len(stack) > 0 {
+				return false // don't stray into nested functions
+			}
+		case nil:
+			stack = stack[:len(stack)-1] // pop
+			return true
+		}
+		stack = append(stack, n) // push
+
+		// Look for [{AssignStmt,ValueSpec} CallExpr SelectorExpr]:
+		//
+		//   ctx, cancel    := context.WithCancel(...)
+		//   ctx, cancel     = context.WithCancel(...)
+		//   var ctx, cancel = context.WithCancel(...)
+		//
+		if isContextWithCancel(f, n) && isCall(stack[len(stack)-2]) {
+			var id *ast.Ident // id of cancel var
+			stmt := stack[len(stack)-3]
+			switch stmt := stmt.(type) {
+			case *ast.ValueSpec:
+				if len(stmt.Names) > 1 {
+					id = stmt.Names[1]
+				}
+			case *ast.AssignStmt:
+				if len(stmt.Lhs) > 1 {
+					id, _ = stmt.Lhs[1].(*ast.Ident)
+				}
+			}
+			if id != nil {
+				if id.Name == "_" {
+					f.Badf(id.Pos(), "the cancel function returned by context.%s should be called, not discarded, to avoid a context leak",
+						n.(*ast.SelectorExpr).Sel.Name)
+				} else if v, ok := f.pkg.uses[id].(*types.Var); ok {
+					cancelvars[v] = stmt
+				} else if v, ok := f.pkg.defs[id].(*types.Var); ok {
+					cancelvars[v] = stmt
+				}
+			}
+		}
+
+		return true
+	})
+
+	if len(cancelvars) == 0 {
+		return // no need to build CFG
+	}
+
+	// Tell the CFG builder which functions never return.
+	info := &types.Info{Uses: f.pkg.uses, Selections: f.pkg.selectors}
+	mayReturn := func(call *ast.CallExpr) bool {
+		name := callName(info, call)
+		return !noReturnFuncs[name]
+	}
+
+	// Build the CFG.
+	var g *cfg.CFG
+	switch node := node.(type) {
+	case *ast.FuncDecl:
+		g = cfg.New(node.Body, mayReturn)
+	case *ast.FuncLit:
+		g = cfg.New(node.Body, mayReturn)
+	}
+
+	// Print CFG.
+	if debugLostCancel {
+		fmt.Println(g.Format(f.fset))
+	}
+
+	// Examine the CFG for each variable in turn.
+	// (It would be more efficient to analyze all cancelvars in a
+	// single pass over the AST, but seldom is there more than one.)
+	for v, stmt := range cancelvars {
+		if ret := lostCancelPath(f, g, v, stmt); ret != nil {
+			lineno := f.fset.Position(stmt.Pos()).Line
+			f.Badf(stmt.Pos(), "the %s function is not used on all paths (possible context leak)", v.Name())
+			f.Badf(ret.Pos(), "this return statement may be reached without using the %s var defined on line %d", v.Name(), lineno)
+		}
+	}
+}
+
+func isCall(n ast.Node) bool { _, ok := n.(*ast.CallExpr); return ok }
+
+func hasImport(f *ast.File, path string) bool {
+	for _, imp := range f.Imports {
+		v, _ := strconv.Unquote(imp.Path.Value)
+		if v == path {
+			return true
+		}
+	}
+	return false
+}
+
+// isContextWithCancel reports whether n is one of the qualified identifiers
+// context.With{Cancel,Timeout,Deadline}.
+func isContextWithCancel(f *File, n ast.Node) bool {
+	if sel, ok := n.(*ast.SelectorExpr); ok {
+		switch sel.Sel.Name {
+		case "WithCancel", "WithTimeout", "WithDeadline":
+			if x, ok := sel.X.(*ast.Ident); ok {
+				if pkgname, ok := f.pkg.uses[x].(*types.PkgName); ok {
+					return pkgname.Imported().Path() == contextPackage
+				}
+				// Import failed, so we can't check package path.
+				// Just check the local package name (heuristic).
+				return x.Name == "context"
+			}
+		}
+	}
+	return false
+}
+
+// lostCancelPath finds a path through the CFG, from stmt (which defines
+// the 'cancel' variable v) to a return statement, that doesn't "use" v.
+// If it finds one, it returns the return statement (which may be synthetic).
+func lostCancelPath(f *File, g *cfg.CFG, v *types.Var, stmt ast.Node) *ast.ReturnStmt {
+	// uses reports whether stmts contain a "use" of variable v.
+	uses := func(f *File, v *types.Var, stmts []ast.Node) bool {
+		found := false
+		for _, stmt := range stmts {
+			ast.Inspect(stmt, func(n ast.Node) bool {
+				if id, ok := n.(*ast.Ident); ok {
+					if f.pkg.uses[id] == v {
+						found = true
+					}
+				}
+				return !found
+			})
+		}
+		return found
+	}
+
+	// blockUses computes "uses" for each block, caching the result.
+	memo := make(map[*cfg.Block]bool)
+	blockUses := func(f *File, v *types.Var, b *cfg.Block) bool {
+		res, ok := memo[b]
+		if !ok {
+			res = uses(f, v, b.Nodes)
+			memo[b] = res
+		}
+		return res
+	}
+
+	// Find the var's defining block in the CFG,
+	// plus the rest of the statements of that block.
+	var defblock *cfg.Block
+	var rest []ast.Node
+outer:
+	for _, b := range g.Blocks {
+		for i, n := range b.Nodes {
+			if n == stmt {
+				defblock = b
+				rest = b.Nodes[i+1:]
+				break outer
+			}
+		}
+	}
+	if defblock == nil {
+		panic("internal error: can't find defining block for cancel var")
+	}
+
+	// Is v "used" in the remainder of its defining block?
+	if uses(f, v, rest) {
+		return nil
+	}
+
+	// Does the defining block return without using v?
+	if ret := defblock.Return(); ret != nil {
+		return ret
+	}
+
+	// Search the CFG depth-first for a path, from defblock to a
+	// return block, in which v is never "used".
+	seen := make(map[*cfg.Block]bool)
+	var search func(blocks []*cfg.Block) *ast.ReturnStmt
+	search = func(blocks []*cfg.Block) *ast.ReturnStmt {
+		for _, b := range blocks {
+			if !seen[b] {
+				seen[b] = true
+
+				// Prune the search if the block uses v.
+				if blockUses(f, v, b) {
+					continue
+				}
+
+				// Found path to return statement?
+				if ret := b.Return(); ret != nil {
+					if debugLostCancel {
+						fmt.Printf("found path to return in block %s\n", b)
+					}
+					return ret // found
+				}
+
+				// Recur
+				if ret := search(b.Succs); ret != nil {
+					if debugLostCancel {
+						fmt.Printf(" from block %s\n", b)
+					}
+					return ret
+				}
+			}
+		}
+		return nil
+	}
+	return search(defblock.Succs)
+}
+
+var noReturnFuncs = map[string]bool{
+	"(*testing.common).FailNow": true,
+	"(*testing.common).Fatal":   true,
+	"(*testing.common).Fatalf":  true,
+	"(*testing.common).Skip":    true,
+	"(*testing.common).SkipNow": true,
+	"(*testing.common).Skipf":   true,
+	"log.Fatal":                 true,
+	"log.Fatalf":                true,
+	"log.Fatalln":               true,
+	"os.Exit":                   true,
+	"panic":                     true,
+	"runtime.Goexit":            true,
+}
+
+// callName returns the canonical name of the builtin, method, or
+// function called by call, if known.
+func callName(info *types.Info, call *ast.CallExpr) string {
+	switch fun := call.Fun.(type) {
+	case *ast.Ident:
+		// builtin, e.g. "panic"
+		if obj, ok := info.Uses[fun].(*types.Builtin); ok {
+			return obj.Name()
+		}
+	case *ast.SelectorExpr:
+		if sel, ok := info.Selections[fun]; ok {
+			// method call, e.g. "(*testing.common).Fatal"
+			meth := sel.Obj()
+			return fmt.Sprintf("(%s).%s",
+				meth.Type().(*types.Signature).Recv().Type(),
+				meth.Name())
+		}
+		if obj, ok := info.Uses[fun.Sel]; ok {
+			// qualified identifier, e.g. "os.Exit"
+			return fmt.Sprintf("%s.%s",
+				obj.Pkg().Path(),
+				obj.Name())
+		}
+	}
+
+	// function with no name, or defined in missing imported package
+	return ""
+}

src/cmd/vet/testdata/lostcancel.go

+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package testdata
+
+import (
+	"context"
+	"log"
+	"os"
+	"testing"
+)
+
+// Check the three functions and assignment forms (var, :=, =) we look for.
+// (Do these early: line numbers are fragile.)
+func _() {
+	var ctx, cancel = context.WithCancel() // ERROR "the cancel function is not used on all paths \(possible context leak\)"
+} // ERROR "this return statement may be reached without using the cancel var defined on line 17"
+
+func _() {
+	ctx, cancel2 := context.WithDeadline() // ERROR "the cancel2 function is not used..."
+} // ERROR "may be reached without using the cancel2 var defined on line 21"
+
+func _() {
+	var ctx context.Context
+	var cancel3 func()
+	ctx, cancel3 = context.WithTimeout() // ERROR "function is not used..."
+} // ERROR "this return statement may be reached without using the cancel3 var defined on line 27"
+
+func _() {
+	ctx, _ := context.WithCancel()  // ERROR "the cancel function returned by context.WithCancel should be called, not discarded, to avoid a context leak"
+	ctx, _ = context.WithTimeout()  // ERROR "the cancel function returned by context.WithTimeout should be called, not discarded, to avoid a context leak"
+	ctx, _ = context.WithDeadline() // ERROR "the cancel function returned by context.WithDeadline should be called, not discarded, to avoid a context leak"
+}
+
+func _() {
+	ctx, cancel := context.WithCancel()
+	defer cancel() // ok
+}
+
+func _() {
+	ctx, cancel := context.WithCancel() // ERROR "not used on all paths"
+	if condition {
+		cancel()
+	}
+	return // ERROR "this return statement may be reached without using the cancel var"
+}
+
+func _() {
+	ctx, cancel := context.WithCancel()
+	if condition {
+		cancel()
+	} else {
+		// ok: infinite loop
+		for {
+			print(0)
+		}
+	}
+}
+
+func _() {
+	ctx, cancel := context.WithCancel() // ERROR "not used on all paths"
+	if condition {
+		cancel()
+	} else {
+		for i := 0; i < 10; i++ {
+			print(0)
+		}
+	}
+} // ERROR "this return statement may be reached without using the cancel var"
+
+func _() {
+	ctx, cancel := context.WithCancel()
+	// ok: used on all paths
+	switch someInt {
+	case 0:
+		new(testing.T).FailNow()
+	case 1:
+		log.Fatal()
+	case 2:
+		cancel()
+	case 3:
+		print("hi")
+		fallthrough
+	default:
+		os.Exit(1)
+	}
+}
+
+func _() {
+	ctx, cancel := context.WithCancel() // ERROR "not used on all paths"
+	switch someInt {
+	case 0:
+		new(testing.T).FailNow()
+	case 1:
+		log.Fatal()
+	case 2:
+		cancel()
+	case 3:
+		print("hi") // falls through to implicit return
+	default:
+		os.Exit(1)
+	}
+} // ERROR "this return statement may be reached without using the cancel var"
+
+func _(ch chan int) int {
+	ctx, cancel := context.WithCancel() // ERROR "not used on all paths"
+	select {
+	case <-ch:
+		new(testing.T).FailNow()
+	case y <- ch:
+		print("hi") // falls through to implicit return
+	case ch <- 1:
+		cancel()
+	default:
+		os.Exit(1)
+	}
+} // ERROR "this return statement may be reached without using the cancel var"
+
+func _(ch chan int) int {
+	ctx, cancel := context.WithCancel()
+	// A blocking select must execute one of its cases.
+	select {
+	case <-ch:
+		panic()
+	}
+}
+
+func _() {
+	go func() {
+		ctx, cancel := context.WithCancel() // ERROR "not used on all paths"
+		print(ctx)
+	}() // ERROR "may be reached without using the cancel var"
+}
+
+var condition bool
+var someInt int