core/authorize: add support for rego print statements (#5049)

authorize/evaluator/headers_evaluator.go

@@ -105,6 +105,7 @@ func NewHeadersEvaluator(ctx context.Context, store *store.Store) (*HeadersEvalu
 		rego.Store(store),
 		rego.Module("pomerium.headers", opa.HeadersRego),
 		rego.Query("result := data.pomerium.headers"),
+		rego.EnablePrintStatements(true),
 		getGoogleCloudServerlessHeadersRegoOption,
 		variableSubstitutionFunctionRegoOption,
 		store.GetDataBrokerRecordOption(),

authorize/evaluator/log.go

@@ -0,0 +1,19 @@
+package evaluator
+
+import (
+	"github.com/open-policy-agent/opa/topdown/print"
+	"github.com/rs/zerolog"
+)
+
+type regoPrintHook struct {
+	logger zerolog.Logger
+}
+
+var _ print.Hook = (*regoPrintHook)(nil)
+
+func (h regoPrintHook) Print(ctx print.Context, msg string) error {
+	h.logger.Debug().
+		Any("location", ctx.Location).
+		Msg("rego: " + msg)
+	return nil
+}

authorize/evaluator/log_test.go

@@ -0,0 +1,49 @@
+package evaluator
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPrintHook(t *testing.T) {
+	t.Parallel()
+
+	ctx, clearTimeout := context.WithTimeout(context.Background(), time.Second*10)
+	t.Cleanup(clearTimeout)
+
+	r := rego.New(
+		rego.Module("policy.rego", `
+package pomerium.policy
+
+import rego.v1
+
+allow if {
+	print("HELLO WORLD")
+	true
+}
+		`),
+		rego.EnablePrintStatements(true),
+		rego.Query("data.pomerium.policy.allow"),
+	)
+	q, err := r.PrepareForEval(ctx)
+	require.NoError(t, err)
+
+	var buf bytes.Buffer
+	logger := zerolog.New(&buf).Level(zerolog.DebugLevel)
+
+	rs, err := q.Eval(ctx, rego.EvalPrintHook(regoPrintHook{
+		logger: logger,
+	}))
+	require.NoError(t, err)
+	assert.True(t, rs.Allowed())
+
+	assert.Equal(t, `{"level":"debug","location":{"file":"policy.rego","row":7,"col":2},"message":"rego: HELLO WORLD"}`, strings.TrimSpace(buf.String()))
+}

authorize/evaluator/policy_evaluator.go

@@ -157,6 +157,7 @@ func NewPolicyEvaluator(
 			rego.Store(store),
 			rego.Module("pomerium.policy", e.queries[i].script),
 			rego.Query("result = data.pomerium.policy"),
+			rego.EnablePrintStatements(true),
 			getGoogleCloudServerlessHeadersRegoOption,
 			store.GetDataBrokerRecordOption(),
 		)
@@ -168,6 +169,7 @@ func NewPolicyEvaluator(
 				rego.Store(store),
 				rego.Module("pomerium.policy", "package pomerium.policy\n\n"+e.queries[i].script),
 				rego.Query("result = data.pomerium.policy"),
+				rego.EnablePrintStatements(true),
 				getGoogleCloudServerlessHeadersRegoOption,
 				store.GetDataBrokerRecordOption(),
 			)
@@ -210,7 +212,11 @@ func (e *PolicyEvaluator) evaluateQuery(ctx context.Context, req *PolicyRequest,
 	defer span.End()
 	span.AddAttributes(octrace.StringAttribute("script_checksum", query.checksum()))
 
-	rs, err := safeEval(ctx, query.PreparedEvalQuery, rego.EvalInput(req))
+	rs, err := safeEval(ctx, query.PreparedEvalQuery,
+		rego.EvalInput(req),
+		rego.EvalPrintHook(regoPrintHook{
+			logger: *log.Logger(),
+		}))
 	if err != nil {
 		return nil, fmt.Errorf("authorize: error evaluating policy.rego: %w", err)
 	}