update proxy_on_request_headers and proxy_on_response_headers

yskopets · Aug 27, 2020 · f97d84f · f97d84f
1 parent 810381d
commit f97d84f
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 14 deletions.
diff --git a/examples/http_auth_random.rs b/examples/http_auth_random.rs
@@ -28,7 +28,7 @@ pub fn _start() {
 struct HttpAuthRandom;
 
 impl HttpContext for HttpAuthRandom {
-    fn on_http_request_headers(&mut self, _: usize) -> Action {
+    fn on_http_request_headers(&mut self, _: usize, _: bool) -> Action {
         self.dispatch_http_call(
             "httpbin",
             vec![
@@ -44,7 +44,7 @@ impl HttpContext for HttpAuthRandom {
         Action::Pause
     }
 
-    fn on_http_response_headers(&mut self, _: usize) -> Action {
+    fn on_http_response_headers(&mut self, _: usize, _: bool) -> Action {
         self.set_http_response_header("Powered-By", Some("proxy-wasm"));
         Action::Continue
     }

diff --git a/examples/http_headers.rs b/examples/http_headers.rs
@@ -33,7 +33,7 @@ struct HttpHeaders {
 impl Context for HttpHeaders {}
 
 impl HttpContext for HttpHeaders {
-    fn on_http_request_headers(&mut self, _: usize) -> Action {
+    fn on_http_request_headers(&mut self, _: usize, _: bool) -> Action {
         for (name, value) in &self.get_http_request_headers() {
             trace!("#{} -> {}: {}", self.context_id, name, value);
         }
@@ -51,7 +51,7 @@ impl HttpContext for HttpHeaders {
         }
     }
 
-    fn on_http_response_headers(&mut self, _: usize) -> Action {
+    fn on_http_response_headers(&mut self, _: usize, _: bool) -> Action {
         for (name, value) in &self.get_http_response_headers() {
             trace!("#{} <- {}: {}", self.context_id, name, value);
         }

diff --git a/src/dispatcher.rs b/src/dispatcher.rs
@@ -311,10 +311,15 @@ impl Dispatcher {
         }
     }
 
-    fn on_http_request_headers(&self, context_id: u32, num_headers: usize) -> Action {
+    fn on_http_request_headers(
+        &self,
+        context_id: u32,
+        num_headers: usize,
+        end_of_stream: bool,
+    ) -> Action {
         if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
-            http_stream.on_http_request_headers(num_headers)
+            http_stream.on_http_request_headers(num_headers, end_of_stream)
         } else {
             panic!("invalid context_id")
         }
@@ -343,10 +348,15 @@ impl Dispatcher {
         }
     }
 
-    fn on_http_response_headers(&self, context_id: u32, num_headers: usize) -> Action {
+    fn on_http_response_headers(
+        &self,
+        context_id: u32,
+        num_headers: usize,
+        end_of_stream: bool,
+    ) -> Action {
         if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
-            http_stream.on_http_response_headers(num_headers)
+            http_stream.on_http_response_headers(num_headers, end_of_stream)
         } else {
             panic!("invalid context_id")
         }
@@ -477,8 +487,14 @@ pub extern "C" fn proxy_on_upstream_connection_close(context_id: u32, peer_type:
 }
 
 #[no_mangle]
-pub extern "C" fn proxy_on_request_headers(context_id: u32, num_headers: usize) -> Action {
-    DISPATCHER.with(|dispatcher| dispatcher.on_http_request_headers(context_id, num_headers))
+pub extern "C" fn proxy_on_request_headers(
+    context_id: u32,
+    num_headers: usize,
+    end_of_stream: bool,
+) -> Action {
+    DISPATCHER.with(|dispatcher| {
+        dispatcher.on_http_request_headers(context_id, num_headers, end_of_stream)
+    })
 }
 
 #[no_mangle]
@@ -497,8 +513,14 @@ pub extern "C" fn proxy_on_request_trailers(context_id: u32, num_trailers: usize
 }
 
 #[no_mangle]
-pub extern "C" fn proxy_on_response_headers(context_id: u32, num_headers: usize) -> Action {
-    DISPATCHER.with(|dispatcher| dispatcher.on_http_response_headers(context_id, num_headers))
+pub extern "C" fn proxy_on_response_headers(
+    context_id: u32,
+    num_headers: usize,
+    end_of_stream: bool,
+) -> Action {
+    DISPATCHER.with(|dispatcher| {
+        dispatcher.on_http_response_headers(context_id, num_headers, end_of_stream)
+    })
 }
 
 #[no_mangle]

diff --git a/src/traits.rs b/src/traits.rs
@@ -157,7 +157,7 @@ pub trait StreamContext: Context {
 }
 
 pub trait HttpContext: Context {
-    fn on_http_request_headers(&mut self, _num_headers: usize) -> Action {
+    fn on_http_request_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
         Action::Continue
     }
 
@@ -217,7 +217,7 @@ pub trait HttpContext: Context {
         hostcalls::continue_stream(StreamType::Request).unwrap()
     }
 
-    fn on_http_response_headers(&mut self, _num_headers: usize) -> Action {
+    fn on_http_response_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
         Action::Continue
     }