feat: audit trail in stdout logs

core/opaevaluator.go

@@ -16,12 +16,14 @@ package core
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 	"time"
 
 	"github.com/rond-authz/rond/custom_builtins"
+	"github.com/rond-authz/rond/internal/audit"
 	"github.com/rond-authz/rond/internal/opatranslator"
 	"github.com/rond-authz/rond/internal/utils"
 	"github.com/rond-authz/rond/logging"
@@ -33,6 +35,10 @@ import (
 	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
+var (
+	ErrModuleCompilationFailed = errors.New("failed module compilation")
+)
+
 type RondConfig struct {
 	RequestFlow  RequestFlow       `json:"requestFlow"`
 	ResponseFlow ResponseFlow      `json:"responseFlow"`
@@ -75,6 +81,7 @@ type OPAEvaluatorOptions struct {
 	EnablePrintStatements bool
 	MongoClient           custom_builtins.IMongoClient
 	Logger                logging.Logger
+	Builtins              []func(*rego.Rego)
 }
 
 func (opaEval *OPAEvaluator) partiallyEvaluate(logger logging.Logger, input EvalInput, options *PolicyEvaluationOptions) (primitive.M, error) {
@@ -230,12 +237,21 @@ type Module struct {
 	Content string
 }
 
-func NewOPAModuleConfig(modules []Module) (*OPAModuleConfig, error) {
-	compiler := ast.NewCompiler().WithBuiltins(map[string]*ast.Builtin{
+type BuiltinDeclarations map[string]*ast.Builtin
+
+func NewOPAModuleConfigWithBuiltins(modules []Module, builtinsDeclarations BuiltinDeclarations) (*OPAModuleConfig, error) {
+	builtins := map[string]*ast.Builtin{
 		custom_builtins.GetHeaderDecl.Name:     custom_builtins.GetHeaderDecl,
 		custom_builtins.MongoFindOneDecl.Name:  custom_builtins.MongoFindOneDecl,
 		custom_builtins.MongoFindManyDecl.Name: custom_builtins.MongoFindManyDecl,
-	})
+		audit.SetLabelsDecl.Name:               audit.SetLabelsDecl,
+	}
+
+	for name, decl := range builtinsDeclarations {
+		builtins[name] = decl
+	}
+
+	compiler := ast.NewCompiler().WithBuiltins(builtins)
 
 	modulesToCompile := map[string]*ast.Module{}
 	for _, m := range modules {
@@ -245,7 +261,7 @@ func NewOPAModuleConfig(modules []Module) (*OPAModuleConfig, error) {
 	compiler.Compile(modulesToCompile)
 
 	if compiler.Failed() {
-		return nil, fmt.Errorf("fails to compile the module: %s", compiler.Errors)
+		return nil, fmt.Errorf("%w: %s", ErrModuleCompilationFailed, compiler.Errors)
 	}
 
 	return &OPAModuleConfig{
@@ -254,6 +270,11 @@ func NewOPAModuleConfig(modules []Module) (*OPAModuleConfig, error) {
 	}, nil
 }
 
+func NewOPAModuleConfig(modules []Module) (*OPAModuleConfig, error) {
+	return NewOPAModuleConfigWithBuiltins(modules, nil)
+}
+
+// MustNewOPAModuleConfig is an helper function for tests purposes only!!
 func MustNewOPAModuleConfig(modules []Module) *OPAModuleConfig {
 	opaModule, err := NewOPAModuleConfig(modules)
 	if err != nil {
@@ -262,6 +283,15 @@ func MustNewOPAModuleConfig(modules []Module) *OPAModuleConfig {
 	return opaModule
 }
 
+// MustNewOPAModuleConfigWithBuiltins is an helper function for tests purposes only!!
+func MustNewOPAModuleConfigWithBuiltins(modules []Module, builtins BuiltinDeclarations) *OPAModuleConfig {
+	opaModule, err := NewOPAModuleConfigWithBuiltins(modules, builtins)
+	if err != nil {
+		panic(err)
+	}
+	return opaModule
+}
+
 type PermissionOnResourceKey string
 
 type PermissionsOnResourceMap map[PermissionOnResourceKey]bool

core/partialevaluators.go

@@ -48,7 +48,7 @@ func (policyEvaluators PartialResultsEvaluators) AddFromConfig(ctx context.Conte
 	}
 
 	if _, ok := policyEvaluators[allowPolicy]; !ok {
-		evaluator, err := createPartialEvaluator(ctx, logger, allowPolicy, opaModuleConfig, options, isFilterQuery)
+		evaluator, err := createEvaluator(ctx, logger, allowPolicy, opaModuleConfig, options, isFilterQuery)
 		if err != nil {
 			return fmt.Errorf("%w: %s", ErrEvaluatorCreationFailed, err.Error())
 		}
@@ -57,7 +57,7 @@ func (policyEvaluators PartialResultsEvaluators) AddFromConfig(ctx context.Conte
 
 	if responsePolicy != "" {
 		if _, ok := policyEvaluators[responsePolicy]; !ok {
-			evaluator, err := createPartialEvaluator(ctx, logger, responsePolicy, opaModuleConfig, options, false)
+			evaluator, err := createEvaluator(ctx, logger, responsePolicy, opaModuleConfig, options, false)
 			if err != nil {
 				return fmt.Errorf("%w: %s", ErrEvaluatorCreationFailed, err.Error())
 			}
@@ -87,7 +87,7 @@ func (partialEvaluators PartialResultsEvaluators) GetEvaluatorFromPolicy(ctx con
 	return nil, fmt.Errorf("%w: %s", ErrEvaluatorNotFound, policy)
 }
 
-func createPartialEvaluator(ctx context.Context, logger logging.Logger, policy string, opaModuleConfig *OPAModuleConfig, options *OPAEvaluatorOptions, isPartial bool) (*PartialEvaluator, error) {
+func createEvaluator(ctx context.Context, logger logging.Logger, policy string, opaModuleConfig *OPAModuleConfig, options *OPAEvaluatorOptions, isPartial bool) (*PartialEvaluator, error) {
 	logger.WithField("policyName", policy).Info("precomputing rego policy")
 
 	preparedPartialEvaluator := &PartialEvaluator{}

core/partialevaluators_test.go

@@ -24,6 +24,9 @@ import (
 	"github.com/rond-authz/rond/logging"
 	"github.com/rond-authz/rond/metrics"
 
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -406,4 +409,61 @@ func TestPartialResultEvaluators(t *testing.T) {
 		require.Empty(t, query)
 		require.Empty(t, res)
 	})
+
+	t.Run("evaluates policy with custom builtins", func(t *testing.T) {
+		partialEvaluators := PartialResultsEvaluators{}
+		rondConfig := &RondConfig{
+			RequestFlow: RequestFlow{
+				PolicyName: "check_metadata",
+			},
+		}
+
+		builtinDecl := &ast.Builtin{
+			Name: "some_builtin",
+			Decl: types.NewFunction(types.Args(types.A), types.A),
+		}
+
+		opaModule := MustNewOPAModuleConfigWithBuiltins([]Module{
+			{
+				Name: "example.rego",
+				Content: `package policies
+			check_metadata {
+				some_builtin("test")
+			}`,
+			},
+		}, BuiltinDeclarations{
+			"some_builtin": builtinDecl,
+		})
+
+		someBuiltinInvoked := false
+		builtins := []func(*rego.Rego){
+			rego.Function1(
+				&rego.Function{
+					Name: "some_builtin",
+					Decl: types.NewFunction(types.Args(types.A), types.A),
+				},
+				func(bctx rego.BuiltinContext, dataToStore *ast.Term) (*ast.Term, error) {
+					someBuiltinInvoked = true
+					return ast.BooleanTerm(true), nil
+				},
+			),
+		}
+		evalOpts := OPAEvaluatorOptions{
+			Logger:   logger,
+			Builtins: builtins,
+		}
+
+		err := partialEvaluators.AddFromConfig(context.Background(), logger, opaModule, rondConfig, &evalOpts)
+		require.NoError(t, err)
+
+		input, err := CreateRegoQueryInput(logger, rondInput, RegoInputOptions{})
+		require.NoError(t, err)
+		evaluator, err := partialEvaluators.GetEvaluatorFromPolicy(context.Background(), "check_metadata", &evalOpts)
+		require.NoError(t, err)
+		res, query, err := evaluator.PolicyEvaluation(logger, input, nil)
+		require.NoError(t, err)
+		require.Empty(t, query)
+		require.Empty(t, res)
+		require.True(t, someBuiltinInvoked)
+	})
 }

core/rego.go

@@ -22,6 +22,7 @@ import (
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/rond-authz/rond/custom_builtins"
+	"github.com/rond-authz/rond/internal/audit"
 )
 
 func newRegoInstanceBuilder(policy string, opaModuleConfig *OPAModuleConfig, evaluatorOptions *OPAEvaluatorOptions) (*rego.Rego, error) {
@@ -44,7 +45,13 @@ func newRegoInstanceBuilder(policy string, opaModuleConfig *OPAModuleConfig, eva
 		rego.PrintHook(NewPrintHook(os.Stdout, policy)),
 		rego.Capabilities(ast.CapabilitiesForThisVersion()),
 		custom_builtins.GetHeaderFunction,
+		audit.SetLabels,
 	}
+
+	if len(evaluatorOptions.Builtins) > 0 {
+		options = append(options, evaluatorOptions.Builtins...)
+	}
+
 	if evaluatorOptions.MongoClient != nil {
 		options = append(options, custom_builtins.MongoFindOne, custom_builtins.MongoFindMany)
 	}

internal/audit/agent.go

@@ -0,0 +1,34 @@
+// Copyright 2024 Mia srl
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package audit
+
+import "context"
+
+type Labels = map[string]any
+
+type Agent interface {
+	Trace(context.Context, Audit)
+	Cache() AuditCache
+	SetGlobalLabels(labels Labels)
+}
+
+// noopAgent is a lazy agent that does nothing :(
+type noopAgent struct{}
+
+func NewNoopAgent() Agent { return &noopAgent{} }
+
+func (a *noopAgent) Trace(context.Context, Audit)  {}
+func (a *noopAgent) Cache() AuditCache             { return &SingleRecordCache{} }
+func (a *noopAgent) SetGlobalLabels(labels Labels) {}

internal/audit/agent_log.go

@@ -0,0 +1,59 @@
+// Copyright 2024 Mia srl
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package audit
+
+import (
+	"context"
+
+	"github.com/rond-authz/rond/internal/utils"
+	"github.com/rond-authz/rond/logging"
+)
+
+type logAgent struct {
+	l       logging.Logger
+	cache   AuditCache
+	globals map[string]any
+}
+
+func NewLogAgent(l logging.Logger) Agent {
+	return &logAgent{
+		l:     l,
+		cache: &SingleRecordCache{},
+	}
+}
+func (a *logAgent) SetGlobalLabels(labels Labels) {
+	a.globals = labels
+}
+
+func (a *logAgent) Trace(_ context.Context, auditInput Audit) {
+	data := a.cache.Load()
+
+	auditData := auditInput.toPrint()
+	if a.globals != nil {
+		auditData.applyDataFromPolicy(a.globals)
+	}
+
+	if data != nil {
+		auditData.applyDataFromPolicy(data)
+	}
+
+	a.l.
+		WithField("trail", utils.ToMap(auditSerializerTagAnnotation, auditData)).
+		Info("audit trail")
+}
+
+func (a *logAgent) Cache() AuditCache {
+	return a.cache
+}