OPA: Allow prepare query optimisation via partial evaluation

config/config.go

@@ -289,14 +289,15 @@ type Config struct {
 	LuaModules *listFlag `yaml:"lua-modules"`
 	LuaSources *listFlag `yaml:"lua-sources"`
 
-	EnableOpenPolicyAgent                bool          `yaml:"enable-open-policy-agent"`
-	OpenPolicyAgentConfigTemplate        string        `yaml:"open-policy-agent-config-template"`
-	OpenPolicyAgentEnvoyMetadata         string        `yaml:"open-policy-agent-envoy-metadata"`
-	OpenPolicyAgentCleanerInterval       time.Duration `yaml:"open-policy-agent-cleaner-interval"`
-	OpenPolicyAgentStartupTimeout        time.Duration `yaml:"open-policy-agent-startup-timeout"`
-	OpenPolicyAgentRequestBodyBufferSize int64         `yaml:"open-policy-agent-request-body-buffer-size"`
-	OpenPolicyAgentMaxRequestBodySize    int64         `yaml:"open-policy-agent-max-request-body-size"`
-	OpenPolicyAgentMaxMemoryBodyParsing  int64         `yaml:"open-policy-agent-max-memory-body-parsing"`
+	EnableOpenPolicyAgent                          bool          `yaml:"enable-open-policy-agent"`
+	OpenPolicyAgentConfigTemplate                  string        `yaml:"open-policy-agent-config-template"`
+	OpenPolicyAgentEnvoyMetadata                   string        `yaml:"open-policy-agent-envoy-metadata"`
+	OpenPolicyAgentCleanerInterval                 time.Duration `yaml:"open-policy-agent-cleaner-interval"`
+	OpenPolicyAgentStartupTimeout                  time.Duration `yaml:"open-policy-agent-startup-timeout"`
+	OpenPolicyAgentRequestBodyBufferSize           int64         `yaml:"open-policy-agent-request-body-buffer-size"`
+	OpenPolicyAgentMaxRequestBodySize              int64         `yaml:"open-policy-agent-max-request-body-size"`
+	OpenPolicyAgentMaxMemoryBodyParsing            int64         `yaml:"open-policy-agent-max-memory-body-parsing"`
+	EnableOpenPolicyAgentPreevaluationOptimization bool          `yaml:"enable-open-policy-agent-preevaluation-optimization"`
 
 	PassiveHealthCheck mapFlags `yaml:"passive-health-check"`
 }
@@ -523,6 +524,7 @@ func NewConfig() *Config {
 	flag.Int64Var(&cfg.OpenPolicyAgentMaxRequestBodySize, "open-policy-agent-max-request-body-size", openpolicyagent.DefaultMaxRequestBodySize, "Maximum number of bytes from a http request body that are passed as input to the policy")
 	flag.Int64Var(&cfg.OpenPolicyAgentRequestBodyBufferSize, "open-policy-agent-request-body-buffer-size", openpolicyagent.DefaultRequestBodyBufferSize, "Read buffer size for the request body")
 	flag.Int64Var(&cfg.OpenPolicyAgentMaxMemoryBodyParsing, "open-policy-agent-max-memory-body-parsing", openpolicyagent.DefaultMaxMemoryBodyParsing, "Total number of bytes used to parse http request bodies across all requests. Once the limit is met, requests will be rejected.")
+	flag.BoolVar(&cfg.EnableOpenPolicyAgentPreevaluationOptimization, "enable-open-policy-agent-preevaluation-optimization", false, "As an optimization, apply partial evaluation during preparation of the OPA policy. This can reduce the time to evaluate the policy at runtime but consumes more memory.")
 
 	// TLS client certs
 	flag.StringVar(&cfg.ClientKeyFile, "client-tls-key", "", "TLS Key file for backend connections, multiple keys may be given comma separated - the order must match the certs")
@@ -955,14 +957,15 @@ func (c *Config) ToOptions() skipper.Options {
 		LuaModules: c.LuaModules.values,
 		LuaSources: c.LuaSources.values,
 
-		EnableOpenPolicyAgent:                c.EnableOpenPolicyAgent,
-		OpenPolicyAgentConfigTemplate:        c.OpenPolicyAgentConfigTemplate,
-		OpenPolicyAgentEnvoyMetadata:         c.OpenPolicyAgentEnvoyMetadata,
-		OpenPolicyAgentCleanerInterval:       c.OpenPolicyAgentCleanerInterval,
-		OpenPolicyAgentStartupTimeout:        c.OpenPolicyAgentStartupTimeout,
-		OpenPolicyAgentMaxRequestBodySize:    c.OpenPolicyAgentMaxRequestBodySize,
-		OpenPolicyAgentRequestBodyBufferSize: c.OpenPolicyAgentRequestBodyBufferSize,
-		OpenPolicyAgentMaxMemoryBodyParsing:  c.OpenPolicyAgentMaxMemoryBodyParsing,
+		EnableOpenPolicyAgent:                          c.EnableOpenPolicyAgent,
+		OpenPolicyAgentConfigTemplate:                  c.OpenPolicyAgentConfigTemplate,
+		OpenPolicyAgentEnvoyMetadata:                   c.OpenPolicyAgentEnvoyMetadata,
+		OpenPolicyAgentCleanerInterval:                 c.OpenPolicyAgentCleanerInterval,
+		OpenPolicyAgentStartupTimeout:                  c.OpenPolicyAgentStartupTimeout,
+		OpenPolicyAgentMaxRequestBodySize:              c.OpenPolicyAgentMaxRequestBodySize,
+		OpenPolicyAgentRequestBodyBufferSize:           c.OpenPolicyAgentRequestBodyBufferSize,
+		OpenPolicyAgentMaxMemoryBodyParsing:            c.OpenPolicyAgentMaxMemoryBodyParsing,
+		EnableOpenPolicyAgentPreevaluationOptimization: c.EnableOpenPolicyAgentPreevaluationOptimization,
 
 		PassiveHealthCheck: c.PassiveHealthCheck.values,
 	}

filters/openpolicyagent/evaluation.go

@@ -3,17 +3,17 @@ package openpolicyagent
 import (
 	"context"
 	"fmt"
+	"time"
+
 	ext_authz_v3_core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	ext_authz_v3 "github.com/envoyproxy/go-control-plane/envoy/service/auth/v3"
 	"github.com/open-policy-agent/opa-envoy-plugin/envoyauth"
 	"github.com/open-policy-agent/opa-envoy-plugin/opa/decisionlog"
 	"github.com/open-policy-agent/opa/ast"
-	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/server"
 	"github.com/open-policy-agent/opa/topdown"
 	"github.com/opentracing/opentracing-go"
 	pbstruct "google.golang.org/protobuf/types/known/structpb"
-	"time"
 )
 
 func (opa *OpenPolicyAgentInstance) Eval(ctx context.Context, req *ext_authz_v3.CheckRequest) (*envoyauth.EvalResult, error) {
@@ -70,7 +70,7 @@ func (opa *OpenPolicyAgentInstance) Eval(ctx context.Context, req *ext_authz_v3.
 		return nil, err
 	}
 
-	err = envoyauth.Eval(ctx, opa, inputValue, result, rego.DistributedTracingOpts(opa.DistributedTracing()))
+	err = envoyauth.Eval(ctx, opa, inputValue, result)
 	if err != nil {
 		return nil, err
 	}

filters/openpolicyagent/opaauthorizerequest/opaauthorizerequest_test.go

@@ -623,6 +623,7 @@ const (
 
 	certPath = "../../../skptesting/cert.pem"
 	keyPath  = "../../../skptesting/key.pem"
+	dataPath = "../../../skptesting/data.json"
 )
 
 func BenchmarkAuthorizeRequest(b *testing.B) {
@@ -806,6 +807,119 @@ func BenchmarkAuthorizeRequest(b *testing.B) {
 			assert.False(b, ctx.FServed)
 		}
 	})
+
+	b.Run("authorize-request-jwt-validation-with-pre-evaluation", func(b *testing.B) {
+
+		publicKey, err := os.ReadFile(certPath)
+		if err != nil {
+			log.Fatalf("Failed to read public key: %v", err)
+		}
+
+		dataFile, err := os.ReadFile(dataPath)
+		if err != nil {
+			log.Fatalf("Failed to read data.json: %v", err)
+		}
+
+		opaControlPlane := opasdktest.MustNewServer(
+			opasdktest.MockBundle("/bundles/somebundle.tar.gz", map[string]string{
+				"main.rego": fmt.Sprintf(`
+					package envoy.authz
+
+					import future.keywords.if
+
+					default allow = false
+
+					public_key_cert := %q
+
+					bearer_token := t if {
+						v := input.attributes.request.http.headers.authorization
+						startswith(v, "Bearer ")
+						t := substring(v, count("Bearer "), -1)
+					}
+
+					allow if {
+						[valid, _, payload] :=  io.jwt.decode_verify(bearer_token, {
+							"cert": public_key_cert,
+							"aud": "nqz3xhorr5"
+						})
+
+						valid
+						sub := payload.sub
+						has_user_logged_in(sub)
+						is_admin(sub)
+						has_some_read_permission(sub)
+					}
+
+					has_user_logged_in(sub) {
+						data.users[sub].history[_].action == "login"
+					}
+
+					is_admin(sub) {
+						data.users[sub].role == "admin"
+					}
+
+					has_some_read_permission(sub) {
+                     	count(data.users[sub].permissions["read"]) > 0
+					}
+
+				`, publicKey),
+				"data.json": string(dataFile),
+			}),
+		)
+
+		f, err := createOpaFilterWithPreEvaluation(opaControlPlane)
+		assert.NoError(b, err)
+
+		url, err := url.Parse("http://opa-authorized.test/somepath")
+		assert.NoError(b, err)
+
+		claims := jwt.MapClaims{
+			"iss":   "https://some.identity.acme.com",
+			"sub":   "5974934733",
+			"aud":   "nqz3xhorr5",
+			"iat":   time.Now().Add(-time.Minute).UTC().Unix(),
+			"exp":   time.Now().Add(tokenExp).UTC().Unix(),
+			"email": "[email protected]",
+		}
+
+		token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
+
+		privKey, err := os.ReadFile(keyPath)
+		if err != nil {
+			log.Fatalf("Failed to read priv key: %v", err)
+		}
+
+		key, err := jwt.ParseRSAPrivateKeyFromPEM([]byte(privKey))
+		if err != nil {
+			log.Fatalf("Failed to parse RSA PEM: %v", err)
+		}
+
+		// Sign and get the complete encoded token as a string using the secret
+		signedToken, err := token.SignedString(key)
+		if err != nil {
+			log.Fatalf("Failed to sign token: %v", err)
+		}
+
+		ctx := &filtertest.Context{
+			FStateBag: map[string]interface{}{},
+			FResponse: &http.Response{},
+			FRequest: &http.Request{
+				Header: map[string][]string{
+					"Authorization": {fmt.Sprintf("Bearer %s", signedToken)},
+				},
+				URL: url,
+			},
+			FMetrics: &metricstest.MockMetrics{},
+		}
+
+		b.ResetTimer()
+		b.ReportAllocs()
+
+		for i := 0; i < b.N; i++ {
+			f.Request(ctx)
+			assert.False(b, ctx.FServed)
+		}
+	})
 }
 
 func createOpaFilter(opaControlPlane *opasdktest.Server) (filters.Filter, error) {
@@ -815,6 +929,13 @@ func createOpaFilter(opaControlPlane *opasdktest.Server) (filters.Filter, error)
 	return spec.CreateFilter([]interface{}{"somebundle.tar.gz"})
 }
 
+func createOpaFilterWithPreEvaluation(opaControlPlane *opasdktest.Server) (filters.Filter, error) {
+	config := generateConfig(opaControlPlane, "envoy/authz/allow")
+	opaFactory := openpolicyagent.NewOpenPolicyAgentRegistry(openpolicyagent.WithPreevaluationOptimization(true))
+	spec := NewOpaAuthorizeRequestSpec(opaFactory, openpolicyagent.WithConfigTemplate(config))
+	return spec.CreateFilter([]interface{}{"somebundle.tar.gz"})
+}
+
 func createBodyBasedOpaFilter(opaControlPlane *opasdktest.Server) (filters.Filter, error) {
 	config := generateConfig(opaControlPlane, "envoy/authz/allow")
 	opaFactory := openpolicyagent.NewOpenPolicyAgentRegistry()

filters/openpolicyagent/openpolicyagent.go

@@ -17,6 +17,7 @@ import (
 
 	ext_authz_v3_core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	"github.com/google/uuid"
+	"github.com/open-policy-agent/opa-envoy-plugin/envoyauth"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/config"
 	"github.com/open-policy-agent/opa/logging"
@@ -74,6 +75,8 @@ type OpenPolicyAgentRegistry struct {
 	bodyReadBufferSize      int64
 
 	tracer opentracing.Tracer
+
+	preevaluationOptimization bool
 }
 
 type OpenPolicyAgentFilter interface {
@@ -122,6 +125,13 @@ func WithTracer(tracer opentracing.Tracer) func(*OpenPolicyAgentRegistry) error
 	}
 }
 
+func WithPreevaluationOptimization(enable bool) func(*OpenPolicyAgentRegistry) error {
+	return func(cfg *OpenPolicyAgentRegistry) error {
+		cfg.preevaluationOptimization = enable
+		return nil
+	}
+}
+
 func NewOpenPolicyAgentRegistry(opts ...func(*OpenPolicyAgentRegistry) error) *OpenPolicyAgentRegistry {
 	registry := &OpenPolicyAgentRegistry{
 		reuseDuration:       defaultReuseDuration,
@@ -369,6 +379,7 @@ type OpenPolicyAgentInstance struct {
 	bundleName             string
 	preparedQuery          *rego.PreparedEvalQuery
 	preparedQueryDoOnce    *sync.Once
+	preparedQueryErr       error
 	interQueryBuiltinCache iCache.InterQueryCache
 	once                   sync.Once
 	stopped                bool
@@ -740,22 +751,11 @@ func (opa *OpenPolicyAgentInstance) Runtime() *ast.Term { return opa.manager.Inf
 // Logger is an implementation of the envoyauth.EvalContext interface
 func (opa *OpenPolicyAgentInstance) Logger() logging.Logger { return opa.manager.Logger() }
 
-// PreparedQueryDoOnce is an implementation of the envoyauth.EvalContext interface
-func (opa *OpenPolicyAgentInstance) PreparedQueryDoOnce() *sync.Once { return opa.preparedQueryDoOnce }
-
 // InterQueryBuiltinCache is an implementation of the envoyauth.EvalContext interface
 func (opa *OpenPolicyAgentInstance) InterQueryBuiltinCache() iCache.InterQueryCache {
 	return opa.interQueryBuiltinCache
 }
 
-// PreparedQuery is an implementation of the envoyauth.EvalContext interface
-func (opa *OpenPolicyAgentInstance) PreparedQuery() *rego.PreparedEvalQuery { return opa.preparedQuery }
-
-// SetPreparedQuery is an implementation of the envoyauth.EvalContext interface
-func (opa *OpenPolicyAgentInstance) SetPreparedQuery(q *rego.PreparedEvalQuery) {
-	opa.preparedQuery = q
-}
-
 // Config is an implementation of the envoyauth.EvalContext interface
 func (opa *OpenPolicyAgentInstance) Config() *config.Config { return opa.opaConfig }
 
@@ -764,6 +764,27 @@ func (opa *OpenPolicyAgentInstance) DistributedTracing() opatracing.Options {
 	return buildTracingOptions(opa.registry.tracer, opa.bundleName, opa.manager)
 }
 
+// CreatePreparedQueryOnce is an implementation of the envoyauth.EvalContext interface
+func (opa *OpenPolicyAgentInstance) CreatePreparedQueryOnce(opts envoyauth.PrepareQueryOpts) (*rego.PreparedEvalQuery, error) {
+	opa.preparedQueryDoOnce.Do(func() {
+		regoOpts := append(opts.Opts, rego.DistributedTracingOpts(opa.DistributedTracing()))
+
+		prepareOpts := opts.PrepareOpts
+
+		if opa.registry.preevaluationOptimization {
+			prepareOpts = append(prepareOpts, rego.WithPartialEval())
+		}
+
+		pq, err := rego.New(regoOpts...).
+			PrepareForEval(context.Background(), prepareOpts...)
+
+		opa.preparedQuery = &pq
+		opa.preparedQueryErr = err
+	})
+
+	return opa.preparedQuery, opa.preparedQueryErr
+}
+
 // logging.Logger that does not pollute info with debug logs
 type QuietLogger struct {
 	target logging.Logger