From 6c6e929b9329b8fc05abd6a058e701f2f5ba5cd0 Mon Sep 17 00:00:00 2001
From: James Chacon <james.chacon@snowflake.com>
Date: Thu, 3 Mar 2022 14:14:33 -0800
Subject: [PATCH] Add print() statement support to policy. (#83)

* Add print() statement support to policy.

Just info log at v1 if we generate any data.

Convert input to json before we debug log it in rpcauth. Otherwise what rego sees is not what we think.

* Don't marshal unless we're logging at this level

* Update tests to add logging to trigger coverage.
---
 auth/opa/opa.go                  | 16 +++++++++++++++-
 auth/opa/rpcauth/rpcauth.go      | 19 +++++++++++++++++--
 auth/opa/rpcauth/rpcauth_test.go |  8 ++++++++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/auth/opa/opa.go b/auth/opa/opa.go
index 6d66bf69..a3dc15ef 100644
--- a/auth/opa/opa.go
+++ b/auth/opa/opa.go
@@ -19,11 +19,14 @@
 package opa
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 
+	"github.com/go-logr/logr"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/topdown"
 )
 
 const (
@@ -47,6 +50,7 @@ var (
 // a sansshell rego policy file.
 type AuthzPolicy struct {
 	query rego.PreparedEvalQuery
+	b     *bytes.Buffer
 }
 
 type policyOptions struct {
@@ -94,25 +98,35 @@ func NewAuthzPolicy(ctx context.Context, policy string, opts ...Option) (*AuthzP
 		return nil, fmt.Errorf("policy has invalid package '%s' (must be '%s')", module.Package, sansshellPackage)
 	}
 
+	b := &bytes.Buffer{}
 	r := rego.New(
 		rego.Query(options.query),
 		rego.ParsedModule(module),
+		rego.EnablePrintStatements(true),
+		rego.PrintHook(topdown.NewPrintHook(b)),
 	)
 
 	prepared, err := r.PrepareForEval(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("rego: PrepareForEval() error: %w", err)
 	}
-	return &AuthzPolicy{query: prepared}, nil
+	return &AuthzPolicy{
+		query: prepared,
+		b:     b,
+	}, nil
 }
 
 // Eval evaluates this policy using the provided input, returning 'true'
 // iff the evaulation was successful, and the operation represented by
 // `input` is permitted by the policy.
 func (q *AuthzPolicy) Eval(ctx context.Context, input interface{}) (bool, error) {
+	logger := logr.FromContextOrDiscard(ctx)
 	results, err := q.query.Eval(ctx, rego.EvalInput(input))
 	if err != nil {
 		return false, fmt.Errorf("authz policy evaluation error: %w", err)
 	}
+	if q.b.Len() > 0 {
+		logger.V(1).Info("print statements", "buffer", q.b.String())
+	}
 	return results.Allowed(), nil
 }
diff --git a/auth/opa/rpcauth/rpcauth.go b/auth/opa/rpcauth/rpcauth.go
index b3a40888..5ed42f36 100644
--- a/auth/opa/rpcauth/rpcauth.go
+++ b/auth/opa/rpcauth/rpcauth.go
@@ -20,6 +20,7 @@ package rpcauth
 
 import (
 	"context"
+	"encoding/json"
 
 	"github.com/go-logr/logr"
 	"google.golang.org/grpc"
@@ -75,7 +76,14 @@ func NewWithPolicy(ctx context.Context, policy string, authzHooks ...RPCAuthzHoo
 func (g *Authorizer) Eval(ctx context.Context, input *RPCAuthInput) error {
 	logger := logr.FromContextOrDiscard(ctx)
 	if input != nil {
-		logger.V(1).Info("evaluating authz policy", "input.message", string(input.Message), "input", input)
+		if logger.V(2).Enabled() {
+			b, err := json.Marshal(input)
+			if err != nil {
+				logger.V(2).Info("marshal", "can't marshal input", err)
+			} else {
+				logger.V(2).Info("evaluating authz policy", "input", string(b))
+			}
+		}
 	}
 	if input == nil {
 		return status.Error(codes.InvalidArgument, "policy input cannot be nil")
@@ -89,7 +97,14 @@ func (g *Authorizer) Eval(ctx context.Context, input *RPCAuthInput) error {
 			return status.Errorf(codes.Internal, "authz hook error: %v", err)
 		}
 	}
-	logger.V(1).Info("evaluating authz policy post hooks", "input.message", string(input.Message), "input", input)
+	if logger.V(1).Enabled() {
+		b, err := json.Marshal(input)
+		if err != nil {
+			logger.V(1).Info("marshal", "can't marshal input", err)
+		} else {
+			logger.V(1).Info("evaluating authz policy post hooks", "input", string(b))
+		}
+	}
 	allowed, err := g.policy.Eval(ctx, input)
 	if err != nil {
 		return status.Errorf(codes.Internal, "authz policy evaluation error: %v", err)
diff --git a/auth/opa/rpcauth/rpcauth_test.go b/auth/opa/rpcauth/rpcauth_test.go
index a559b6c1..705ecd3b 100644
--- a/auth/opa/rpcauth/rpcauth_test.go
+++ b/auth/opa/rpcauth/rpcauth_test.go
@@ -22,8 +22,10 @@ import (
 	"crypto/x509"
 	"encoding/json"
 	"errors"
+	"log"
 	"net"
 	"net/url"
+	"os"
 	"testing"
 
 	"google.golang.org/grpc"
@@ -37,6 +39,8 @@ import (
 
 	"github.com/Snowflake-Labs/sansshell/auth/opa"
 	"github.com/Snowflake-Labs/sansshell/testing/testutil"
+	"github.com/go-logr/logr"
+	"github.com/go-logr/stdr"
 )
 
 var policyString = `
@@ -71,6 +75,10 @@ allow {
 
 func TestAuthzHook(t *testing.T) {
 	ctx := context.Background()
+	logger := stdr.New(log.New(os.Stderr, "", log.LstdFlags|log.Lshortfile))
+	ctx = logr.NewContext(ctx, logger)
+	// This way the tests exercise the logging code.
+	stdr.SetVerbosity(2)
 	policy, err := opa.NewAuthzPolicy(ctx, policyString)
 	testutil.FatalOnErr("NewAuthzPolicy", err, t)