support streaming update_client

Author: yoshidan

modules/service/src/elc.rs

@@ -1,13 +1,16 @@
 use crate::service::AppService;
 use enclave_api::EnclaveProtoAPI;
+use lcp_proto::google::protobuf::Any;
+use lcp_proto::lcp::service::elc::v1::msg_update_client_stream_chunk::Chunk;
 use lcp_proto::lcp::service::elc::v1::{
     msg_server::Msg, query_server::Query, MsgAggregateMessages, MsgAggregateMessagesResponse,
     MsgCreateClient, MsgCreateClientResponse, MsgUpdateClient, MsgUpdateClientResponse,
-    MsgVerifyMembership, MsgVerifyMembershipResponse, MsgVerifyNonMembership,
-    MsgVerifyNonMembershipResponse, QueryClientRequest, QueryClientResponse,
+    MsgUpdateClientStreamChunk, MsgVerifyMembership, MsgVerifyMembershipResponse,
+    MsgVerifyNonMembership, MsgVerifyNonMembershipResponse, QueryClientRequest,
+    QueryClientResponse,
 };
 use store::transaction::CommitStore;
-use tonic::{Request, Response, Status};
+use tonic::{Request, Response, Status, Streaming};
 
 #[tonic::async_trait]
 impl<E, S> Msg for AppService<E, S>
@@ -35,6 +38,69 @@ where
         }
     }
 
+    async fn update_client_stream(
+        &self,
+        request: Request<Streaming<MsgUpdateClientStreamChunk>>,
+    ) -> Result<Response<MsgUpdateClientResponse>, Status> {
+        let mut stream = request.into_inner();
+
+        // read the first message (must be Init)
+        let init = match stream.message().await? {
+            Some(chunk) => match chunk.chunk {
+                Some(Chunk::Init(init)) => init,
+                _ => {
+                    return Err(Status::invalid_argument(
+                        "first message must be of type Init",
+                    ))
+                }
+            },
+            None => {
+                return Err(Status::invalid_argument(
+                    "expected Init message as the first message",
+                ))
+            }
+        };
+
+        // accumulate header chunks
+        let mut header_bytes = Vec::new();
+
+        while let Some(chunk_msg) = stream.message().await? {
+            match chunk_msg.chunk {
+                Some(Chunk::HeaderChunk(header_chunk)) => {
+                    header_bytes.extend(header_chunk.data);
+                }
+                Some(Chunk::Init(_)) => {
+                    return Err(Status::invalid_argument(
+                        "Init must only appear as the first message",
+                    ));
+                }
+                None => {
+                    return Err(Status::invalid_argument("received empty chunk message"));
+                }
+            }
+        }
+
+        if header_bytes.is_empty() {
+            return Err(Status::invalid_argument("no header data received"));
+        }
+
+        // create MsgUpdateClient from Init and collected header data
+        let msg = MsgUpdateClient {
+            client_id: init.client_id,
+            include_state: init.include_state,
+            signer: init.signer,
+            header: Some(Any {
+                type_url: init.type_url,
+                value: header_bytes,
+            }),
+        };
+
+        match self.enclave.proto_update_client(msg) {
+            Ok(res) => Ok(Response::new(res)),
+            Err(e) => Err(Status::aborted(e.to_string())),
+        }
+    }
+
     async fn aggregate_messages(
         &self,
         request: Request<MsgAggregateMessages>,

proto/definitions/lcp/service/elc/v1/tx.proto

@@ -16,6 +16,9 @@ service Msg {
   // UpdateClient defines a rpc handler method for MsgUpdateClient.
   rpc UpdateClient(MsgUpdateClient) returns (MsgUpdateClientResponse);
 
+  // UpdateClientStream defines a rpc handler method for MsgUpdateClient.
+  rpc UpdateClientStream(stream MsgUpdateClientStreamChunk) returns (MsgUpdateClientResponse);
+
   // AggregateMessages defines a rpc handler method for MsgAggregateMessages
   rpc AggregateMessages(MsgAggregateMessages) returns (MsgAggregateMessagesResponse);
 
@@ -134,3 +137,21 @@ message MsgVerifyNonMembershipResponse {
   bytes message = 1;
   bytes signature = 2;
 }
+
+message MsgUpdateClientStreamChunk {
+  oneof chunk {
+    UpdateClientStreamInit init = 1;
+    UpdateClientStreamHeaderChunk header_chunk = 2;
+  }
+}
+
+message UpdateClientStreamInit {
+  string client_id = 1;
+  bool include_state = 2;
+  bytes signer = 3;
+  string type_url = 4;
+}
+
+message UpdateClientStreamHeaderChunk {
+  bytes data = 1;
+}

proto/src/prost/lcp.service.elc.v1.rs

@@ -407,6 +407,45 @@ pub struct MsgVerifyNonMembershipResponse {
     #[prost(bytes = "vec", tag = "2")]
     pub signature: ::prost::alloc::vec::Vec<u8>,
 }
+#[derive(::serde::Serialize, ::serde::Deserialize)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct MsgUpdateClientStreamChunk {
+    #[prost(oneof = "msg_update_client_stream_chunk::Chunk", tags = "1, 2")]
+    pub chunk: ::core::option::Option<msg_update_client_stream_chunk::Chunk>,
+}
+/// Nested message and enum types in `MsgUpdateClientStreamChunk`.
+pub mod msg_update_client_stream_chunk {
+    #[derive(::serde::Serialize, ::serde::Deserialize)]
+    #[allow(clippy::derive_partial_eq_without_eq)]
+    #[derive(Clone, PartialEq, ::prost::Oneof)]
+    pub enum Chunk {
+        #[prost(message, tag = "1")]
+        Init(super::UpdateClientStreamInit),
+        #[prost(message, tag = "2")]
+        HeaderChunk(super::UpdateClientStreamHeaderChunk),
+    }
+}
+#[derive(::serde::Serialize, ::serde::Deserialize)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct UpdateClientStreamInit {
+    #[prost(string, tag = "1")]
+    pub client_id: ::prost::alloc::string::String,
+    #[prost(bool, tag = "2")]
+    pub include_state: bool,
+    #[prost(bytes = "vec", tag = "3")]
+    pub signer: ::prost::alloc::vec::Vec<u8>,
+    #[prost(string, tag = "4")]
+    pub type_url: ::prost::alloc::string::String,
+}
+#[derive(::serde::Serialize, ::serde::Deserialize)]
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct UpdateClientStreamHeaderChunk {
+    #[prost(bytes = "vec", tag = "1")]
+    pub data: ::prost::alloc::vec::Vec<u8>,
+}
 /// Generated client implementations.
 #[cfg(feature = "client")]
 pub mod msg_client {
@@ -518,6 +557,30 @@ pub mod msg_client {
             );
             self.inner.unary(request.into_request(), path, codec).await
         }
+        /// UpdateClientStream defines a rpc handler method for MsgUpdateClient.
+        pub async fn update_client_stream(
+            &mut self,
+            request: impl tonic::IntoStreamingRequest<
+                Message = super::MsgUpdateClientStreamChunk,
+            >,
+        ) -> Result<tonic::Response<super::MsgUpdateClientResponse>, tonic::Status> {
+            self.inner
+                .ready()
+                .await
+                .map_err(|e| {
+                    tonic::Status::new(
+                        tonic::Code::Unknown,
+                        format!("Service was not ready: {}", e.into()),
+                    )
+                })?;
+            let codec = tonic::codec::ProstCodec::default();
+            let path = http::uri::PathAndQuery::from_static(
+                "/lcp.service.elc.v1.Msg/UpdateClientStream",
+            );
+            self.inner
+                .client_streaming(request.into_streaming_request(), path, codec)
+                .await
+        }
         /// AggregateMessages defines a rpc handler method for MsgAggregateMessages
         pub async fn aggregate_messages(
             &mut self,
@@ -604,6 +667,11 @@ pub mod msg_server {
             &self,
             request: tonic::Request<super::MsgUpdateClient>,
         ) -> Result<tonic::Response<super::MsgUpdateClientResponse>, tonic::Status>;
+        /// UpdateClientStream defines a rpc handler method for MsgUpdateClient.
+        async fn update_client_stream(
+            &self,
+            request: tonic::Request<tonic::Streaming<super::MsgUpdateClientStreamChunk>>,
+        ) -> Result<tonic::Response<super::MsgUpdateClientResponse>, tonic::Status>;
         /// AggregateMessages defines a rpc handler method for MsgAggregateMessages
         async fn aggregate_messages(
             &self,
@@ -759,6 +827,49 @@ pub mod msg_server {
                     };
                     Box::pin(fut)
                 }
+                "/lcp.service.elc.v1.Msg/UpdateClientStream" => {
+                    #[allow(non_camel_case_types)]
+                    struct UpdateClientStreamSvc<T: Msg>(pub Arc<T>);
+                    impl<
+                        T: Msg,
+                    > tonic::server::ClientStreamingService<
+                        super::MsgUpdateClientStreamChunk,
+                    > for UpdateClientStreamSvc<T> {
+                        type Response = super::MsgUpdateClientResponse;
+                        type Future = BoxFuture<
+                            tonic::Response<Self::Response>,
+                            tonic::Status,
+                        >;
+                        fn call(
+                            &mut self,
+                            request: tonic::Request<
+                                tonic::Streaming<super::MsgUpdateClientStreamChunk>,
+                            >,
+                        ) -> Self::Future {
+                            let inner = self.0.clone();
+                            let fut = async move {
+                                (*inner).update_client_stream(request).await
+                            };
+                            Box::pin(fut)
+                        }
+                    }
+                    let accept_compression_encodings = self.accept_compression_encodings;
+                    let send_compression_encodings = self.send_compression_encodings;
+                    let inner = self.inner.clone();
+                    let fut = async move {
+                        let inner = inner.0;
+                        let method = UpdateClientStreamSvc(inner);
+                        let codec = tonic::codec::ProstCodec::default();
+                        let mut grpc = tonic::server::Grpc::new(codec)
+                            .apply_compression_config(
+                                accept_compression_encodings,
+                                send_compression_encodings,
+                            );
+                        let res = grpc.client_streaming(method, req).await;
+                        Ok(res)
+                    };
+                    Box::pin(fut)
+                }
                 "/lcp.service.elc.v1.Msg/AggregateMessages" => {
                     #[allow(non_camel_case_types)]
                     struct AggregateMessagesSvc<T: Msg>(pub Arc<T>);