Re-introduce GetResponseHTTPHeaders method (#667)

Signed-off-by: Pushpalanka Jayawardhana <[email protected]>

Author: Pushpalanka

envoyauth/response.go

@@ -4,9 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"net/http"
-	"slices"
-
 	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"
@@ -16,6 +13,8 @@ import (
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/topdown/builtins"
 	"google.golang.org/protobuf/types/known/structpb"
+	"net/http"
+	"slices"
 )
 
 // EvalResult - Captures the result from evaluating a query against an input
@@ -185,20 +184,9 @@ func (result *EvalResult) GetRequestHTTPHeadersToRemove() ([]string, error) {
 	return result.getStringSliceFromDecision("request_headers_to_remove")
 }
 
-func (result *EvalResult) getHeadersFromDecision(fieldName string) ([]*ext_core_v3.HeaderValueOption, error) {
-	switch decision := result.Decision.(type) {
-	case bool:
-		return nil, nil
-	case map[string]interface{}:
-		val, ok := decision[fieldName]
-		if !ok {
-			return nil, nil
-		}
-
-		return transformHeadersToEnvoy(val)
-	default:
-		return nil, result.invalidDecisionErr()
-	}
+// GetResponseHTTPHeaders - returns the http headers to return if they are part of the decision as http header
+func (result *EvalResult) GetResponseHTTPHeaders() (http.Header, error) {
+	return result.getHTTPHeadersFromDecision("headers")
 }
 
 // GetResponseEnvoyHeaderValueOptions - returns the http headers to return if they are part of the decision as envoy header value options
@@ -343,54 +331,112 @@ func makeHeaderValueOption(k, v string) *ext_core_v3.HeaderValueOption {
 	}
 }
 
-func makeEnvoyHeaderValueOptionsFromHeadersMap(hvo []*ext_core_v3.HeaderValueOption, headers map[string]any) ([]*ext_core_v3.HeaderValueOption, error) {
-	hvo = slices.Grow(hvo, len(headers))
+func (result *EvalResult) getHeadersFromDecision(fieldName string) ([]*ext_core_v3.HeaderValueOption, error) {
+	return getHeadersWithTransformation(result.Decision, fieldName, preallocateForEnvoyHeaders, collectEnvoyHeaders)
+}
+
+func (result *EvalResult) getHTTPHeadersFromDecision(fieldName string) (http.Header, error) {
+	return getHeadersWithTransformation(result.Decision, fieldName, preallocateForHTTPHeaders, collectHTTPHeaders)
+}
+
+func getHeadersWithTransformation[T any](
+	decision any,
+	fieldName string,
+	preallocate func(*T, int),
+	collector func(string, string, *T),
+) (T, error) {
+	var result T
+
+	switch decision := decision.(type) {
+	case bool:
+		return result, nil
+	case map[string]interface{}:
+		headersList, err := extractHeadersFromDecision(decision, fieldName)
+		if err != nil {
+			return result, err
+		}
+
+		for _, headers := range headersList {
+			if err := collectHeaderValues(headers, preallocate, collector, &result); err != nil {
+				return result, err
+			}
+		}
+		return result, nil
+	default:
+		return result, fmt.Errorf("illegal value for policy evaluation result: %T", decision)
+	}
+}
+
+func extractHeadersFromDecision(decision map[string]interface{}, fieldName string) ([]map[string]interface{}, error) {
+	val, ok := decision[fieldName]
+	if !ok {
+		return nil, nil
+	}
+
+	switch v := val.(type) {
+	case map[string]interface{}:
+		return []map[string]interface{}{v}, nil
+	case []interface{}:
+		headersList := make([]map[string]interface{}, 0, len(v))
+		for _, item := range v {
+			headers, ok := item.(map[string]interface{})
+			if !ok {
+				return nil, fmt.Errorf("type assertion error, expected headers to be of type 'object' but got '%T'", item)
+			}
+			headersList = append(headersList, headers)
+		}
+		return headersList, nil
+	default:
+		return nil, fmt.Errorf("type assertion error, expected headers to be of type 'object' but got '%T'", v)
+	}
+}
+
+func collectHeaderValues[T any](
+	headers map[string]interface{},
+	preallocate func(*T, int),
+	collector func(string, string, *T),
+	result *T,
+) error {
 	for key, value := range headers {
 		switch val := value.(type) {
 		case string:
-			hvo = append(hvo, makeHeaderValueOption(key, val))
+			preallocate(result, 1)
+			collector(key, val, result)
 		case []string:
-			hvo = slices.Grow(hvo, len(val))
+			preallocate(result, len(val))
 			for _, v := range val {
-				hvo = append(hvo, makeHeaderValueOption(key, v))
+				collector(key, v, result)
 			}
 		case []interface{}:
-			hvo = slices.Grow(hvo, len(val))
+			preallocate(result, len(val))
 			for _, v := range val {
 				s, ok := v.(string)
 				if !ok {
-					return nil, fmt.Errorf("invalid value type %T for header '%s'", v, key)
+					return fmt.Errorf("invalid value type %T for header '%s'", v, key)
 				}
-				hvo = append(hvo, makeHeaderValueOption(key, s))
+				collector(key, s, result)
 			}
 		default:
-			return nil, fmt.Errorf("type assertion error for header '%s'", key)
+			return fmt.Errorf("type assertion error for header '%s'", key)
 		}
 	}
-	return hvo, nil
+	return nil
 }
 
-func transformHeadersToEnvoy(input any) ([]*ext_core_v3.HeaderValueOption, error) {
-	switch input := input.(type) {
-	case []any:
-		var (
-			hvo []*ext_core_v3.HeaderValueOption
-			err error
-		)
-		for _, val := range input {
-			headers, ok := val.(map[string]any)
-			if !ok {
-				return nil, fmt.Errorf("type assertion error, expected headers to be of type 'object' but got '%T'", val)
-			}
+func collectEnvoyHeaders(key, value string, result *[]*ext_core_v3.HeaderValueOption) {
+	*result = append(*result, makeHeaderValueOption(key, value))
+}
 
-			hvo, err = makeEnvoyHeaderValueOptionsFromHeadersMap(hvo, headers)
-			if err != nil {
-				return nil, err
-			}
-		}
-		return hvo, nil
-	case map[string]any:
-		return makeEnvoyHeaderValueOptionsFromHeadersMap(nil, input)
+func collectHTTPHeaders(key, value string, result *http.Header) {
+	result.Add(key, value)
+}
+
+func preallocateForEnvoyHeaders(result *[]*ext_core_v3.HeaderValueOption, additional int) {
+	*result = slices.Grow(*result, additional)
+}
+
+func preallocateForHTTPHeaders(result *http.Header, _ int) {
+	if *result == nil {
+		*result = make(http.Header)
 	}
-	return nil, fmt.Errorf("type assertion error, expected headers to be of type 'object' but got '%T'", input)
 }

envoyauth/response_test.go

@@ -373,7 +373,7 @@ func TestGetResponseHTTPHeadersToAdd(t *testing.T) {
 	}
 }
 
-func TestGetResponseHeaders(t *testing.T) {
+func TestGetResponseHeaderValueOptions(t *testing.T) {
 	input := make(map[string]interface{})
 	er := EvalResult{
 		Decision: input,
@@ -455,6 +455,127 @@ func TestGetResponseHeaders(t *testing.T) {
 	if len(result) != 2 {
 		t.Fatalf("Expected two header but got %v", len(result))
 	}
+
+	testAddHeaders["foo"] = []interface{}{"bar", "baz"}
+	input["headers"] = testAddHeaders
+
+	result, err = er.GetResponseEnvoyHeaderValueOptions()
+
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+
+	if len(result) != 2 {
+		t.Fatalf("Expected two header but got %v", len(result))
+	}
+
+	if seen["bar"] != 1 || seen["baz"] != 1 {
+		t.Errorf("expected 'bar' and 'baz', got %v", seen)
+	}
+}
+
+func TestGetResponseHeaders(t *testing.T) {
+	input := make(map[string]interface{})
+	er := EvalResult{
+		Decision: input,
+	}
+
+	result, err := er.GetResponseHTTPHeaders()
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+
+	if len(result) != 0 {
+		t.Fatal("Expected no headers")
+	}
+
+	badHeader := "test"
+	input["headers"] = badHeader
+
+	_, err = er.GetResponseHTTPHeaders()
+	if err == nil {
+		t.Fatal("Expected error but got nil")
+	}
+
+	testHeaders := make(map[string]interface{})
+	testHeaders["foo"] = "bar"
+	input["headers"] = testHeaders
+
+	result, err = er.GetResponseHTTPHeaders()
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+
+	if len(result) != 1 {
+		t.Fatalf("Expected one header but got %v", len(result))
+	}
+
+	testHeaders["baz"] = 1
+
+	_, err = er.GetResponseHTTPHeaders()
+	if err == nil {
+		t.Fatal("Expected error but got nil")
+	}
+
+	input["headers"] = []interface{}{
+		map[string]interface{}{
+			"foo": "bar",
+		},
+		map[string]interface{}{
+			"foo": "baz",
+		},
+	}
+
+	result, err = er.GetResponseHTTPHeaders()
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+
+	if len(result.Values("foo")) != 2 {
+		t.Fatalf("Expected two header values but got %v", result.Values("foo"))
+	}
+
+	seen := map[string]int{}
+	for _, values := range result {
+		for _, value := range values {
+			seen[value]++
+		}
+	}
+
+	if seen["bar"] != 1 || seen["baz"] != 1 {
+		t.Errorf("expected 'bar' and 'baz', got %v", seen)
+	}
+
+	testAddHeaders := make(map[string]interface{})
+	testAddHeaders["foo"] = []string{"bar", "baz"}
+	input["headers"] = testAddHeaders
+
+	result, err = er.GetResponseHTTPHeaders()
+
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+
+	if len(result.Values("foo")) != 2 {
+		t.Fatalf("Expected two header but got %v", len(result.Values("foo")))
+	}
+
+	testAddHeaders["foo"] = []interface{}{"bar", "baz"}
+	input["headers"] = testAddHeaders
+
+	result, err = er.GetResponseHTTPHeaders()
+
+	if err != nil {
+		t.Fatalf("Expected no error but got %v", err)
+	}
+
+	if len(result.Values("foo")) != 2 {
+		t.Fatalf("Expected two header but got %v", len(result.Values("foo")))
+	}
+
+	if seen["bar"] != 1 || seen["baz"] != 1 {
+		t.Errorf("expected 'bar' and 'baz', got %v", seen)
+	}
 }
 
 func TestGetResponseBody(t *testing.T) {