Add support for gRPC streams.

src/dispatcher.rs

@@ -42,6 +42,10 @@ pub(crate) fn register_grpc_callout(token_id: u32) {
     DISPATCHER.with(|dispatcher| dispatcher.register_grpc_callout(token_id));
 }
 
+pub(crate) fn register_grpc_stream(token_id: u32) {
+    DISPATCHER.with(|dispatcher| dispatcher.register_grpc_stream(token_id));
+}
+
 struct NoopRoot;
 
 impl Context for NoopRoot {}
@@ -57,6 +61,7 @@ struct Dispatcher {
     active_id: Cell<u32>,
     callouts: RefCell<HashMap<u32, u32>>,
     grpc_callouts: RefCell<HashMap<u32, u32>>,
+    grpc_stream_tokens: RefCell<HashMap<u32, u32>>,
 }
 
 impl Dispatcher {
@@ -71,6 +76,7 @@ impl Dispatcher {
             active_id: Cell::new(0),
             callouts: RefCell::new(HashMap::new()),
             grpc_callouts: RefCell::new(HashMap::new()),
+            grpc_stream_tokens: RefCell::new(HashMap::new()),
         }
     }
 
@@ -97,6 +103,17 @@ impl Dispatcher {
         }
     }
 
+    fn register_grpc_stream_tokens(&self, token_id: u32) {
+        if self
+            .grpc_stream_tokens
+            .borrow_mut()
+            .insert(token_id, self.active_id.get())
+            .is_some()
+        {
+            panic!("duplicate token_id")
+        }
+    }
+
     fn register_grpc_callout(&self, token_id: u32) {
         if self
             .grpc_callouts
@@ -400,46 +417,117 @@ impl Dispatcher {
     }
 
     fn on_grpc_receive(&self, token_id: u32, response_size: usize) {
+        if let Some(context_id) = self.grpc_callouts.borrow_mut().remove(&token_id) {
+            if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                http_stream.on_grpc_call_response(token_id, 0, response_size);
+            } else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                stream.on_grpc_call_response(token_id, 0, response_size);
+            } else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                root.on_grpc_call_response(token_id, 0, response_size);
+            }
+        } else if let Some(context_id) = self.grpc_callouts.borrow_mut().get(&token_id) {
+            let context_id = context_id.clone();
+            if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                http_stream.on_grpc_stream_receive_body(token_id, response_size);
+            } else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                stream.on_grpc_stream_receive_body(token_id, response_size);
+            } else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                root.on_grpc_stream_receive_body(token_id, response_size);
+            }
+        } else {
+            panic!("invalid token_id")
+        }
+    }
+
+    fn on_grpc_close(&self, token_id: u32, status_code: u32) {
+        if let Some(context_id) = self.grpc_callouts.borrow_mut().remove(&token_id) {
+            if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                http_stream.on_grpc_call_response(token_id, status_code, 0);
+            } else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                stream.on_grpc_call_response(token_id, status_code, 0);
+            } else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                root.on_grpc_call_response(token_id, status_code, 0);
+            }
+        } else if let Some(context_id) = self.grpc_stream_tokens.borrow_mut().remove(&token_id) {
+            if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                http_stream.on_grpc_stream_close(token_id, status_code)
+            } else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                stream.on_grpc_stream_close(token_id, status_code)
+            } else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
+                self.active_id.set(context_id);
+                hostcalls::set_effective_context(context_id).unwrap();
+                root.on_grpc_stream_close(token_id, status_code)
+            }
+        } else {
+            panic!("invalid token_id")
+        }
+    }
+
+    fn on_grpc_receive_initial_metadata(&self, token_id: u32, headers: u32) {
         let context_id = self
-            .grpc_callouts
+            .grpc_stream_tokens
             .borrow_mut()
-            .remove(&token_id)
-            .expect("invalid token_id");
+            .get(&token_id)
+            .expect("invalid token_id")
+            .clone();
 
         if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
             hostcalls::set_effective_context(context_id).unwrap();
-            http_stream.on_grpc_call_response(token_id, 0, response_size);
+            http_stream.on_grpc_stream_receive_initial_metadata(token_id, headers);
         } else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
             hostcalls::set_effective_context(context_id).unwrap();
-            stream.on_grpc_call_response(token_id, 0, response_size);
+            stream.on_grpc_stream_receive_initial_metadata(token_id, headers);
         } else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
             hostcalls::set_effective_context(context_id).unwrap();
-            root.on_grpc_call_response(token_id, 0, response_size);
+            root.on_grpc_stream_receive_initial_metadata(token_id, headers);
         }
     }
 
-    fn on_grpc_close(&self, token_id: u32, status_code: u32) {
+    fn on_grpc_receive_trailing_metadata(&self, token_id: u32, trailers: u32) {
         let context_id = self
-            .grpc_callouts
+            .grpc_stream_tokens
             .borrow_mut()
-            .remove(&token_id)
-            .expect("invalid token_id");
+            .get(&token_id)
+            .expect("invalid token_id")
+            .clone();
 
         if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
             hostcalls::set_effective_context(context_id).unwrap();
-            http_stream.on_grpc_call_response(token_id, status_code, 0);
+            http_stream.on_grpc_stream_receive_trailing_metadata(token_id, trailers);
         } else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
             hostcalls::set_effective_context(context_id).unwrap();
-            stream.on_grpc_call_response(token_id, status_code, 0);
+            stream.on_grpc_stream_receive_trailing_metadata(token_id, trailers);
         } else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
             self.active_id.set(context_id);
             hostcalls::set_effective_context(context_id).unwrap();
-            root.on_grpc_call_response(token_id, status_code, 0);
+            root.on_grpc_stream_receive_trailing_metadata(token_id, trailers);
         }
     }
 }

src/hostcalls.rs

@@ -704,11 +704,85 @@ pub fn dispatch_grpc_call(
     }
 }
 
+extern "C" {
+    fn proxy_grpc_stream(
+        upstream_data: *const u8,
+        upstream_size: usize,
+        service_name_data: *const u8,
+        service_name_size: usize,
+        method_name_data: *const u8,
+        method_name_size: usize,
+        initial_metadata_data: *const u8,
+        initial_metadata_size: usize,
+        return_stream_id: *mut u32,
+    ) -> Status;
+}
+
+pub fn create_grpc_stream(
+    upstream_name: &str,
+    service_name: &str,
+    method_name: &str,
+    initial_metadata: Vec<(&str, &[u8])>,
+) -> Result<u32, Status> {
+    let mut return_stream_id = 0;
+    let serialized_initial_metadata = utils::serialize_bytes_value_map(initial_metadata);
+    unsafe {
+        match proxy_grpc_stream(
+            upstream_name.as_ptr(),
+            upstream_name.len(),
+            service_name.as_ptr(),
+            service_name.len(),
+            method_name.as_ptr(),
+            method_name.len(),
+            serialized_initial_metadata.as_ptr(),
+            serialized_initial_metadata.len(),
+            &mut return_stream_id,
+        ) {
+            Status::Ok => {
+                dispatcher::register_grpc_stream(return_stream_id);
+                Ok(return_stream_id)
+            }
+            Status::ParseFailure => Err(Status::ParseFailure),
+            Status::InternalFailure => Err(Status::InternalFailure),
+            status => panic!("unexpected status: {}", status as u32),
+        }
+    }
+}
+
+extern "C" {
+    fn proxy_grpc_send(
+        token: u32,
+        message_ptr: *const u8,
+        message_len: usize,
+        end_stream: bool,
+    ) -> Status;
+}
+
+pub fn grpc_stream_send(
+    token: u32,
+    message: Option<&[u8]>,
+    end_stream: bool,
+) -> Result<(), Status> {
+    unsafe {
+        match proxy_grpc_send(
+            token,
+            message.map_or(null(), |message| message.as_ptr()),
+            message.map_or(0, |message| message.len()),
+            end_stream,
+        ) {
+            Status::Ok => Ok(()),
+            Status::BadArgument => Err(Status::BadArgument),
+            Status::NotFound => Err(Status::NotFound),
+            status => panic!("unexpected status: {}", status as u32),
+        }
+    }
+}
+
 extern "C" {
     fn proxy_grpc_cancel(token_id: u32) -> Status;
 }
 
-pub fn cancel_grpc_call(token_id: u32) -> Result<(), Status> {
+pub fn grpc_call_cancel(token_id: u32) -> Result<(), Status> {
     unsafe {
         match proxy_grpc_cancel(token_id) {
             Status::Ok => Ok(()),
@@ -718,6 +792,20 @@ pub fn cancel_grpc_call(token_id: u32) -> Result<(), Status> {
     }
 }
 
+extern "C" {
+    fn proxy_grpc_close(token_id: u32) -> Status;
+}
+
+pub fn grpc_stream_close(token_id: u32) -> Result<(), Status> {
+    unsafe {
+        match proxy_grpc_close(token_id) {
+            Status::Ok => Ok(()),
+            Status::NotFound => Err(Status::NotFound),
+            status => panic!("unexpected status: {}", status as u32),
+        }
+    }
+}
+
 extern "C" {
     fn proxy_set_effective_context(context_id: u32) -> Status;
 }
@@ -893,4 +981,25 @@ mod utils {
         }
         map
     }
+
+    pub(super) fn deserialize_bytes_map(bytes: &[u8]) -> Vec<(String, Vec<u8>)> {
+        let mut map = Vec::new();
+        if bytes.is_empty() {
+            return map;
+        }
+        let size = u32::from_le_bytes(<[u8; 4]>::try_from(&bytes[0..4]).unwrap()) as usize;
+        let mut p = 4 + size * 8;
+        for n in 0..size {
+            let s = 4 + n * 8;
+            let size = u32::from_le_bytes(<[u8; 4]>::try_from(&bytes[s..s + 4]).unwrap()) as usize;
+            let key = bytes[p..p + size].to_vec();
+            p += size + 1;
+            let size =
+                u32::from_le_bytes(<[u8; 4]>::try_from(&bytes[s + 4..s + 8]).unwrap()) as usize;
+            let value = bytes[p..p + size].to_vec();
+            p += size + 1;
+            map.push((String::from_utf8(key).unwrap(), value));
+        }
+        map
+    }
 }

src/traits.rs

@@ -119,6 +119,45 @@ pub trait Context {
         hostcalls::cancel_grpc_call(token_id)
     }
 
+    fn create_grpc_stream(
+        &self,
+        cluster_name: &str,
+        service_name: &str,
+        method_name: &str,
+        initial_metadata: Vec<(&str, &[u8])>,
+    ) -> Result<u32, Status> {
+        hostcalls::create_grpc_stream(cluster_name, service_name, method_name, initial_metadata)
+    }
+
+    fn grpc_stream_send(
+        &self,
+        token_id: u32,
+        message: Option<&[u8]>,
+        end_stream: bool,
+    ) -> Result<(), Status> {
+        hostcalls::grpc_stream_send(token_id, message, end_stream)
+    }
+
+    fn grpc_stream_close(&self, token_id: u32) -> Result<(), Status> {
+        hostcalls::grpc_stream_close(token_id)
+    }
+
+    fn on_grpc_stream_receive_initial_metadata(&mut self, _token_id: u32, _headers: u32) {}
+
+    fn on_grpc_stream_receive_trailing_metadata(&mut self, _token_id: u32, _trailers: u32) {}
+
+    fn on_grpc_stream_receive_body(&mut self, _token_id: u32, _response_size: usize) {}
+
+    fn on_grpc_stream_close(&mut self, _token_id: u32, _status_code: u32) {}
+
+    fn get_grpc_call_initial_metadata(&self) -> Vec<(String, Vec<u8>)> {
+        hostcalls::get_map_bytes(MapType::GrpcReceiveInitialMetadata).unwrap()
+    }
+
+    fn get_grpc_call_trailing_metadata(&self) -> Vec<(String, Vec<u8>)> {
+        hostcalls::get_map_bytes(MapType::GrpcReceiveTrailingMetadata).unwrap()
+    }
+
     fn on_done(&mut self) -> bool {
         true
     }

src/types.rs

@@ -73,6 +73,8 @@ pub enum MapType {
     HttpRequestTrailers = 1,
     HttpResponseHeaders = 2,
     HttpResponseTrailers = 3,
+    GrpcReceiveInitialMetadata = 4,
+    GrpcReceiveTrailingMetadata = 5,
     HttpCallResponseHeaders = 6,
     HttpCallResponseTrailers = 7,
 }