From 090eb274e75c4a23abc7c4006f7e4b55c174d5d9 Mon Sep 17 00:00:00 2001 From: Jason LeBrun Date: Thu, 1 Aug 2024 19:00:53 +0000 Subject: [PATCH] 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 --- .../hello_world/trusted_app/BUILD | 2 + .../hello_world/trusted_app/src/app.rs | 38 ++++++ .../trusted_app/src/app_service.rs | 127 +++--------------- .../hello_world/trusted_app/src/lib.rs | 1 + .../hello_world/trusted_app/src/main.rs | 14 +- oak_containers_sdk/BUILD | 3 + oak_containers_sdk/src/lib.rs | 3 + oak_containers_sdk/src/oak_session_context.rs | 116 ++++++++++++++++ oak_containers_sdk/src/tonic.rs | 48 +++++++ 9 files changed, 239 insertions(+), 113 deletions(-) create mode 100644 oak_containers_examples/hello_world/trusted_app/src/app.rs create mode 100644 oak_containers_sdk/src/oak_session_context.rs create mode 100644 oak_containers_sdk/src/tonic.rs diff --git a/oak_containers_examples/hello_world/trusted_app/BUILD b/oak_containers_examples/hello_world/trusted_app/BUILD index 246ff235f6..3bd71d6d68 100644 --- a/oak_containers_examples/hello_world/trusted_app/BUILD +++ b/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", diff --git a/oak_containers_examples/hello_world/trusted_app/src/app.rs b/oak_containers_examples/hello_world/trusted_app/src/app.rs new file mode 100644 index 0000000000..b128124a32 --- /dev/null +++ b/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, +} + +#[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> { + 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() + ) + } +} diff --git a/oak_containers_examples/hello_world/trusted_app/src/app_service.rs b/oak_containers_examples/hello_world/trusted_app/src/app_service.rs index 88ecf51486..7e9133d590 100644 --- a/oak_containers_examples/hello_world/trusted_app/src/app_service.rs +++ b/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, - 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, + oak_session_context: Arc, } impl TrustedApplicationImplementation { - pub fn new( - application_config: Vec, - 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, - ) -> tonic::Result { - 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> + Send + 'static>>; + type SessionStream = oak_containers_sdk::tonic::OakSessionStream; async fn session( &self, request: tonic::Request>, ) -> Result, 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, - 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 diff --git a/oak_containers_examples/hello_world/trusted_app/src/lib.rs b/oak_containers_examples/hello_world/trusted_app/src/lib.rs index 522fa2fd19..cf16a54c1e 100644 --- a/oak_containers_examples/hello_world/trusted_app/src/lib.rs +++ b/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; diff --git a/oak_containers_examples/hello_world/trusted_app/src/main.rs b/oak_containers_examples/hello_world/trusted_app/src/main.rs index 7e06a4c24a..4380ea3386 100644 --- a/oak_containers_examples/hello_world/trusted_app/src/main.rs +++ b/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> { 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??; diff --git a/oak_containers_sdk/BUILD b/oak_containers_sdk/BUILD index 9e4d59bc1b..7190e3b696 100644 --- a/oak_containers_sdk/BUILD +++ b/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", diff --git a/oak_containers_sdk/src/lib.rs b/oak_containers_sdk/src/lib.rs index 1c490a5188..12478bab8a 100644 --- a/oak_containers_sdk/src/lib.rs +++ b/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; diff --git a/oak_containers_sdk/src/oak_session_context.rs b/oak_containers_sdk/src/oak_session_context.rs new file mode 100644 index 0000000000..e4fb9434a1 --- /dev/null +++ b/oak_containers_sdk/src/oak_session_context.rs @@ -0,0 +1,116 @@ +// +// 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_crypto::encryptor::ServerEncryptor; +use oak_proto_rust::oak::{ + crypto::v1::{EncryptedRequest, EncryptedResponse}, + session::v1::{ + request_wrapper, response_wrapper, EndorsedEvidence, GetEndorsedEvidenceResponse, + InvokeRequest, InvokeResponse, RequestWrapper, ResponseWrapper, + }, +}; + +const EMPTY_ASSOCIATED_DATA: &[u8] = b""; + +use crate::InstanceEncryptionKeyHandle; + +/// An Oak trusted application will write an implementation of +/// `ApplicationHandler` that accepts a serialized request (including di) +#[async_trait::async_trait] +pub trait ApplicationHandler: Send + Sync { + async fn handle(&self, request_bytes: &[u8]) -> anyhow::Result>; +} + +/// The state heeded to handle an Oak crypto session inside of a streaming +/// server handler. This structure contains the +/// server-implementation-independent functionality for encryption, and holds +/// the instance of the class that actually implements the application +/// logic. +pub struct OakSessionContext { + encryption_key_handle: InstanceEncryptionKeyHandle, + endorsed_evidence: EndorsedEvidence, + application_handler: Box, +} + +impl OakSessionContext { + pub fn new( + encryption_key_handle: InstanceEncryptionKeyHandle, + endorsed_evidence: EndorsedEvidence, + application_handler: Box, + ) -> Self { + Self { encryption_key_handle, endorsed_evidence, application_handler } + } + + pub fn endorsed_evidence(&self) -> &EndorsedEvidence { + &self.endorsed_evidence + } + + /// Decrypts the request, passes it to the application handler, and + /// encrypts the resulting response. + async fn handle_encrypted_request( + &self, + encrypted_request: &Option, + ) -> anyhow::Result { + let encrypted_request = + encrypted_request.as_ref().context("encrypted request wasn't present")?; + + // Associated data is ignored. + let (server_encryptor, name_bytes, _) = + ServerEncryptor::decrypt_async(encrypted_request, &self.encryption_key_handle) + .await + .context("couldn't decrypt request")?; + + let response = self + .application_handler + .handle(&name_bytes) + .await + .context("application handler failed")?; + + server_encryptor + .encrypt(&response, EMPTY_ASSOCIATED_DATA) + .context("couldn't encrypt response") + } + + /// Handles an incoming Oak request by either handling it internally if it's + /// an Oak message, or passing it to the provided application handler if + /// it's an application message. + pub(crate) async fn handle_request( + &self, + request_wrapper: &RequestWrapper, + ) -> anyhow::Result { + let request = &request_wrapper + .request + .as_ref() + .ok_or_else(|| anyhow::anyhow!("empty request message"))?; + + let response = match request { + request_wrapper::Request::GetEndorsedEvidenceRequest(_) => { + response_wrapper::Response::GetEndorsedEvidenceResponse( + GetEndorsedEvidenceResponse { + endorsed_evidence: Some(self.endorsed_evidence().clone()), + }, + ) + } + request_wrapper::Request::InvokeRequest(InvokeRequest { encrypted_request }) => { + let encrypted_response = self.handle_encrypted_request(encrypted_request).await?; + response_wrapper::Response::InvokeResponse(InvokeResponse { + encrypted_response: Some(encrypted_response), + }) + } + }; + Ok(ResponseWrapper { response: Some(response) }) + } +} diff --git a/oak_containers_sdk/src/tonic.rs b/oak_containers_sdk/src/tonic.rs new file mode 100644 index 0000000000..7326aa493f --- /dev/null +++ b/oak_containers_sdk/src/tonic.rs @@ -0,0 +1,48 @@ +// +// 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. + +use std::{pin::Pin, sync::Arc}; + +use oak_proto_rust::oak::session::v1::{RequestWrapper, ResponseWrapper}; +use tokio_stream::{Stream, StreamExt}; + +use crate::oak_session_context::OakSessionContext; + +/// Helper for handling streaming requests presented by a tonic server +/// implementation. + +pub type OakSessionStream = + Pin> + Send + 'static>>; + +pub async fn oak_session( + session_context: Arc, + request: tonic::Request>, +) -> Result, tonic::Status> { + let mut request_stream = request.into_inner(); + + let response_stream = async_stream::try_stream! { + while let Some(request) = request_stream.next().await { + let request_wrapper = request + .map_err(|e| tonic::Status::internal(format!("failed to read request from stream: {e:?}")))?; + + yield session_context + .handle_request(&request_wrapper) + .await + .map_err(|e| tonic::Status::internal(format!("failed to handle request: {e:?}")))?; + } + }; + + Ok(tonic::Response::new(Box::pin(response_stream) as OakSessionStream)) +}