From ea035a0f8a3508a5561c8898c24e4d99f3572310 Mon Sep 17 00:00:00 2001
From: Charlie Egan <charlieegan3@users.noreply.github.com>
Date: Tue, 6 Dec 2022 09:09:01 +0000
Subject: [PATCH] Add StrictBuiltinErrors option to Partial

Signed-off-by: Charlie Egan <charlieegan3@users.noreply.github.com>
---
 sdk/opa.go      | 54 +++++++++++++++++++-----------------
 sdk/opa_test.go | 73 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 25 deletions(-)

diff --git a/sdk/opa.go b/sdk/opa.go
index a1f39adc6e..69175a5a4d 100644
--- a/sdk/opa.go
+++ b/sdk/opa.go
@@ -353,16 +353,17 @@ func (opa *OPA) Partial(ctx context.Context, options PartialOptions) (*PartialRe
 		&record,
 		func(s state, result *DecisionResult) {
 			pq, record.InputAST, record.Bundles, record.Error = partial(ctx, partialEvalArgs{
-				runtime:   s.manager.Info,
-				printHook: s.manager.PrintHook(),
-				compiler:  s.manager.GetCompiler(),
-				store:     s.manager.Store,
-				txn:       record.Txn,
-				now:       record.Timestamp,
-				query:     record.Query,
-				unknowns:  options.Unknowns,
-				input:     *record.Input,
-				m:         record.Metrics,
+				runtime:             s.manager.Info,
+				printHook:           s.manager.PrintHook(),
+				compiler:            s.manager.GetCompiler(),
+				store:               s.manager.Store,
+				txn:                 record.Txn,
+				now:                 record.Timestamp,
+				query:               record.Query,
+				unknowns:            options.Unknowns,
+				input:               *record.Input,
+				m:                   record.Metrics,
+				strictBuiltinErrors: options.StrictBuiltinErrors,
 			})
 			if record.Error == nil {
 				result.Result, record.Error = options.Mapper.MapResults(pq)
@@ -397,11 +398,12 @@ type PartialQueryMapper interface {
 
 // PartialOptions contains parameters for partial query evaluation.
 type PartialOptions struct {
-	Now      time.Time          // specifies wallclock time used for time.now_ns(), decision log timestamp, etc.
-	Input    interface{}        // specifies value of the input document to evaluate policy with
-	Query    string             // specifies the query to be partially evaluated
-	Unknowns []string           // specifies the unknown elements of the policy
-	Mapper   PartialQueryMapper // specifies the mapper to use when processing results
+	Now                 time.Time          // specifies wallclock time used for time.now_ns(), decision log timestamp, etc.
+	Input               interface{}        // specifies value of the input document to evaluate policy with
+	Query               string             // specifies the query to be partially evaluated
+	Unknowns            []string           // specifies the unknown elements of the policy
+	Mapper              PartialQueryMapper // specifies the mapper to use when processing results
+	StrictBuiltinErrors bool               // treat built-in function errors as fatal
 }
 
 type PartialResult struct {
@@ -510,16 +512,17 @@ func evaluate(ctx context.Context, args evalArgs) (interface{}, ast.Value, map[s
 }
 
 type partialEvalArgs struct {
-	runtime   *ast.Term
-	compiler  *ast.Compiler
-	printHook print.Hook
-	store     storage.Store
-	txn       storage.Transaction
-	unknowns  []string
-	query     string
-	now       time.Time
-	input     interface{}
-	m         metrics.Metrics
+	runtime             *ast.Term
+	compiler            *ast.Compiler
+	printHook           print.Hook
+	store               storage.Store
+	txn                 storage.Transaction
+	unknowns            []string
+	query               string
+	now                 time.Time
+	input               interface{}
+	m                   metrics.Metrics
+	strictBuiltinErrors bool
 }
 
 func partial(ctx context.Context, args partialEvalArgs) (*rego.PartialQueries, ast.Value, map[string]server.BundleInfo, error) {
@@ -544,6 +547,7 @@ func partial(ctx context.Context, args partialEvalArgs) (*rego.PartialQueries, a
 		rego.Query(args.query),
 		rego.Unknowns(args.unknowns),
 		rego.PrintHook(args.printHook),
+		rego.StrictBuiltinErrors(args.strictBuiltinErrors),
 	)
 
 	pq, err := re.Partial(ctx)
diff --git a/sdk/opa_test.go b/sdk/opa_test.go
index 71d91f7f66..833be5139d 100644
--- a/sdk/opa_test.go
+++ b/sdk/opa_test.go
@@ -356,6 +356,79 @@ func TestPartial(t *testing.T) {
 
 }
 
+func TestPartialWithStrictBuiltinErrors(t *testing.T) {
+
+	ctx := context.Background()
+
+	server := sdktest.MustNewServer(
+		sdktest.MockBundle("/bundles/bundle.tar.gz", map[string]string{
+			"main.rego": `
+				package example
+
+erroring_function(number) = output {
+	output := number / 0
+}
+
+allow {
+	data.unknown.x 
+	erroring_function(input.x) 
+}`,
+		}),
+	)
+
+	defer server.Stop()
+
+	config := fmt.Sprintf(`{
+		"services": {
+			"test": {
+				"url": %q
+			}
+		},
+		"bundles": {
+			"test": {
+				"resource": "/bundles/bundle.tar.gz"
+			}
+		},
+		"decision_logs": {
+			"console": true
+		}
+	}`, server.URL())
+
+	testLogger := loggingtest.New()
+	opa, err := sdk.New(ctx, sdk.Options{
+		Config:        strings.NewReader(config),
+		ConsoleLogger: testLogger,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer opa.Stop(ctx)
+
+	_, err = opa.Partial(ctx, sdk.PartialOptions{
+		Input:               map[string]int{"x": 1},
+		Query:               "data.example.allow",
+		Unknowns:            []string{"data.unknown.x"},
+		Mapper:              &sdk.RawMapper{},
+		Now:                 time.Unix(0, 1619868194450288000).UTC(),
+		StrictBuiltinErrors: true,
+	})
+	if err == nil {
+		t.Fatal("expected error but got nil")
+	}
+
+	actual, ok := err.(*topdown.Error)
+	if !ok || actual.Code != "eval_builtin_error" {
+		t.Fatalf("expected eval_builtin_error but got %v", actual)
+	}
+
+	expectedMessage := "div: divide by zero"
+	if actual.Message != expectedMessage {
+		t.Fatalf("expected %v but got %v", expectedMessage, actual.Message)
+	}
+
+}
+
 func TestUndefinedError(t *testing.T) {
 
 	ctx := context.Background()