From 8f8cd7728f21736dfe483cb82d725412fd9d7f57 Mon Sep 17 00:00:00 2001
From: Pushpalanka Jayawardhana <pushpalanka.jayawardhana@zalando.de>
Date: Fri, 6 Dec 2024 22:56:49 +0100
Subject: [PATCH] Attempted benchmark test with JWT evaluation(no significant
 improvement observed. Need a more complex policy)

Signed-off-by: Pushpalanka Jayawardhana <pushpalanka.jayawardhana@zalando.de>
---
 .../opaauthorizerequest_test.go               | 99 +++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/filters/openpolicyagent/opaauthorizerequest/opaauthorizerequest_test.go b/filters/openpolicyagent/opaauthorizerequest/opaauthorizerequest_test.go
index 89e8f3de50..84affbc850 100644
--- a/filters/openpolicyagent/opaauthorizerequest/opaauthorizerequest_test.go
+++ b/filters/openpolicyagent/opaauthorizerequest/opaauthorizerequest_test.go
@@ -806,6 +806,98 @@ 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)
+		}
+
+		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
+						
+						payload.sub == "5974934733"
+					}				
+				`, publicKey),
+			}),
+		)
+
+		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": "someone@example.org",
+		}
+
+		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 +907,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()