Add convenience functions to remove headers and trailers. (#256)

Signed-off-by: Nandan B N <[email protected]>
proxy-wasm · Jan 15, 2025 · 8604c64 · 8604c64
1 parent f14c4f6
commit 8604c64
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/src/hostcalls.rs b/src/hostcalls.rs
@@ -279,20 +279,29 @@ pub fn get_map_value_bytes(map_type: MapType, key: &str) -> Result<Option<Bytes>
 }
 
 extern "C" {
-    fn proxy_replace_header_map_value(
+    fn proxy_remove_header_map_value(
         map_type: MapType,
         key_data: *const u8,
         key_size: usize,
-        value_data: *const u8,
-        value_size: usize,
     ) -> Status;
 }
 
+pub fn remove_map_value(map_type: MapType, key: &str) -> Result<(), Status> {
+    unsafe {
+        match proxy_remove_header_map_value(map_type, key.as_ptr(), key.len()) {
+            Status::Ok => Ok(()),
+            status => panic!("unexpected status: {}", status as u32),
+        }
+    }
+}
+
 extern "C" {
-    fn proxy_remove_header_map_value(
+    fn proxy_replace_header_map_value(
         map_type: MapType,
         key_data: *const u8,
         key_size: usize,
+        value_data: *const u8,
+        value_size: usize,
     ) -> Status;
 }
 

diff --git a/src/traits.rs b/src/traits.rs
@@ -351,6 +351,10 @@ pub trait HttpContext: Context {
         hostcalls::add_map_value_bytes(MapType::HttpRequestHeaders, name, value).unwrap()
     }
 
+    fn remove_http_request_header(&self, name: &str) {
+        hostcalls::remove_map_value(MapType::HttpRequestHeaders, name).unwrap()
+    }
+
     fn on_http_request_body(&mut self, _body_size: usize, _end_of_stream: bool) -> Action {
         Action::Continue
     }
@@ -407,6 +411,10 @@ pub trait HttpContext: Context {
         hostcalls::add_map_value_bytes(MapType::HttpRequestTrailers, name, value).unwrap()
     }
 
+    fn remove_http_request_trailer(&self, name: &str) {
+        hostcalls::remove_map_value(MapType::HttpRequestTrailers, name).unwrap()
+    }
+
     fn resume_http_request(&self) {
         hostcalls::resume_http_request().unwrap()
     }
@@ -459,6 +467,10 @@ pub trait HttpContext: Context {
         hostcalls::add_map_value_bytes(MapType::HttpResponseHeaders, name, value).unwrap()
     }
 
+    fn remove_http_response_header(&self, name: &str) {
+        hostcalls::remove_map_value(MapType::HttpResponseHeaders, name).unwrap()
+    }
+
     fn on_http_response_body(&mut self, _body_size: usize, _end_of_stream: bool) -> Action {
         Action::Continue
     }
@@ -515,6 +527,10 @@ pub trait HttpContext: Context {
         hostcalls::add_map_value_bytes(MapType::HttpResponseTrailers, name, value).unwrap()
     }
 
+    fn remove_http_response_trailer(&self, name: &str) {
+        hostcalls::remove_map_value(MapType::HttpResponseTrailers, name).unwrap()
+    }
+
     fn resume_http_response(&self) {
         hostcalls::resume_http_response().unwrap()
     }