Add response headers in correct phase

Signed-off-by: Adam Cattermole <acatterm@redhat.com>
Kuadrant · Nov 13, 2024 · 36a5c42 · 36a5c42
1 parent 5e73c10
commit 36a5c42
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 18 deletions.
diff --git a/src/filter/http_context.rs b/src/filter/http_context.rs
@@ -142,7 +142,13 @@ impl Context for Filter {
 
         match op_res {
             Ok(operation) => {
-                if GrpcService::process_grpc_response(operation, resp_size).is_ok() {
+                if GrpcService::process_grpc_response(
+                    operation,
+                    resp_size,
+                    &mut self.response_headers_to_add,
+                )
+                .is_ok()
+                {
                     // call the next op
                     match self.operation_dispatcher.borrow_mut().next() {
                         Ok(some_op) => {

diff --git a/src/service.rs b/src/service.rs
@@ -54,6 +54,7 @@ impl GrpcService {
     pub fn process_grpc_response(
         operation: Rc<Operation>,
         resp_size: usize,
+        response_headers_to_add: &mut Vec<(String, String)>,
     ) -> Result<(), StatusCode> {
         let failure_mode = operation.get_failure_mode();
         if let Some(res_body_bytes) =
@@ -62,9 +63,11 @@ impl GrpcService {
             match GrpcMessageResponse::new(operation.get_service_type(), &res_body_bytes) {
                 Ok(res) => match operation.get_service_type() {
                     ServiceType::Auth => AuthService::process_auth_grpc_response(res, failure_mode),
-                    ServiceType::RateLimit => {
-                        RateLimitService::process_ratelimit_grpc_response(res, failure_mode)
-                    }
+                    ServiceType::RateLimit => RateLimitService::process_ratelimit_grpc_response(
+                        res,
+                        failure_mode,
+                        response_headers_to_add,
+                    ),
                 },
                 Err(e) => {
                     warn!(

diff --git a/src/service/rate_limit.rs b/src/service/rate_limit.rs
@@ -7,7 +7,7 @@ use crate::service::GrpcService;
 use log::warn;
 use protobuf::{Message, RepeatedField};
 use proxy_wasm::hostcalls;
-use proxy_wasm::types::{Bytes, MapType};
+use proxy_wasm::types::Bytes;
 
 pub const RATELIMIT_SERVICE_NAME: &str = "envoy.service.ratelimit.v3.RateLimitService";
 pub const RATELIMIT_METHOD_NAME: &str = "ShouldRateLimit";
@@ -38,6 +38,7 @@ impl RateLimitService {
     pub fn process_ratelimit_grpc_response(
         rl_resp: GrpcMessageResponse,
         failure_mode: FailureMode,
+        response_headers_to_add: &mut Vec<(String, String)>,
     ) -> Result<(), StatusCode> {
         match rl_resp {
             GrpcMessageResponse::RateLimit(RateLimitResponse {
@@ -65,14 +66,9 @@ 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,
-                        header.get_key(),
-                        header.get_value(),
-                    )
-                    .unwrap()
+                    response_headers_to_add
+                        .push((header.get_key().to_owned(), header.get_value().to_owned()))
                 });
                 Ok(())
             }

diff --git a/tests/rate_limited.rs b/tests/rate_limited.rs
@@ -353,6 +353,12 @@ fn it_passes_additional_headers() {
         )
         .expect_get_buffer_bytes(Some(BufferType::GrpcReceiveBuffer))
         .returning(Some(&grpc_response))
+        .execute_and_expect(ReturnType::None)
+        .unwrap();
+
+    module
+        .call_proxy_on_response_headers(http_context, 0, false)
+        .expect_log(Some(LogLevel::Debug), Some("#2 on_http_response_headers"))
         .expect_add_header_map_value(
             Some(MapType::HttpResponseHeaders),
             Some("test"),
@@ -363,12 +369,6 @@ fn it_passes_additional_headers() {
             Some("other"),
             Some("header value"),
         )
-        .execute_and_expect(ReturnType::None)
-        .unwrap();
-
-    module
-        .call_proxy_on_response_headers(http_context, 0, false)
-        .expect_log(Some(LogLevel::Debug), Some("#2 on_http_response_headers"))
         .execute_and_expect(ReturnType::Action(Action::Continue))
         .unwrap();
 }