Extract boilerplate from hello_world example into SDK

This CR refactors the hello world example to move as much functionality
as possible into the Oak SDK.

Now to start your Oak trusted application, you need to:
1) Create an ApplicationHandler that accepts opaque bytes, and returns
   opaque bytes (which is all that Oak cares about). The interpretation
   is left to the application designer and is outside of the scope of Oak.

2) On startup, create an OakHandler containing your application handler,
   the endorsed evidence provided by the orchestrator (or similar), and
   an InstanceEncryptionKeyHandle.

3) Implement the service using the library of your choice (we use tonic
   in our example). In the handler for the oak-enabled method, you just
   need to write:

   `oak_session(self.oak_handler.clone(), request).await`

   This will start the stream processing, and either handle oak messages
   or pass application messages to your handler.

* Hello world app is split into app, and app_service
  * app_service.rs is just the generic service implementation, including a
    call into the Oak server SDK helper.
  * app.rs contains the actual application business logic.
  * Notice the reduced amount of boilerplate in app_service.

* oak_handler.rs contains server-tech-independent Oak message processing
  logic.

* tonic.rs contains the helper to connect tonic requests to the
  oak_handler logic.

Bug: b/356858826
Change-Id: I5dc3d95605f7ec8111190174a432c0a80dc7fbc3

oak_containers_examples/hello_world/trusted_app/BUILD

@@ -27,10 +27,12 @@ package(
 rust_library(
     name = "lib",
     srcs = [
+        "src/app.rs",
         "src/app_service.rs",
         "src/lib.rs",
     ],
     crate_name = "oak_containers_hello_world_trusted_app",
+    proc_macro_deps = ["@oak_crates_index//:async-trait"],
     deps = [
         "//oak_containers_examples/hello_world/proto:hello_world_rust_proto",
         "//oak_containers_sdk",

oak_containers_examples/hello_world/trusted_app/src/app.rs

@@ -0,0 +1,38 @@
+//
+// 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 oak_containers_sdk::ApplicationHandler;
+
+/// The actual business logic for the hello world application.
+pub struct HelloWorldApplicationHandler {
+    pub application_config: Vec<u8>,
+}
+
+#[async_trait::async_trait]
+impl ApplicationHandler for HelloWorldApplicationHandler {
+    /// This implementation is quite simple, since there's just a single request
+    /// that is a string. In a real implementation, we'd probably
+    /// deserialize into a proto, and dispatch to various handlers from
+    /// there.
+    async fn handle(&self, request_bytes: &[u8]) -> anyhow::Result<Vec<u8>> {
+        let name = String::from_utf8_lossy(request_bytes);
+        let config_len = self.application_config.len();
+        Ok(
+            format!(
+                "Hello from the trusted side, {name}! Btw, the Trusted App has a config with a length of {config_len} bytes.",
+            ).into_bytes()
+        )
+    }
+}

oak_containers_examples/hello_world/trusted_app/src/app_service.rs

@@ -13,144 +13,53 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{pin::Pin, sync::Arc};
+/// App service shows an example of the boilerplate needed to implement the
+/// Oak-related machinery for session message stream handling.
+///
+/// Oak doesn't provide full service implementation conveniences in the SDK,
+/// since typically, Oak-enabled methods will be part of a larger service that
+/// may need to be configured in specific ways.
+use std::sync::Arc;
 
 use anyhow::anyhow;
-use oak_containers_sdk::InstanceEncryptionKeyHandle;
-use oak_crypto::encryptor::ServerEncryptor;
+use oak_containers_sdk::OakSessionContext;
 use oak_hello_world_proto::oak::containers::example::trusted_application_server::{
     TrustedApplication, TrustedApplicationServer,
 };
-use oak_proto_rust::oak::{
-    crypto::v1::{EncryptedRequest, EncryptedResponse},
-    session::v1::{
-        request_wrapper, response_wrapper, EndorsedEvidence, GetEndorsedEvidenceResponse,
-        InvokeRequest, InvokeResponse, RequestWrapper, ResponseWrapper,
-    },
-};
+use oak_proto_rust::oak::session::v1::RequestWrapper;
 use tokio::net::TcpListener;
-use tokio_stream::{wrappers::TcpListenerStream, Stream, StreamExt};
-
-const EMPTY_ASSOCIATED_DATA: &[u8] = b"";
-
-// Things that need to be available to every incoming streaming session.
-struct SessionContext {
-    application_config: Vec<u8>,
-    encryption_key_handle: InstanceEncryptionKeyHandle,
-    endorsed_evidence: EndorsedEvidence,
-}
+use tokio_stream::wrappers::TcpListenerStream;
 
+/// The struct that will hold the gRPC TrustedApplication implementation.
 struct TrustedApplicationImplementation {
-    session_context: Arc<SessionContext>,
+    oak_session_context: Arc<OakSessionContext>,
 }
 
 impl TrustedApplicationImplementation {
-    pub fn new(
-        application_config: Vec<u8>,
-        encryption_key_handle: InstanceEncryptionKeyHandle,
-        endorsed_evidence: EndorsedEvidence,
-    ) -> Self {
-        Self {
-            session_context: Arc::new(SessionContext {
-                application_config,
-                encryption_key_handle,
-                endorsed_evidence,
-            }),
-        }
-    }
-}
-
-impl SessionContext {
-    fn handle_hello(&self, request_bytes: &[u8]) -> String {
-        let name = String::from_utf8_lossy(request_bytes);
-        format!(
-            "Hello from the trusted side, {}! Btw, the Trusted App has a config with a length of {} bytes.",
-            name,
-            self.application_config.len()
-        )
-    }
-
-    async fn handle_request(
-        &self,
-        encrypted_request: Option<EncryptedRequest>,
-    ) -> tonic::Result<EncryptedResponse> {
-        let encrypted_request = encrypted_request
-            .ok_or(tonic::Status::internal("encrypted request wasn't provided"))?;
-
-        // Associated data is ignored.
-        let (server_encryptor, name_bytes, _) =
-            ServerEncryptor::decrypt_async(&encrypted_request, &self.encryption_key_handle)
-                .await
-                .map_err(|error| {
-                tonic::Status::internal(format!("couldn't decrypt request: {:?}", error))
-            })?;
-
-        let response = self.handle_hello(&name_bytes);
-
-        server_encryptor.encrypt(response.as_bytes(), EMPTY_ASSOCIATED_DATA).map_err(|error| {
-            tonic::Status::internal(format!("couldn't encrypt response: {:?}", error))
-        })
+    pub fn new(oak_session_context: OakSessionContext) -> Self {
+        Self { oak_session_context: Arc::new(oak_session_context) }
     }
 }
 
 #[tonic::async_trait]
 impl TrustedApplication for TrustedApplicationImplementation {
-    type SessionStream =
-        Pin<Box<dyn Stream<Item = Result<ResponseWrapper, tonic::Status>> + Send + 'static>>;
+    type SessionStream = oak_containers_sdk::tonic::OakSessionStream;
 
     async fn session(
         &self,
         request: tonic::Request<tonic::Streaming<RequestWrapper>>,
     ) -> Result<tonic::Response<Self::SessionStream>, tonic::Status> {
-        let mut request_stream = request.into_inner();
-
-        let session_context = self.session_context.clone();
-
-        let response_stream = async_stream::try_stream! {
-            while let Some(request) = request_stream.next().await {
-                let request = request
-                    .map_err(|err| {
-                        tonic::Status::internal(format!("error reading message from request stream: {err}"))
-                    })?
-                    .request
-                    .ok_or_else(|| tonic::Status::invalid_argument("empty request message"))?;
-
-                let response = match request {
-                    request_wrapper::Request::GetEndorsedEvidenceRequest(_) => {
-                        response_wrapper::Response::GetEndorsedEvidenceResponse(GetEndorsedEvidenceResponse {
-                            endorsed_evidence: Some(session_context.endorsed_evidence.clone()),
-                        })
-                    }
-                    request_wrapper::Request::InvokeRequest(InvokeRequest { encrypted_request }) => {
-                        let encrypted_response = session_context.handle_request(encrypted_request)
-                            .await
-                            .map_err(|err| tonic::Status::internal(format!("hello failed: {err:?}")))?;
-                        response_wrapper::Response::InvokeResponse(
-                            InvokeResponse { encrypted_response: Some(encrypted_response) }
-                        )
-                    }
-                };
-                yield ResponseWrapper {
-                    response: Some(response),
-                };
-            }
-        };
-
-        Ok(tonic::Response::new(Box::pin(response_stream) as Self::SessionStream))
+        oak_containers_sdk::tonic::oak_session(self.oak_session_context.clone(), request).await
     }
 }
 
 pub async fn create(
     listener: TcpListener,
-    application_config: Vec<u8>,
-    encryption_key_handle: InstanceEncryptionKeyHandle,
-    endorsed_evidence: EndorsedEvidence,
+    oak_session_context: OakSessionContext,
 ) -> Result<(), anyhow::Error> {
     tonic::transport::Server::builder()
         .add_service(TrustedApplicationServer::new(TrustedApplicationImplementation::new(
-            application_config,
-            encryption_key_handle,
-            endorsed_evidence,
+            oak_session_context,
         )))
         .serve_with_incoming(TcpListenerStream::new(listener))
         .await

oak_containers_examples/hello_world/trusted_app/src/lib.rs

@@ -13,4 +13,5 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod app;
 pub mod app_service;

oak_containers_examples/hello_world/trusted_app/src/main.rs

@@ -16,7 +16,7 @@
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
 use anyhow::Context;
-use oak_containers_sdk::{InstanceEncryptionKeyHandle, OrchestratorClient};
+use oak_containers_sdk::{InstanceEncryptionKeyHandle, OakSessionContext, OrchestratorClient};
 use tokio::net::TcpListener;
 
 const TRUSTED_APP_PORT: u16 = 8080;
@@ -41,13 +41,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let encryption_key_handle = InstanceEncryptionKeyHandle::create()
         .await
         .context("couldn't create encryption key handle: {:?}")?;
+
     let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), TRUSTED_APP_PORT);
     let listener = TcpListener::bind(addr).await?;
+
     let join_handle = tokio::spawn(oak_containers_hello_world_trusted_app::app_service::create(
         listener,
-        application_config,
-        encryption_key_handle,
-        endorsed_evidence,
+        OakSessionContext::new(
+            encryption_key_handle,
+            endorsed_evidence,
+            Box::new(oak_containers_hello_world_trusted_app::app::HelloWorldApplicationHandler {
+                application_config,
+            }),
+        ),
     ));
     orchestrator_client.notify_app_ready().await.context("failed to notify that app is ready")?;
     join_handle.await??;

oak_containers_sdk/BUILD

@@ -26,7 +26,9 @@ rust_library(
     srcs = [
         "src/crypto.rs",
         "src/lib.rs",
+        "src/oak_session_context.rs",
         "src/orchestrator_client.rs",
+        "src/tonic.rs",
     ],
     proc_macro_deps = [
         "@oak_crates_index//:async-trait",
@@ -36,6 +38,7 @@ rust_library(
         "//oak_proto_rust",
         "//oak_proto_rust/grpc",
         "@oak_crates_index//:anyhow",
+        "@oak_crates_index//:async-stream",
         "@oak_crates_index//:prost",
         "@oak_crates_index//:prost-types",
         "@oak_crates_index//:tokio",

oak_containers_sdk/src/lib.rs

@@ -14,7 +14,9 @@
 // limitations under the License.
 
 pub mod crypto;
+pub mod oak_session_context;
 pub mod orchestrator_client;
+pub mod tonic;
 
 // Unix Domain Sockets do not use URIs, hence this URI will never be used.
 // It is defined purely since in order to create a channel, since a URI has to
@@ -29,4 +31,5 @@ const ORCHESTRATOR_IPC_SOCKET: &str = "/oak_utils/orchestrator_ipc";
 
 // Re-export structs so that they are available at the top level of the SDK.
 pub use crypto::InstanceEncryptionKeyHandle;
+pub use oak_session_context::{ApplicationHandler, OakSessionContext};
 pub use orchestrator_client::OrchestratorClient;