From 5ec67ccd3a605e143e2f3df7c98caed57a91b937 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Fri, 1 Oct 2021 10:23:03 -0700
Subject: [PATCH] rego: Do not rewrite/capture void function call return values

This commit prevents the rewriting step from attempting to capture the
result of void function calls. This avoids generating invalid queries
that will fail to type check.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
---
 rego/rego.go      |  3 ++-
 rego/rego_test.go | 16 ++++++++++++++++
 types/types.go    | 17 +++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/rego/rego.go b/rego/rego.go
index 6601b42fcd..d0a09f24e4 100644
--- a/rego/rego.go
+++ b/rego/rego.go
@@ -2167,7 +2167,8 @@ func (r *Rego) rewriteQueryToCaptureValue(qc ast.QueryCompiler, query ast.Body)
 			expr.Terms = ast.Equality.Expr(terms, capture).Terms
 			r.capture[expr] = capture.Value.(ast.Var)
 		case []*ast.Term:
-			if r.compiler.GetArity(expr.Operator()) == len(terms)-1 {
+			tpe := r.compiler.TypeEnv.Get(terms[0])
+			if !types.Void(tpe) && types.Arity(tpe) == len(terms)-1 {
 				capture = r.generateTermVar()
 				expr.Terms = append(terms, capture)
 				r.capture[expr] = capture.Value.(ast.Var)
diff --git a/rego/rego_test.go b/rego/rego_test.go
index 4e5ce3b764..1f669122d1 100644
--- a/rego/rego_test.go
+++ b/rego/rego_test.go
@@ -216,6 +216,22 @@ func TestRegoRewrittenVarsCapture(t *testing.T) {
 
 }
 
+func TestRegoDoNotCaptureVoidCalls(t *testing.T) {
+
+	ctx := context.Background()
+
+	r := New(Query("print(1)"))
+
+	rs, err := r.Eval(ctx)
+	if err != nil || len(rs) != 1 {
+		t.Fatal(err, "rs:", rs)
+	}
+
+	if !rs[0].Expressions[0].Value.(bool) {
+		t.Fatal("expected expression value to be true")
+	}
+}
+
 func TestRegoCancellation(t *testing.T) {
 
 	ast.RegisterBuiltin(&ast.Builtin{
diff --git a/types/types.go b/types/types.go
index fdfc1f879c..256efcc05c 100644
--- a/types/types.go
+++ b/types/types.go
@@ -467,6 +467,23 @@ func Args(x ...Type) []Type {
 	return x
 }
 
+// Void returns true if the function has no return value. This function returns
+// false if tpe is not a function.
+func Void(x Type) bool {
+	f, ok := x.(*Function)
+	return ok && f.Result() == nil
+}
+
+// Arity returns the number of arguments in the function signature. This
+// function returns -1 if tpe is not a function.
+func Arity(x Type) int {
+	f, ok := x.(*Function)
+	if !ok {
+		return -1
+	}
+	return len(f.FuncArgs().Args)
+}
+
 // NewFunction returns a new Function object where xs[:len(xs)-1] are arguments
 // and xs[len(xs)-1] is the result type.
 func NewFunction(args []Type, result Type) *Function {