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

@@ -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
 }

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)

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)