From 8c49c01d83649ec8070db7e8e40cb7d4d4bdf08f Mon Sep 17 00:00:00 2001 From: Magnus Jungsbluth Date: Mon, 14 Oct 2024 11:01:44 +0200 Subject: [PATCH] refactor creation of prepared queries Before this refactoring only the PreparedQuery options could be passed when using this plugin as a library. There are two more options (prepare options and evaluation options) that can be passed when using the OPA library directly. This change breaks the interface but also allows for extensibility of creating the prepared queries in the future. Signed-off-by: Magnus Jungsbluth --- envoyauth/evaluation.go | 60 +++++++++++++++++------------------- envoyauth/evaluation_test.go | 12 ++++++++ internal/internal.go | 12 ++++++++ 3 files changed, 52 insertions(+), 32 deletions(-) diff --git a/envoyauth/evaluation.go b/envoyauth/evaluation.go index 719eb7a00..31d726841 100644 --- a/envoyauth/evaluation.go +++ b/envoyauth/evaluation.go @@ -3,13 +3,11 @@ package envoyauth import ( "context" "fmt" - "sync" "github.com/open-policy-agent/opa/ast" "github.com/open-policy-agent/opa/bundle" "github.com/open-policy-agent/opa/config" "github.com/open-policy-agent/opa/logging" - "github.com/open-policy-agent/opa/metrics" "github.com/open-policy-agent/opa/rego" "github.com/open-policy-agent/opa/storage" "github.com/open-policy-agent/opa/topdown/builtins" @@ -25,17 +23,21 @@ type EvalContext interface { Store() storage.Store Compiler() *ast.Compiler Runtime() *ast.Term - PreparedQueryDoOnce() *sync.Once InterQueryBuiltinCache() iCache.InterQueryCache - PreparedQuery() *rego.PreparedEvalQuery - SetPreparedQuery(*rego.PreparedEvalQuery) Logger() logging.Logger Config() *config.Config DistributedTracing() tracing.Options + CreatePreparedQueryOnce(opts PrepareQueryOpts) (*rego.PreparedEvalQuery, error) +} + +// PrepareQueryOpts - Options to prepare a Rego query to be passed to the CreatePreparedQueryOnce method +type PrepareQueryOpts struct { + Opts []func(*rego.Rego) + PrepareOpts []rego.PrepareOption } // Eval - Evaluates an input against a provided EvalContext and yields result -func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result *EvalResult, opts ...func(*rego.Rego)) error { +func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result *EvalResult, evalOpts ...rego.EvalOption) error { var err error logger := evalContext.Logger() @@ -64,7 +66,19 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result "txn": result.TxnID, }).Debug("Executing policy query.") - err = constructPreparedQuery(evalContext, result.Txn, result.Metrics, opts) + pq, err := evalContext.CreatePreparedQueryOnce( + PrepareQueryOpts{ + Opts: []func(*rego.Rego){ + rego.Metrics(result.Metrics), + rego.ParsedQuery(evalContext.ParsedQuery()), + rego.Compiler(evalContext.Compiler()), + rego.Store(evalContext.Store()), + rego.Transaction(result.Txn), + rego.Runtime(evalContext.Runtime()), + rego.EnablePrintStatements(true), + rego.DistributedTracingOpts(evalContext.DistributedTracing()), + }, + }) if err != nil { return err } @@ -76,9 +90,7 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result ndbCache = builtins.NDBCache{} } - var rs rego.ResultSet - rs, err = evalContext.PreparedQuery().Eval( - ctx, + evalOpts = append(evalOpts, rego.EvalParsedInput(input), rego.EvalTransaction(result.Txn), rego.EvalMetrics(result.Metrics), @@ -87,6 +99,12 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result rego.EvalNDBuiltinCache(ndbCache), ) + var rs rego.ResultSet + rs, err = pq.Eval( + ctx, + evalOpts..., + ) + switch { case err != nil: return err @@ -101,28 +119,6 @@ func Eval(ctx context.Context, evalContext EvalContext, input ast.Value, result return nil } -func constructPreparedQuery(evalContext EvalContext, txn storage.Transaction, m metrics.Metrics, opts []func(*rego.Rego)) error { - var err error - var pq rego.PreparedEvalQuery - evalContext.PreparedQueryDoOnce().Do(func() { - opts = append(opts, - rego.Metrics(m), - rego.ParsedQuery(evalContext.ParsedQuery()), - rego.Compiler(evalContext.Compiler()), - rego.Store(evalContext.Store()), - rego.Transaction(txn), - rego.Runtime(evalContext.Runtime()), - rego.EnablePrintStatements(true), - rego.DistributedTracingOpts(evalContext.DistributedTracing()), - ) - - pq, err = rego.New(opts...).PrepareForEval(context.Background()) - evalContext.SetPreparedQuery(&pq) - }) - - return err -} - func getRevision(ctx context.Context, store storage.Store, txn storage.Transaction, result *EvalResult) error { revisions := map[string]string{} diff --git a/envoyauth/evaluation_test.go b/envoyauth/evaluation_test.go index 496bde6fc..6514dcb61 100644 --- a/envoyauth/evaluation_test.go +++ b/envoyauth/evaluation_test.go @@ -254,6 +254,7 @@ type mockExtAuthzGrpcServer struct { manager *plugins.Manager preparedQuery *rego.PreparedEvalQuery preparedQueryDoOnce *sync.Once + preparedQueryErr error distributedTracingOpts tracing.Options } @@ -301,6 +302,17 @@ func (m *mockExtAuthzGrpcServer) DistributedTracing() tracing.Options { return m.distributedTracingOpts } +func (m *mockExtAuthzGrpcServer) CreatePreparedQueryOnce(opts PrepareQueryOpts) (*rego.PreparedEvalQuery, error) { + m.preparedQueryDoOnce.Do(func() { + pq, err := rego.New(opts.Opts...).PrepareForEval(context.Background()) + + m.preparedQuery = &pq + m.preparedQueryErr = err + }) + + return m.preparedQuery, m.preparedQueryErr +} + type testPlugin struct { events []logs.EventV1 } diff --git a/internal/internal.go b/internal/internal.go index 8e72cc808..b5dffdf53 100644 --- a/internal/internal.go +++ b/internal/internal.go @@ -225,6 +225,7 @@ type envoyExtAuthzGrpcServer struct { manager *plugins.Manager preparedQuery *rego.PreparedEvalQuery preparedQueryDoOnce *sync.Once + preparedQueryErr error interQueryBuiltinCache iCache.InterQueryCache distributedTracingOpts tracing.Options metricAuthzDuration prometheus.HistogramVec @@ -279,6 +280,17 @@ func (p *envoyExtAuthzGrpcServer) DistributedTracing() tracing.Options { return p.distributedTracingOpts } +func (p *envoyExtAuthzGrpcServer) CreatePreparedQueryOnce(opts envoyauth.PrepareQueryOpts) (*rego.PreparedEvalQuery, error) { + p.preparedQueryDoOnce.Do(func() { + pq, err := rego.New(opts.Opts...).PrepareForEval(context.Background()) + + p.preparedQuery = &pq + p.preparedQueryErr = err + }) + + return p.preparedQuery, p.preparedQueryErr +} + func (p *envoyExtAuthzGrpcServer) Start(ctx context.Context) error { p.manager.UpdatePluginStatus(PluginName, &plugins.Status{State: plugins.StateNotReady}) go p.listen()