fmt: report wrong arity for built-in functions

This commit fixes the index out of range error discussed in #5646 and adds error
handling to avoid that `fmt` panics when such errors are encountered by the formatter

The error handling procedure introduced is similar to the one used by the Scanner struct
responsible for the parsing of a `rego` file, since it uses a slice which is filled with
all the eventual errors that can be found during the format procedure

Fixes: #5646
Signed-off-by: Gianluca Oldani <[email protected]>

Author: Trolloldem

cmd/fmt_test.go

@@ -2,12 +2,15 @@ package cmd
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/format"
 	"github.com/open-policy-agent/opa/util/test"
 )
 
@@ -30,6 +33,27 @@ const unformatted = `
 
 `
 
+const singleWrongArity = `package test
+
+p {
+	a := 1
+	b := 2
+	plus(a, b, c) == 3
+}
+`
+
+const MultipleWrongArity = `package test
+
+p {
+	x:=5
+	y:=7
+	z:=6
+	plus([x, y]) == 3
+	and(true, false, false) == false
+	plus(a, x, y, z)
+}
+`
+
 func TestFmtFormatFile(t *testing.T) {
 	params := fmtCommandParams{}
 	var stdout bytes.Buffer
@@ -269,3 +293,89 @@ func TestFmtFailFileChangesDiff(t *testing.T) {
 		}
 	})
 }
+
+func TestFmtSingleWrongArityError(t *testing.T) {
+	params := fmtCommandParams{}
+	var stdout bytes.Buffer
+
+	files := map[string]string{
+		"policy.rego": singleWrongArity,
+	}
+
+	test.WithTempFS(files, func(path string) {
+		policyFile := filepath.Join(path, "policy.rego")
+		info, err := os.Stat(policyFile)
+		err = formatFile(&params, &stdout, policyFile, info, err)
+		if err == nil {
+			t.Fatalf("Expected error but did not receive one")
+		}
+
+		loc := ast.Location{File: policyFile, Row: 6}
+		errExp := ast.NewError(ast.TypeErr, &loc, "%s: %s", "plus", "arity mismatch")
+		errExp.Details = &format.ArityFormatErrDetail{
+			Have: []string{"var", "var", "var"},
+			Want: []string{"number", "number"},
+		}
+		expectedErrs := ast.Errors(make([]*ast.Error, 1))
+		expectedErrs[0] = errExp
+		expectedSingleWrongArityErr := newError("failed to parse Rego source file: %v", fmt.Errorf("%s: %v", policyFile, expectedErrs))
+
+		if err != expectedSingleWrongArityErr {
+			t.Fatalf("Expected:%s\n\nGot:%s\n\n", expectedSingleWrongArityErr, err)
+		}
+	})
+}
+
+func TestFmtMultipleWrongArityError(t *testing.T) {
+	params := fmtCommandParams{}
+	var stdout bytes.Buffer
+
+	files := map[string]string{
+		"policy.rego": MultipleWrongArity,
+	}
+
+	test.WithTempFS(files, func(path string) {
+		policyFile := filepath.Join(path, "policy.rego")
+		info, err := os.Stat(policyFile)
+		err = formatFile(&params, &stdout, policyFile, info, err)
+		if err == nil {
+			t.Fatalf("Expected error but did not receive one")
+		}
+
+		locations := []ast.Location{
+			{File: policyFile, Row: 7},
+			{File: policyFile, Row: 8},
+			{File: policyFile, Row: 9},
+		}
+		haveStrings := [][]string{
+			{"array"},
+			{"boolean", "boolean", "boolean"},
+			{"var", "var", "var", "var"},
+		}
+		wantStrings := [][]string{
+			{"number", "number"},
+			{"set[any]", "set[any]"},
+			{"number", "number"},
+		}
+		operators := []string{
+			"plus",
+			"and",
+			"plus",
+		}
+		expectedErrs := ast.Errors(make([]*ast.Error, 3))
+		for i := 0; i < 3; i++ {
+			loc := locations[i]
+			errExp := ast.NewError(ast.TypeErr, &loc, "%s: %s", operators[i], "arity mismatch")
+			errExp.Details = &format.ArityFormatErrDetail{
+				Have: haveStrings[i],
+				Want: wantStrings[i],
+			}
+			expectedErrs[i] = errExp
+		}
+		expectedMultipleWrongArityErr := newError("failed to parse Rego source file: %v", fmt.Errorf("%s: %v", policyFile, expectedErrs))
+
+		if err != expectedMultipleWrongArityErr {
+			t.Fatalf("Expected:%s\n\nGot:%s\n\n", expectedMultipleWrongArityErr, err)
+		}
+	})
+}

format/format.go

@@ -10,9 +10,11 @@ import (
 	"fmt"
 	"regexp"
 	"sort"
+	"strings"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/internal/future"
+	"github.com/open-policy-agent/opa/types"
 )
 
 // Opts lets you control the code formatting via `AstWithOpts()`.
@@ -37,6 +39,7 @@ func Source(filename string, src []byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+
 	formatted, err := Ast(module)
 	if err != nil {
 		return nil, fmt.Errorf("%s: %v", filename, err)
@@ -142,6 +145,7 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 
 	w := &writer{
 		indent: "\t",
+		errs:   make([]*ast.Error, 0),
 	}
 
 	switch x := x.(type) {
@@ -178,6 +182,9 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 		return nil, fmt.Errorf("not an ast element: %v", x)
 	}
 
+	if len(w.errs) > 0 {
+		return nil, w.errs
+	}
 	return squashTrailingNewlines(w.buf.Bytes()), nil
 }
 
@@ -228,6 +235,7 @@ type writer struct {
 	inline    bool
 	beforeEnd *ast.Comment
 	delay     bool
+	errs      ast.Errors
 }
 
 func (w *writer) writeModule(module *ast.Module, o fmtOpts) {
@@ -587,7 +595,7 @@ func (w *writer) writeSomeDecl(decl *ast.SomeDecl, comments []*ast.Comment) []*a
 				w.write(",")
 			}
 		case ast.Call:
-			comments = w.writeInOperator(false, val[1:], comments)
+			comments = w.writeInOperator(false, val[1:], comments, decl.Location, ast.BuiltinMap[val[0].String()].Decl)
 		}
 	}
 
@@ -622,7 +630,7 @@ func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) []*a
 
 	switch operator {
 	case ast.Member.Name, ast.MemberWithKey.Name:
-		return w.writeInOperator(false, terms[1:], comments)
+		return w.writeInOperator(false, terms[1:], comments, terms[0].Location, ast.BuiltinMap[terms[0].String()].Decl)
 	}
 
 	bi, ok := ast.BuiltinMap[operator]
@@ -647,6 +655,9 @@ func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) []*a
 		comments = w.writeTerm(terms[2], comments)
 		return comments
 	}
+	// NOTE(Trolloldem): in this point we are operating with a built-in function with the
+	// wrong arity even when the assignment notation is used
+	w.errs = append(w.errs, ArityFormatMismatchError(terms[1:], terms[0].String(), terms[0].Location, bi.Decl))
 	return w.writeFunctionCallPlain(terms, comments)
 }
 
@@ -708,7 +719,7 @@ func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Co
 	case ast.Var:
 		w.write(w.formatVar(x))
 	case ast.Call:
-		comments = w.writeCall(parens, x, comments)
+		comments = w.writeCall(parens, x, term.Location, comments)
 	case fmt.Stringer:
 		w.write(x.String())
 	}
@@ -760,7 +771,7 @@ func (w *writer) formatVar(v ast.Var) string {
 	return v.String()
 }
 
-func (w *writer) writeCall(parens bool, x ast.Call, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeCall(parens bool, x ast.Call, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
 	bi, ok := ast.BuiltinMap[x[0].String()]
 	if !ok || bi.Infix == "" {
 		return w.writeFunctionCallPlain(x, comments)
@@ -769,13 +780,22 @@ func (w *writer) writeCall(parens bool, x ast.Call, comments []*ast.Comment) []*
 	if bi.Infix == "in" {
 		// NOTE(sr): `in` requires special handling, mirroring what happens in the parser,
 		// since there can be one or two lhs arguments.
-		return w.writeInOperator(true, x[1:], comments)
+		return w.writeInOperator(true, x[1:], comments, loc, bi.Decl)
 	}
 
 	// TODO(tsandall): improve to consider precedence?
 	if parens {
 		w.write("(")
 	}
+
+	// NOTE(Trolloldem): writeCall is only invoked when the function call is a term
+	// of another function. The only valid arity is the one of the
+	// built-in function
+	if len(bi.Decl.Args()) != len(x)-1 {
+		w.errs = append(w.errs, ArityFormatMismatchError(x[1:], x[0].String(), loc, bi.Decl))
+		return comments
+	}
+
 	comments = w.writeTermParens(true, x[1], comments)
 	w.write(" " + bi.Infix + " ")
 	comments = w.writeTermParens(true, x[2], comments)
@@ -786,7 +806,16 @@ func (w *writer) writeCall(parens bool, x ast.Call, comments []*ast.Comment) []*
 	return comments
 }
 
-func (w *writer) writeInOperator(parens bool, operands []*ast.Term, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeInOperator(parens bool, operands []*ast.Term, comments []*ast.Comment, loc *ast.Location, f *types.Function) []*ast.Comment {
+	if len(operands) != len(f.Args()) {
+		// The number of operands does not math the arity of the `in` operator
+		operator := ast.Member.Name
+		if len(f.Args()) == 3 {
+			operator = ast.MemberWithKey.Name
+		}
+		w.errs = append(w.errs, ArityFormatMismatchError(operands, operator, loc, f))
+		return comments
+	}
 	kw := "in"
 	switch len(operands) {
 	case 2:
@@ -1356,3 +1385,36 @@ func ensureFutureKeywordImport(imps []*ast.Import, kw string) []*ast.Import {
 	imp.Location = defaultLocation(imp)
 	return append(imps, imp)
 }
+
+// ArgErrDetail but for `fmt` checks since compiler has not run yet.
+type ArityFormatErrDetail struct {
+	Have []string `json:"have"`
+	Want []string `json:"want"`
+}
+
+// arityMismatchError but for `fmt` checks since the compiler has not run yet.
+func ArityFormatMismatchError(operands []*ast.Term, operator string, loc *ast.Location, f *types.Function) *ast.Error {
+	want := make([]string, len(f.Args()))
+	for i := range f.Args() {
+		want[i] = types.Sprint(f.Args()[i])
+	}
+
+	have := make([]string, len(operands))
+	for i := 0; i < len(operands); i++ {
+		have[i] = ast.TypeName(operands[i].Value)
+	}
+	err := ast.NewError(ast.TypeErr, loc, "%s: %s", operator, "arity mismatch")
+	err.Details = &ArityFormatErrDetail{
+		Have: have,
+		Want: want,
+	}
+	return err
+}
+
+// Lines returns the string representation of the detail.
+func (d *ArityFormatErrDetail) Lines() []string {
+	return []string{
+		"have: " + "(" + strings.Join(d.Have, ",") + ")",
+		"want: " + "(" + strings.Join(d.Want, ",") + ")",
+	}
+}