Move tonic-specific oak_client pieces to separate crate

Avoids the need for a feature flag in oak_client.

Change-Id: I51b3c0d612536e135cf8025768ab3dbd02c9d650

Cargo.toml

@@ -12,6 +12,7 @@ members = [
   "oak_attestation_verification",
   "oak_channel",
   "oak_client",
+  "oak_client/tonic",
   "oak_containers/agent",
   "oak_containers/examples/hello_world/untrusted_app",
   "oak_containers/examples/hello_world/web_client",
@@ -103,6 +104,7 @@ oak_attestation_explain = { path = "./oak_attestation_explain" }
 oak_attestation_verification = { path = "./oak_attestation_verification" }
 oak_channel = { path = "./oak_channel" }
 oak_client = { path = "./oak_client" }
+oak_client_tonic = { path = "./oak_client/tonic" }
 oak_containers_agent = { path = "./oak_containers/agent" }
 oak_containers_orchestrator = { path = "./oak_containers/orchestrator" }
 oak_containers_attestation = { path = "./oak_containers/attestation" }

oak_client/BUILD

@@ -29,9 +29,7 @@ rust_library(
         "src/transport.rs",
         "src/verifier.rs",
     ],
-    crate_features = [
-        "grpc_streaming_transport_implementation",
-    ],
+    crate_features = [],
     proc_macro_deps = [
         "@oak_crates_index//:async-trait",
     ],
@@ -43,7 +41,5 @@ rust_library(
         "@oak_crates_index//:anyhow",
         "@oak_crates_index//:futures",
         "@oak_crates_index//:futures-util",
-        "@oak_crates_index//:prost",
-        "@oak_crates_index//:tonic",
     ],
 )

oak_client/Cargo.toml

@@ -5,10 +5,6 @@ authors = ["Ivan Petrov <[email protected]>"]
 edition = "2021"
 license = "Apache-2.0"
 
-[features]
-default = ["grpc_streaming_transport_implementation"]
-grpc_streaming_transport_implementation = ["tonic"]
-
 [dependencies]
 anyhow = "*"
 async-trait = "*"

oak_client/src/transport.rs

@@ -14,88 +14,11 @@
 // limitations under the License.
 //
 
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-use std::future::Future;
-
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-use anyhow::Context;
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-use futures::channel::mpsc;
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-use oak_proto_rust::oak::session::v1::{
-    request_wrapper, response_wrapper, RequestWrapper, ResponseWrapper,
-};
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-use oak_proto_rust::oak::session::v1::{
-    GetEndorsedEvidenceRequest, GetEndorsedEvidenceResponse, InvokeRequest, InvokeResponse,
-};
 use oak_proto_rust::oak::{
     crypto::v1::{EncryptedRequest, EncryptedResponse},
     session::v1::EndorsedEvidence,
 };
 
-/// A [Transport] implementation that uses a single gRPC streaming session to
-/// get the evidence and then invokve the desired request.
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-pub struct GrpcStreamingTransport {
-    response_stream: tonic::Streaming<ResponseWrapper>,
-    request_tx_channel: mpsc::Sender<RequestWrapper>,
-}
-
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-impl GrpcStreamingTransport {
-    /// Create a new [GrpcStreamingTransport].
-    ///
-    /// The provided stream_creator will be immediately invokved to create a new
-    /// stream session, and all actions on the newly created instance will use
-    /// that stream. If the stream dies, a new client/transport pair will need
-    /// to be created.
-    ///
-    /// The `stream_creator` is a closure to create a new stream. Typically this
-    /// will be used with a gRPC client instance that exposes a
-    /// bi-directional streaming method.
-    ///
-    /// For example, if you have a gRPC service like:
-    /// ```
-    /// service {
-    ///     rpc MyMethod(stream RequestWrapper) returns (stream ResponseWrapper);
-    /// }
-    /// ```
-    ///
-    /// You will call:
-    /// ```
-    /// let transport =
-    ///   GrpcStreamingTransport::new(|rx| client.my_method(rx))
-    /// ```
-    pub async fn new<Fut>(
-        stream_creator: impl FnOnce(mpsc::Receiver<RequestWrapper>) -> Fut,
-    ) -> anyhow::Result<Self>
-    where
-        Fut: Future<Output = tonic::Result<tonic::Response<tonic::Streaming<ResponseWrapper>>>>,
-    {
-        let (tx, rx) = mpsc::channel(10);
-
-        let response_stream =
-            stream_creator(rx).await.context("couldn't send stream request")?.into_inner();
-
-        Ok(Self { response_stream, request_tx_channel: tx })
-    }
-
-    async fn send_and_receive(
-        &mut self,
-        request: request_wrapper::Request,
-    ) -> anyhow::Result<ResponseWrapper> {
-        self.request_tx_channel
-            .try_send(RequestWrapper { request: Some(request) })
-            .context("Couldn't send request")?;
-        self.response_stream
-            .message()
-            .await
-            .context("received empty response stream")?
-            .context("empty response")
-    }
-}
-
 #[async_trait::async_trait(?Send)]
 pub trait Transport {
     async fn invoke(
@@ -104,54 +27,7 @@ pub trait Transport {
     ) -> anyhow::Result<EncryptedResponse>;
 }
 
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-#[async_trait::async_trait(?Send)]
-impl Transport for GrpcStreamingTransport {
-    async fn invoke(
-        &mut self,
-        encrypted_request: &EncryptedRequest,
-    ) -> anyhow::Result<EncryptedResponse> {
-        let response_wrapper: ResponseWrapper = self
-            .send_and_receive(request_wrapper::Request::InvokeRequest(InvokeRequest {
-                encrypted_request: Some(encrypted_request.clone()),
-            }))
-            .await
-            .context("Sending invoke request")?;
-
-        match response_wrapper.response {
-            Some(response_wrapper::Response::InvokeResponse(InvokeResponse {
-                encrypted_response: Some(encrypted_response),
-            })) => Ok(encrypted_response),
-            _ => Err(anyhow::anyhow!(
-                "response_wrapper does not have a valid invoke_response message"
-            )),
-        }
-    }
-}
-
 #[async_trait::async_trait(?Send)]
 pub trait EvidenceProvider {
     async fn get_endorsed_evidence(&mut self) -> anyhow::Result<EndorsedEvidence>;
 }
-
-#[cfg(feature = "grpc_streaming_transport_implementation")]
-#[async_trait::async_trait(?Send)]
-impl EvidenceProvider for GrpcStreamingTransport {
-    async fn get_endorsed_evidence(&mut self) -> anyhow::Result<EndorsedEvidence> {
-        let response_wrapper = self
-            .send_and_receive(request_wrapper::Request::GetEndorsedEvidenceRequest(
-                GetEndorsedEvidenceRequest {},
-            ))
-            .await
-            .context("Sending get_endorsed_evidence request")?;
-
-        match response_wrapper.response {
-            Some(response_wrapper::Response::GetEndorsedEvidenceResponse(
-                GetEndorsedEvidenceResponse { endorsed_evidence: Some(endorsed_evidence) },
-            )) => Ok(endorsed_evidence),
-            _ => Err(anyhow::anyhow!(
-                "response_wrapper does not have a valid invoke_response message"
-            )),
-        }
-    }
-}

oak_client/tonic/BUILD

@@ -0,0 +1,44 @@
+#
+# Copyright 2024 The Project Oak Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+load("@rules_rust//rust:defs.bzl", "rust_library")
+
+package(
+    default_visibility = ["//:internal"],
+    licenses = ["notice"],
+)
+
+rust_library(
+    name = "oak_client_tonic",
+    srcs = [
+        "lib.rs",
+        "transport.rs",
+    ],
+    proc_macro_deps = [
+        "@oak_crates_index//:async-trait",
+    ],
+    deps = [
+        "//oak_attestation_verification",
+        "//oak_client",
+        "//oak_crypto",
+        "//oak_proto_rust",
+        "//oak_proto_rust/grpc",
+        "@oak_crates_index//:anyhow",
+        "@oak_crates_index//:futures",
+        "@oak_crates_index//:futures-util",
+        "@oak_crates_index//:tonic",
+    ],
+)

oak_client/tonic/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "oak_client_tonic"
+version = "0.1.0"
+authors = ["Ivan Petrov <[email protected]>"]
+edition = "2021"
+license = "Apache-2.0"
+
+[lib]
+path = "lib.rs"
+
+[dependencies]
+anyhow = "*"
+async-trait = "*"
+futures = "*"
+log = "*"
+oak_client = { workspace = true }
+oak_proto_rust = { workspace = true }
+prost = { workspace = true }
+tonic = { workspace = true }

oak_client/tonic/lib.rs

@@ -0,0 +1,16 @@
+//
+// Copyright 2024 The Project Oak Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pub mod transport;