Allow whitelist for the context parameter check (#616)

* Allow a whitelist for the context parameter check

This allows users to configure a set of types that may appear before
`context.Context`.

Notably, I think this rule is useful for allowing the `*testing.T` type
to come before `context.Context`, though there may be other uses (such
as putting a tracer before it, etc).

See #605 for a little more context on this.

Fixes #605

* Save a level of indentation in context-as-arg validation

We can unindent if we make the above check more specific

* refactoring taking into account chavacava's review

Co-authored-by: chavacava <salvadorcavadini+github@gmail.com>

RULES_DESCRIPTIONS.md

@@ -187,7 +187,16 @@ _Configuration_: N/A
 
 _Description_: By [convention](https://github.com/golang/go/wiki/CodeReviewComments#contexts), `context.Context` should be the first parameter of a function. This rule spots function declarations that do not follow the convention.
 
-_Configuration_: N/A
+_Configuration_:
+
+* `allowTypesBefore` : (string) comma-separated list of types that may be before 'context.Context'
+
+Example:
+
+```toml
+[rule.context-as-argument]
+  arguments = [{allowTypesBefore = "*testing.T,*github.com/user/repo/testing.Harness"}]
+```
 
 ## context-keys-type
 

rule/context-as-argument.go

@@ -1,28 +1,34 @@
 package rule
 
 import (
+	"fmt"
 	"go/ast"
+	"strings"
 
 	"github.com/mgechev/revive/lint"
 )
 
 // ContextAsArgumentRule lints given else constructs.
-type ContextAsArgumentRule struct{}
+type ContextAsArgumentRule struct {
+	allowTypesLUT map[string]struct{}
+}
 
 // Apply applies the rule to given file.
-func (r *ContextAsArgumentRule) Apply(file *lint.File, _ lint.Arguments) []lint.Failure {
-	var failures []lint.Failure
+func (r *ContextAsArgumentRule) Apply(file *lint.File, args lint.Arguments) []lint.Failure {
 
-	fileAst := file.AST
+	if r.allowTypesLUT == nil {
+		r.allowTypesLUT = getAllowTypesFromArguments(args)
+	}
+
+	var failures []lint.Failure
 	walker := lintContextArguments{
-		file:    file,
-		fileAst: fileAst,
+		allowTypesLUT: r.allowTypesLUT,
 		onFailure: func(failure lint.Failure) {
 			failures = append(failures, failure)
 		},
 	}
 
-	ast.Walk(walker, fileAst)
+	ast.Walk(walker, file.AST)
 
 	return failures
 }
@@ -33,8 +39,7 @@ func (r *ContextAsArgumentRule) Name() string {
 }
 
 type lintContextArguments struct {
-	file      *lint.File
-	fileAst   *ast.File
+	allowTypesLUT map[string]struct{}
 	onFailure     func(lint.Failure)
 }
 
@@ -43,12 +48,15 @@ func (w lintContextArguments) Visit(n ast.Node) ast.Visitor {
 	if !ok || len(fn.Type.Params.List) <= 1 {
 		return w
 	}
+
+	fnArgs := fn.Type.Params.List
+
 	// A context.Context should be the first parameter of a function.
 	// Flag any that show up after the first.
-	previousArgIsCtx := isPkgDot(fn.Type.Params.List[0].Type, "context", "Context")
-	for _, arg := range fn.Type.Params.List[1:] {
+	isCtxStillAllowed := true
+	for _, arg := range fnArgs {
 		argIsCtx := isPkgDot(arg.Type, "context", "Context")
-		if argIsCtx && !previousArgIsCtx {
+		if argIsCtx && !isCtxStillAllowed {
 			w.onFailure(lint.Failure{
 				Node:       arg,
 				Category:   "arg-order",
@@ -57,7 +65,41 @@ func (w lintContextArguments) Visit(n ast.Node) ast.Visitor {
 			})
 			break // only flag one
 		}
-		previousArgIsCtx = argIsCtx
+
+		typeName := gofmt(arg.Type)
+		// a parameter of type context.Context is still allowed if the current arg type is in the LUT
+		_, isCtxStillAllowed = w.allowTypesLUT[typeName]
 	}
-	return w
+
+	return nil // avoid visiting the function body
+}
+
+func getAllowTypesFromArguments(args lint.Arguments) map[string]struct{} {
+	allowTypesBefore := []string{}
+	if len(args) >= 1 {
+		argKV, ok := args[0].(map[string]interface{})
+		if !ok {
+			panic(fmt.Sprintf("Invalid argument to the context-as-argument rule. Expecting a k,v map, got %T", args[0]))
+		}
+		for k, v := range argKV {
+			switch k {
+			case "allowTypesBefore":
+				typesBefore, ok := v.(string)
+				if !ok {
+					panic(fmt.Sprintf("Invalid argument to the context-as-argument.allowTypesBefore rule. Expecting a string, got %T", v))
+				}
+				allowTypesBefore = append(allowTypesBefore, strings.Split(typesBefore, ",")...)
+			default:
+				panic(fmt.Sprintf("Invalid argument to the context-as-argument rule. Unrecognized key %s", k))
+			}
+		}
+	}
+
+	result := make(map[string]struct{}, len(allowTypesBefore))
+	for _, v := range allowTypesBefore {
+		result[v] = struct{}{}
+	}
+
+	result["context.Context"] = struct{}{} // context.Context is always allowed before another context.Context
+	return result
 }

rule/utils.go

@@ -92,6 +92,7 @@ func validType(T types.Type) bool {
 		!strings.Contains(T.String(), "invalid type") // good but not foolproof
 }
 
+// isPkgDot checks if the expression is <pkg>.<name>
 func isPkgDot(expr ast.Expr, pkg, name string) bool {
 	sel, ok := expr.(*ast.SelectorExpr)
 	return ok && isIdent(sel.X, pkg) && isIdent(sel.Sel, name)
@@ -132,14 +133,6 @@ func pick(n ast.Node, fselect func(n ast.Node) bool, f func(n ast.Node) []ast.No
 	return result
 }
 
-func pickFromExpList(l []ast.Expr, fselect func(n ast.Node) bool, f func(n ast.Node) []ast.Node) []ast.Node {
-	result := make([]ast.Node, 0)
-	for _, e := range l {
-		result = append(result, pick(e, fselect, f)...)
-	}
-	return result
-}
-
 type picker struct {
 	fselect  func(n ast.Node) bool
 	onSelect func(n ast.Node)

test/context-as-argument_test.go

@@ -0,0 +1,19 @@
+package test
+
+import (
+	"testing"
+
+	"github.com/mgechev/revive/lint"
+	"github.com/mgechev/revive/rule"
+)
+
+func TestContextAsArgument(t *testing.T) {
+	testRule(t, "context-as-argument", &rule.ContextAsArgumentRule{}, &lint.RuleConfig{
+		Arguments: []interface{}{
+			map[string]interface{}{
+				"allowTypesBefore": "AllowedBeforeType,AllowedBeforeStruct,*AllowedBeforePtrStruct,*testing.T",
+			},
+		},
+	},
+	)
+}

test/golint_test.go

@@ -31,7 +31,6 @@ var rules = []lint.Rule{
 	&rule.UnexportedReturnRule{},
 	&rule.TimeNamingRule{},
 	&rule.ContextKeysType{},
-	&rule.ContextAsArgumentRule{},
 }
 
 func TestAll(t *testing.T) {

testdata/context-as-argument.go

@@ -5,8 +5,18 @@ package foo
 
 import (
 	"context"
+	"testing"
 )
 
+// AllowedBeforeType is a type that is configured to be allowed before context.Context
+type AllowedBeforeType string
+
+// AllowedBeforeStruct is a type that is configured to be allowed before context.Context
+type AllowedBeforeStruct struct{}
+
+// AllowedBeforePtrStruct is a type that is configured to be allowed before context.Context
+type AllowedBeforePtrStruct struct{}
+
 // A proper context.Context location
 func x(ctx context.Context) { // ok
 }
@@ -15,6 +25,22 @@ func x(ctx context.Context) { // ok
 func x(ctx context.Context, s string) { // ok
 }
 
+// *testing.T is permitted in the linter config for the test
+func x(t *testing.T, ctx context.Context) { // ok
+}
+
+func x(_ AllowedBeforeType, _ AllowedBeforeType, ctx context.Context) { // ok
+}
+
+func x(_, _ AllowedBeforeType, ctx context.Context) { // ok
+}
+
+func x(_ *AllowedBeforePtrStruct, ctx context.Context) { // ok
+}
+
+func x(_ AllowedBeforePtrStruct, ctx context.Context) { // MATCH /context.Context should be the first parameter of a function/
+}
+
 // An invalid context.Context location
 func y(s string, ctx context.Context) { // MATCH /context.Context should be the first parameter of a function/
 }