From 78ec0485452102a57250d3f0d28ba2ded1403334 Mon Sep 17 00:00:00 2001
From: tjons <tyler.schade@solo.io>
Date: Mon, 6 Nov 2023 02:37:22 +0000
Subject: [PATCH] implement dynamic_metadata api

Signed-off-by: tjons <tyler.schade@solo.io>
---
 envoyauth/response.go     | 29 +++++++++++++++++++++++++++++
 internal/internal.go      |  7 ++++++-
 internal/internal_test.go | 35 +++++++++++++++++++++++++++++++++--
 3 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/envoyauth/response.go b/envoyauth/response.go
index 10650ad7c..a9872dd69 100644
--- a/envoyauth/response.go
+++ b/envoyauth/response.go
@@ -8,10 +8,12 @@ import (
 
 	ext_core_v3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	ext_type_v3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
+	_structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/open-policy-agent/opa-envoy-plugin/internal/util"
 	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/topdown/builtins"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 // EvalResult - Captures the result from evaluating a query against an input
@@ -280,6 +282,33 @@ func (result *EvalResult) GetResponseHTTPStatus() (int, error) {
 	return http.StatusForbidden, result.invalidDecisionErr()
 }
 
+// GetDynamicMetadata returns the dynamic metadata to return if part of the decision
+func (result *EvalResult) GetDynamicMetadata() (*_structpb.Struct, error) {
+	var (
+		val interface{}
+		ok  bool
+	)
+	switch decision := result.Decision.(type) {
+	case bool:
+		if decision {
+			return nil, fmt.Errorf("dynamic metadata undefined for simple 'allow'")
+		}
+	case map[string]interface{}:
+		if val, ok = decision["dynamic_metadata"]; !ok {
+			return nil, nil
+		}
+
+		metadata, ok := val.(map[string]interface{})
+		if !ok {
+			return nil, fmt.Errorf("type assertion error")
+		}
+
+		return structpb.NewStruct(metadata)
+	}
+
+	return nil, nil
+}
+
 // GetResponseEnvoyHTTPStatus returns the http status to return if they are part of the decision
 func (result *EvalResult) GetResponseEnvoyHTTPStatus() (*ext_type_v3.HttpStatus, error) {
 	status := &ext_type_v3.HttpStatus{
diff --git a/internal/internal.go b/internal/internal.go
index 4aff1f123..aa131b15d 100644
--- a/internal/internal.go
+++ b/internal/internal.go
@@ -464,8 +464,13 @@ func (p *envoyExtAuthzGrpcServer) check(ctx context.Context, req interface{}) (*
 			return nil, stop, &internalErr
 		}
 
-		if status == int32(code.Code_OK) {
+		dynamicMetadata, err := result.GetDynamicMetadata()
+		if err != nil {
+			return nil, stop, errors.Wrap(err, "failed to get dynamic metadata")
+		}
+		resp.DynamicMetadata = dynamicMetadata
 
+		if status == int32(code.Code_OK) {
 			var headersToRemove []string
 			headersToRemove, err = result.GetRequestHTTPHeadersToRemove()
 			if err != nil {
diff --git a/internal/internal_test.go b/internal/internal_test.go
index e2d85ff84..3b2dc76b7 100644
--- a/internal/internal_test.go
+++ b/internal/internal_test.go
@@ -19,8 +19,10 @@ import (
 	ext_core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	ext_authz_v2 "github.com/envoyproxy/go-control-plane/envoy/service/auth/v2"
 	ext_authz "github.com/envoyproxy/go-control-plane/envoy/service/auth/v3"
+	_structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/prometheus/client_golang/prometheus"
 	"google.golang.org/genproto/googleapis/rpc/code"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/open-policy-agent/opa-envoy-plugin/envoyauth"
 	"github.com/open-policy-agent/opa/ast"
@@ -1463,6 +1465,11 @@ func TestCheckAllowObjectDecision(t *testing.T) {
 	expectedHeaders[http.CanonicalHeaderKey("y")] = "world"
 
 	assertHeaders(t, headers, expectedHeaders)
+
+	dynamicMetadata := output.GetDynamicMetadata()
+	if dynamicMetadata == nil {
+		t.Fatal("Expected DynamicMetadata struct but got nil")
+	}
 }
 
 func TestCheckDenyObjectDecision(t *testing.T) {
@@ -1563,6 +1570,21 @@ func TestCheckAllowWithDryRunObjectDecision(t *testing.T) {
 	expectedHeaders[http.CanonicalHeaderKey("y")] = "world"
 
 	assertHeaders(t, headers, expectedHeaders)
+
+	assertDynamicMetadata(t, &_structpb.Struct{
+		Fields: map[string]*_structpb.Value{
+			"test": {
+				Kind: &_structpb.Value_StringValue{
+					StringValue: "foo",
+				},
+			},
+			"bar": {
+				Kind: &_structpb.Value_StringValue{
+					StringValue: "baz",
+				},
+			},
+		},
+	}, output.GetDynamicMetadata())
 }
 
 func TestPluginStatusLifeCycle(t *testing.T) {
@@ -1741,14 +1763,16 @@ func testAuthzServerWithObjectDecision(customConfig *Config, customPluginFuncs .
 		  "allowed": false,
 		  "headers": {"foo": "bar", "baz": "taz"},
 		  "body": "Unauthorized Request",
-		  "http_status": 301
+		  "http_status": 301,
+		  "dynamic_metadata": {"test": "foo", "bar": "baz"}
 		}
 
 		allow = response {
 			input.parsed_path = ["my", "test", "path"]
 		    response := {
 				"allowed": true,
-				"headers": {"x": "hello", "y": "world"}
+				"headers": {"x": "hello", "y": "world"},
+				"dynamic_metadata": {"test": "foo", "bar": "baz"}
 		    }
 		}`
 
@@ -2016,6 +2040,13 @@ func assertErrorCounterMetric(t *testing.T, server *envoyExtAuthzGrpcServer, lab
 	}
 }
 
+func assertDynamicMetadata(t *testing.T, expectedMetadata, actualMetadata *_structpb.Struct) {
+	t.Helper()
+	if !proto.Equal(expectedMetadata, actualMetadata) {
+		t.Fatalf("Expected metadata %v but got %v", expectedMetadata, actualMetadata)
+	}
+}
+
 type testPlugin struct {
 	events []logs.EventV1
 }