From 8ded8c63f43e7d0b1a1041d1d0f7b00d44067d89 Mon Sep 17 00:00:00 2001
From: Adam Cattermole <acatterm@redhat.com>
Date: Tue, 12 Nov 2024 16:52:58 +0000
Subject: [PATCH] Correctly add auth headers that should be sent to the
 upstream

Signed-off-by: Adam Cattermole <acatterm@redhat.com>
---
 src/service/auth.rs       | 42 +++++++++++++++++++--------------------
 src/service/rate_limit.rs |  1 +
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/src/service/auth.rs b/src/service/auth.rs
index c759986..a1e4e96 100644
--- a/src/service/auth.rs
+++ b/src/service/auth.rs
@@ -133,28 +133,26 @@ impl AuthService {
             match check_response.http_response {
                 Some(CheckResponse_oneof_http_response::ok_response(ok_response)) => {
                     debug!("process_auth_grpc_response: received OkHttpResponse");
-                    ok_response
-                        .get_headers_to_remove()
-                        .iter()
-                        .for_each(|header| {
-                            hostcalls::set_map_value(
-                                MapType::HttpResponseHeaders,
-                                header.as_str(),
-                                None,
-                            )
-                            .unwrap()
-                        });
-                    ok_response
-                        .get_response_headers_to_add()
-                        .iter()
-                        .for_each(|header| {
-                            hostcalls::add_map_value(
-                                MapType::HttpResponseHeaders,
-                                header.get_header().get_key(),
-                                header.get_header().get_value(),
-                            )
-                            .unwrap()
-                        });
+                    if !ok_response.get_response_headers_to_add().is_empty() {
+                        panic!("process_auth_grpc_response: response contained response_headers_to_add which is unsupported!")
+                    }
+                    if !ok_response.get_headers_to_remove().is_empty() {
+                        panic!("process_auth_grpc_response: response contained headers_to_remove which is unsupported!")
+                    }
+                    if !ok_response.get_query_parameters_to_set().is_empty() {
+                        panic!("process_auth_grpc_response: response contained query_parameters_to_set which is unsupported!")
+                    }
+                    if !ok_response.get_query_parameters_to_remove().is_empty() {
+                        panic!("process_auth_grpc_response: response contained query_parameters_to_remove which is unsupported!")
+                    }
+                    ok_response.get_headers().iter().for_each(|header| {
+                        hostcalls::add_map_value(
+                            MapType::HttpRequestHeaders,
+                            header.get_header().get_key(),
+                            header.get_header().get_value(),
+                        )
+                        .unwrap()
+                    });
                     Ok(())
                 }
                 Some(CheckResponse_oneof_http_response::denied_response(denied_response)) => {
diff --git a/src/service/rate_limit.rs b/src/service/rate_limit.rs
index 01c0e33..2f97f31 100644
--- a/src/service/rate_limit.rs
+++ b/src/service/rate_limit.rs
@@ -65,6 +65,7 @@ impl RateLimitService {
                 response_headers_to_add: additional_headers,
                 ..
             }) => {
+                // TODO: This should not be sent to the upstream!
                 additional_headers.iter().for_each(|header| {
                     hostcalls::add_map_value(
                         MapType::HttpResponseHeaders,