Add match call by type

core/helpers.go

@@ -88,6 +88,33 @@ func MatchCallByPackage(n ast.Node, c *Context, pkg string, names ...string) (*a
 	return nil, false
 }
 
+// MatchCallByType ensures that the node is a call expression to a
+// specific object type.
+//
+// Usage:
+// 	node, matched := MatchCallByType(n, ctx, "bytes.Buffer", "WriteTo", "Write")
+//
+func MatchCallByType(n ast.Node, ctx *Context, requiredType string, calls ...string) (*ast.CallExpr, bool) {
+	switch callExpr := n.(type) {
+	case *ast.CallExpr:
+		switch fn := callExpr.Fun.(type) {
+		case *ast.SelectorExpr:
+			switch expr := fn.X.(type) {
+			case *ast.Ident:
+				t := ctx.Info.TypeOf(expr)
+				if t != nil && t.String() == requiredType {
+					for _, call := range calls {
+						if fn.Sel.Name == call {
+							return callExpr, true
+						}
+					}
+				}
+			}
+		}
+	}
+	return nil, false
+}
+
 // MatchCompLit will match an ast.CompositeLit if its string value obays the given regex.
 func MatchCompLit(n ast.Node, r *regexp.Regexp) *ast.CompositeLit {
 	t := reflect.TypeOf(&ast.CompositeLit{})

core/helpers_test.go

@@ -0,0 +1,59 @@
+package core
+
+import (
+	"go/ast"
+	"testing"
+)
+
+type dummyCallback func(ast.Node, *Context, string, ...string) (*ast.CallExpr, bool)
+
+type dummyRule struct {
+	MetaData
+	pkgOrType      string
+	funcsOrMethods []string
+	callback       dummyCallback
+	callExpr       []ast.Node
+	matched        int
+}
+
+func (r *dummyRule) Match(n ast.Node, c *Context) (gi *Issue, err error) {
+	if callexpr, matched := r.callback(n, c, r.pkgOrType, r.funcsOrMethods...); matched {
+		r.matched += 1
+		r.callExpr = append(r.callExpr, callexpr)
+	}
+	return nil, nil
+}
+
+func TestMatchCallByType(t *testing.T) {
+	config := map[string]interface{}{"ignoreNosec": false}
+	analyzer := NewAnalyzer(config, nil)
+	rule := &dummyRule{
+		MetaData: MetaData{
+			Severity:   Low,
+			Confidence: Low,
+			What:       "A dummy rule",
+		},
+		pkgOrType:      "bytes.Buffer",
+		funcsOrMethods: []string{"Write"},
+		callback:       MatchCallByType,
+		callExpr:       []ast.Node{},
+		matched:        0,
+	}
+	analyzer.AddRule(rule, []ast.Node{(*ast.CallExpr)(nil)})
+	source := `
+	package main
+	import (
+		"bytes"
+		"fmt"
+	)
+	func main() {
+		var b bytes.Buffer
+		b.Write([]byte("Hello "))
+		fmt.Fprintf(&b, "world!")
+	}`
+
+	analyzer.ProcessSource("dummy.go", source)
+	if rule.matched != 1 || len(rule.callExpr) != 1 {
+		t.Errorf("Expected to match a bytes.Buffer.Write call")
+	}
+}