From 8604c648a929754d228f9b082fe8606772afea08 Mon Sep 17 00:00:00 2001
From: Nandan B N <its.nandanbn@gmail.com>
Date: Wed, 15 Jan 2025 19:37:37 +0530
Subject: [PATCH] Add convenience functions to remove headers and trailers.
 (#256)

Signed-off-by: Nandan B N <nandan_bn@intuit.com>
---
 src/hostcalls.rs | 17 +++++++++++++----
 src/traits.rs    | 16 ++++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/hostcalls.rs b/src/hostcalls.rs
index f8252c56..8dbc335c 100644
--- 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
index bd54bcbe..e0ae9477 100644
--- 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()
     }