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
 
@@ -253,7 +262,7 @@ _Configuration_: N/A
 _Description_: In GO it is idiomatic to minimize nesting statements, a typical example is to avoid if-then-else constructions. This rule spots constructions like
 ```go
 if cond {
-  // do something 
+  // do something
 } else {
   // do other thing
   return ...
@@ -263,7 +272,7 @@ that can be rewritten into more idiomatic:
 ```go
 if ! cond {
   // do other thing
-  return ... 
+  return ...
 }
 
 // do something
@@ -315,8 +324,8 @@ _Description_: Exported function and methods should have comments. This warns on
 
 More information [here](https://github.com/golang/go/wiki/CodeReviewComments#doc-comments)
 
-_Configuration_: ([]string) rule flags. 
-Please notice that without configuration, the default behavior of the rule is that of its `golint` counterpart. 
+_Configuration_: ([]string) rule flags.
+Please notice that without configuration, the default behavior of the rule is that of its `golint` counterpart.
 Available flags are:
 
 * _checkPrivateReceivers_ enables checking public methods of private types
@@ -426,8 +435,8 @@ Example:
 ```
 ## import-shadowing
 
-_Description_: In GO it is possible to declare identifiers (packages, structs, 
-interfaces, parameters, receivers, variables, constants...) that conflict with the 
+_Description_: In GO it is possible to declare identifiers (packages, structs,
+interfaces, parameters, receivers, variables, constants...) that conflict with the
 name of an imported package. This rule spots identifiers that shadow an import.
 
 _Configuration_: N/A
@@ -517,7 +526,7 @@ _Configuration_: N/A
 
 ## range-val-address
 
-_Description_: Range variables in a loop are reused at each iteration. This rule warns when assigning the address of the variable, passing the address to append() or using it in a map. 
+_Description_: Range variables in a loop are reused at each iteration. This rule warns when assigning the address of the variable, passing the address to append() or using it in a map.
 
 _Configuration_: N/A
 
@@ -535,7 +544,7 @@ Even if possible, redefining these built in names can lead to bugs very difficul
 _Configuration_: N/A
 
 ## string-of-int
-_Description_:  explicit type conversion `string(i)` where `i` has an integer type other than `rune` might behave not as expected by the developer (e.g. `string(42)` is not `"42"`). This rule spot that kind of suspicious conversions. 
+_Description_:  explicit type conversion `string(i)` where `i` has an integer type other than `rune` might behave not as expected by the developer (e.g. `string(42)` is not `"42"`). This rule spot that kind of suspicious conversions.
 
 _Configuration_: N/A
 
@@ -548,7 +557,7 @@ _Configuration_: Each argument is a slice containing 2-3 strings: a scope, a reg
 
 1. The first string defines a scope. This controls which string literals the regex will apply to, and is defined as a function argument. It must contain at least a function name (`core.WriteError`). Scopes may optionally contain a number specifying which argument in the function to check (`core.WriteError[1]`), as well as a struct field (`core.WriteError[1].Message`, only works for top level fields). Function arguments are counted starting at 0, so `[0]` would refer to the first argument, `[1]` would refer to the second, etc. If no argument number is provided, the first argument will be used (same as `[0]`).
 
-2. The second string is a regular expression (beginning and ending with a `/` character), which will be used to check the string literals in the scope. 
+2. The second string is a regular expression (beginning and ending with a `/` character), which will be used to check the string literals in the scope.
 
 3. The third string (optional) is a message containing the purpose for the regex, which will be used in lint errors.
 
@@ -663,8 +672,8 @@ _Configuration_: N/A
 
 ## useless-break
 
-_Description_: This rule warns on useless `break` statements in case clauses of switch and select statements. GO, unlike other programming languages like C, only executes statements of the selected case while ignoring the subsequent case clauses. 
-Therefore, inserting a `break` at the end of a case clause has no effect. 
+_Description_: This rule warns on useless `break` statements in case clauses of switch and select statements. GO, unlike other programming languages like C, only executes statements of the selected case while ignoring the subsequent case clauses.
+Therefore, inserting a `break` at the end of a case clause has no effect.
 
 Because `break` statements are rarely used in case clauses, when switch or select statements are inside a for-loop, the programmer might wrongly assume that a `break` in a case clause will take the control out of the loop.
 The rule emits a specific warning for such cases.

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,22 +39,24 @@ func (r *ContextAsArgumentRule) Name() string {
 }
 
 type lintContextArguments struct {
-	file      *lint.File
-	fileAst   *ast.File
-	onFailure func(lint.Failure)
+	allowTypesLUT map[string]struct{}
+	onFailure     func(lint.Failure)
 }
 
 func (w lintContextArguments) Visit(n ast.Node) ast.Visitor {
 	fn, ok := n.(*ast.FuncDecl)
 	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/golint/context-as-argument.go → 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/
 }