diff --git a/oak_containers/examples/hello_world/host_app/BUILD b/oak_containers/examples/hello_world/host_app/BUILD index 9ecd90c036..b0f6fbd1eb 100644 --- a/oak_containers/examples/hello_world/host_app/BUILD +++ b/oak_containers/examples/hello_world/host_app/BUILD @@ -24,7 +24,6 @@ package( rust_library( name = "oak_containers_examples_hello_world_host_app", srcs = [ - "src/app_client.rs", "src/demo_transport.rs", "src/http_service.rs", "src/launcher_args.rs", diff --git a/oak_containers/examples/hello_world/host_app/src/app_client.rs b/oak_containers/examples/hello_world/host_app/src/app_client.rs deleted file mode 100644 index 8eca9e4362..0000000000 --- a/oak_containers/examples/hello_world/host_app/src/app_client.rs +++ /dev/null @@ -1,53 +0,0 @@ -// -// Copyright 2023 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. - -use anyhow::Context; -use oak_hello_world_proto::oak::containers::example::enclave_application_client::EnclaveApplicationClient as GrpcEnclaveApplicationClient; -use oak_proto_rust::oak::session::v1::{RequestWrapper, ResponseWrapper}; -use tokio::time::Duration; -use tonic::transport::Endpoint; - -/// Utility struct used to interface with the launcher -pub struct EnclaveApplicationClient { - inner: GrpcEnclaveApplicationClient, -} - -impl EnclaveApplicationClient { - pub async fn create(server_addr: String) -> Result> { - let inner: GrpcEnclaveApplicationClient = { - let channel = Endpoint::from_shared(server_addr) - .context("couldn't form channel")? - .connect_timeout(Duration::from_secs(120)) - .connect() - .await - .context("couldn't connect to enclave app")?; - GrpcEnclaveApplicationClient::new(channel) - }; - Ok(Self { inner }) - } - - pub async fn legacy_session( - &mut self, - request: impl tonic::IntoStreamingRequest, - ) -> anyhow::Result> { - Ok(self - .inner - // How to safely map this request stream? - .legacy_session(request) - .await - .context("couldn't send hello request")? - .into_inner()) - } -} diff --git a/oak_containers/examples/hello_world/host_app/src/http_service.rs b/oak_containers/examples/hello_world/host_app/src/http_service.rs index 508825482d..3df3792830 100644 --- a/oak_containers/examples/hello_world/host_app/src/http_service.rs +++ b/oak_containers/examples/hello_world/host_app/src/http_service.rs @@ -22,24 +22,25 @@ use bytes::Bytes; use http_body_util::{BodyExt, Full}; use hyper::{body, server::conn::http1, service::service_fn, Request, Response}; use hyper_util::rt::{TokioIo, TokioTimer}; +use oak_hello_world_proto::oak::containers::example::enclave_application_client::EnclaveApplicationClient; use oak_proto_rust::oak::session::v1::{RequestWrapper, ResponseWrapper}; use prost::Message; -use tokio::{net::TcpListener, sync::Mutex}; - -use crate::app_client::EnclaveApplicationClient; +use tokio::{net::TcpListener, sync::Mutex, time::Duration}; +use tonic::transport::{channel::Channel, Endpoint}; async fn handle_request( request: RequestWrapper, - enclave_app: Arc>, + enclave_app: Arc>>, ) -> tonic::Result { // This is not how we should actually use the streaming interface, but it // works for HPKE, as long as all requests go to the same machine. - let mut response_stream = + let response_stream = enclave_app.lock().await.legacy_session(tokio_stream::iter(vec![request])).await.map_err( |err| tonic::Status::internal(format!("starting streaming session failed: {err:?}")), )?; response_stream + .into_inner() .message() .await? .context("no response wrapper was returned") @@ -58,9 +59,13 @@ pub async fn serve( .get_trusted_app_address() .await .map_err(|error| anyhow!("Failed to get app address: {error:?}"))?; - let app_client = EnclaveApplicationClient::create(format!("http://{enclave_app_address}")) + let channel = Endpoint::from_shared(format!("http://{enclave_app_address}")) + .context("couldn't form channel")? + .connect_timeout(Duration::from_secs(120)) + .connect() .await - .map_err(|error| anyhow!("Failed to create enclave application client: {error:?}"))?; + .context("couldn't connect to enclave app")?; + let app_client = EnclaveApplicationClient::new(channel); let app_client = Arc::new(Mutex::new(app_client)); diff --git a/oak_containers/examples/hello_world/host_app/src/lib.rs b/oak_containers/examples/hello_world/host_app/src/lib.rs index 83eb137a8d..4174f8f68a 100644 --- a/oak_containers/examples/hello_world/host_app/src/lib.rs +++ b/oak_containers/examples/hello_world/host_app/src/lib.rs @@ -13,7 +13,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod app_client; pub mod demo_transport; pub mod http_service; pub mod launcher_args; diff --git a/oak_containers/examples/hello_world/host_app/src/service.rs b/oak_containers/examples/hello_world/host_app/src/service.rs index 1d5b4e9976..07122644cb 100644 --- a/oak_containers/examples/hello_world/host_app/src/service.rs +++ b/oak_containers/examples/hello_world/host_app/src/service.rs @@ -15,24 +15,24 @@ use std::{future::Future, pin::Pin, sync::Arc}; -use anyhow::anyhow; +use anyhow::{anyhow, Context}; use futures::{channel::mpsc, Stream, StreamExt}; -use oak_hello_world_proto::oak::containers::example::host_application_server::{ - HostApplication, HostApplicationServer, +use oak_hello_world_proto::oak::containers::example::{ + enclave_application_client::EnclaveApplicationClient, + host_application_server::{HostApplication, HostApplicationServer}, }; use oak_proto_rust::oak::session::v1::{RequestWrapper, ResponseWrapper}; -use tokio::{net::TcpListener, sync::Mutex}; +use tokio::{net::TcpListener, sync::Mutex, time::Duration}; use tokio_stream::wrappers::TcpListenerStream; - -use crate::app_client::EnclaveApplicationClient; +use tonic::transport::{channel::Channel, Endpoint}; /// The sample application's implementation of Oak's streaming service protocol. struct HostApplicationImpl { - enclave_app: Arc>, + enclave_app: Arc>>, } impl HostApplicationImpl { - pub fn new(enclave_app: EnclaveApplicationClient) -> Self { + pub fn new(enclave_app: EnclaveApplicationClient) -> Self { Self { enclave_app: Arc::new(Mutex::new(enclave_app)) } } } @@ -48,12 +48,12 @@ async fn forward_stream( upstream_starter: impl FnOnce(mpsc::Receiver) -> Fut, ) -> Result>, tonic::Status> where - Fut: Future, tonic::Status>>, + Fut: Future>, tonic::Status>>, { let mut request_stream = request_stream; let (mut tx, rx) = mpsc::channel(10); - let mut upstream = upstream_starter(rx).await?; + let mut upstream = upstream_starter(rx).await?.into_inner(); Ok(async_stream::try_stream! { loop { @@ -120,9 +120,13 @@ pub async fn create( .get_trusted_app_address() .await .map_err(|error| anyhow!("Failed to get app address: {error:?}"))?; - let app_client = EnclaveApplicationClient::create(format!("http://{enclave_app_address}")) + let channel = Endpoint::from_shared(format!("http://{enclave_app_address}")) + .context("couldn't form channel")? + .connect_timeout(Duration::from_secs(120)) + .connect() .await - .map_err(|error| anyhow!("Failed to create enclave application client: {error:?}"))?; + .context("couldn't connect to enclave app")?; + let app_client = EnclaveApplicationClient::new(channel); tonic::transport::Server::builder() .add_service(HostApplicationServer::new(HostApplicationImpl::new(app_client))) .serve_with_incoming(TcpListenerStream::new(listener))