rego+topdown: allow custom builtins via rego pkg to halt

Before, functions provided via arguments to `rego.New()` hadn't been
able to halt the topdown evaluation.

Now, they do, if they return a `*rego.HaltError`.

Fixes open-policy-agent#3534.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

rego/errors.go

@@ -0,0 +1,18 @@
+package rego
+
+// HaltError is an error type to return from a custom function implementation
+// that will abort the evaluation process (analogous to topdown.Halt).
+type HaltError struct {
+	err error
+}
+
+// Error delegates to the wrapped error
+func (h *HaltError) Error() string {
+	return h.err.Error()
+}
+
+// NewHaltError wraps an error such that the evaluation process will stop
+// when it occurs.
+func NewHaltError(err error) error {
+	return &HaltError{err: err}
+}

rego/rego.go

@@ -2457,6 +2457,14 @@ func parseStringsToRefs(s []string) ([]ast.Ref, error) {
 // was defined.
 func finishFunction(name string, bctx topdown.BuiltinContext, result *ast.Term, err error, iter func(*ast.Term) error) error {
 	if err != nil {
+		var e *HaltError
+		if errors.As(err, &e) {
+			return topdown.Halt{Err: &topdown.Error{
+				Code:     topdown.BuiltinErr,
+				Message:  fmt.Sprintf("%v: %v", name, e.Error()),
+				Location: bctx.Location,
+			}}
+		}
 		return &topdown.Error{
 			Code:     topdown.BuiltinErr,
 			Message:  fmt.Sprintf("%v: %v", name, err.Error()),

rego/rego_test.go

@@ -8,6 +8,8 @@ package rego
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	"fmt"
 	"log"
 	"net/http"
 	"net/http/httptest"
@@ -19,6 +21,7 @@ import (
 	"time"
 
 	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/ast/location"
 	"github.com/open-policy-agent/opa/bundle"
 	"github.com/open-policy-agent/opa/internal/storage/mock"
 	"github.com/open-policy-agent/opa/metrics"
@@ -236,8 +239,37 @@ func TestRegoCancellation(t *testing.T) {
 
 	if err == nil {
 		t.Fatalf("Expected cancellation error but got: %v", rs)
-	} else if topdownErr, ok := err.(*topdown.Error); !ok || topdownErr.Code != topdown.CancelErr {
-		t.Fatalf("Got unexpected error: %v", err)
+	}
+	exp := topdown.Error{Code: topdown.CancelErr, Message: "caller cancelled query execution"}
+	if !errors.Is(err, &exp) {
+		t.Errorf("error: expected %v, got: %v", exp, err)
+	}
+}
+
+func TestRegoCustomBuiltinHalt(t *testing.T) {
+
+	funOpt := Function1(
+		&Function{
+			Name: "halt_func",
+			Decl: types.NewFunction(
+				types.Args(types.S),
+				types.NewNull(),
+			),
+		},
+		func(BuiltinContext, *ast.Term) (*ast.Term, error) {
+			return nil, NewHaltError(fmt.Errorf("stop"))
+		},
+	)
+	r := New(Query(`halt_func("")`), funOpt)
+	rs, err := r.Eval(context.Background())
+	if err == nil {
+		t.Fatalf("Expected halt error but got: %v", rs)
+	}
+	// exp is the error topdown returns after unwrapping the Halt
+	exp := topdown.Error{Code: topdown.BuiltinErr, Message: "halt_func: stop",
+		Location: location.NewLocation([]byte(`halt_func("")`), "", 1, 1)}
+	if !errors.Is(err, &exp) {
+		t.Fatalf("error: expected %v, got: %v", exp, err)
 	}
 }
 

topdown/eval.go

@@ -2,6 +2,7 @@ package topdown
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"sort"
@@ -1528,11 +1529,12 @@ func (e evalBuiltin) eval(iter unifyIterator) error {
 	})
 
 	if err != nil {
-		if h, ok := err.(Halt); !ok {
+		var t Halt
+		if errors.As(err, &t) {
+			err = t.Err
+		} else {
 			e.e.builtinErrors.errs = append(e.e.builtinErrors.errs, err)
 			err = nil
-		} else {
-			err = h.Err
 		}
 	}