From cffb2bfb69ba9ed1654bffdfe857c8a2caf3667c Mon Sep 17 00:00:00 2001 From: Jason LeBrun Date: Wed, 12 Feb 2025 17:11:20 +0000 Subject: [PATCH] Move Orchestrator Channel creation to helper in oak_sdk/containers Provides a consistent interface for creating a new Oak Orchestrator client channel. This reduces the number of channels created on startup, and simplifies constructors for many other types. Bug: b/357921050 Change-Id: Id7e06df3704778f67ec8514dd93807ffe13c0373 --- .../examples/hello_world/enclave_app/BUILD | 1 + .../hello_world/enclave_app/src/main.rs | 10 ++--- oak_containers_sdk/BUILD | 2 +- oak_containers_sdk/src/crypto.rs | 37 +++-------------- oak_containers_sdk/src/lib.rs | 11 ----- oak_containers_sdk/src/orchestrator_client.rs | 21 ++-------- oak_functions_containers_app/BUILD | 2 + oak_functions_containers_app/src/main.rs | 16 ++++---- oak_sdk/containers/BUILD | 41 +++++++++++++++++++ oak_sdk/containers/lib.rs | 17 ++++++++ oak_sdk/containers/orchestrator_channel.rs | 38 +++++++++++++++++ 11 files changed, 122 insertions(+), 74 deletions(-) create mode 100644 oak_sdk/containers/BUILD create mode 100644 oak_sdk/containers/lib.rs create mode 100644 oak_sdk/containers/orchestrator_channel.rs diff --git a/oak_containers/examples/hello_world/enclave_app/BUILD b/oak_containers/examples/hello_world/enclave_app/BUILD index 724ecd689a..90f3e32c20 100644 --- a/oak_containers/examples/hello_world/enclave_app/BUILD +++ b/oak_containers/examples/hello_world/enclave_app/BUILD @@ -78,6 +78,7 @@ rust_binary( deps = [ ":oak_containers_examples_hello_world_enclave_app", "//oak_containers_sdk", + "//oak_sdk/containers:oak_sdk_containers", "//oak_sdk/server/v1:oak_sdk_server_v1", "//oak_session", "@oak_crates_index//:anyhow", diff --git a/oak_containers/examples/hello_world/enclave_app/src/main.rs b/oak_containers/examples/hello_world/enclave_app/src/main.rs index 888b70b4c7..df3060a911 100644 --- a/oak_containers/examples/hello_world/enclave_app/src/main.rs +++ b/oak_containers/examples/hello_world/enclave_app/src/main.rs @@ -17,6 +17,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use anyhow::Context; use oak_containers_sdk::{InstanceEncryptionKeyHandle, OrchestratorClient}; +use oak_sdk_containers::default_orchestrator_channel; use oak_sdk_server_v1::OakApplicationContext; use tokio::net::TcpListener; @@ -25,9 +26,10 @@ const ENCLAVE_APP_PORT: u16 = 8080; #[tokio::main] async fn main() -> Result<(), Box> { println!("Logging!"); + let orchestrator_channel = + default_orchestrator_channel().await.context("failed to create orchestrator channel")?; - let mut orchestrator_client = - OrchestratorClient::create().await.context("Could't create orchestrator client")?; + let mut orchestrator_client = OrchestratorClient::create(&orchestrator_channel); let application_config = orchestrator_client .get_application_config() @@ -39,9 +41,7 @@ async fn main() -> Result<(), Box> { .await .context("failed to get endorsed evidence")?; - let encryption_key_handle = InstanceEncryptionKeyHandle::create() - .await - .context("couldn't create encryption key handle")?; + let encryption_key_handle = InstanceEncryptionKeyHandle::create(&orchestrator_channel); let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), ENCLAVE_APP_PORT); let listener = TcpListener::bind(addr).await?; diff --git a/oak_containers_sdk/BUILD b/oak_containers_sdk/BUILD index f0e4ed5b62..be083d6401 100644 --- a/oak_containers_sdk/BUILD +++ b/oak_containers_sdk/BUILD @@ -41,6 +41,7 @@ rust_library( "//oak_dice", "//oak_proto_rust", "//oak_proto_rust/grpc", + "//oak_sdk/containers:oak_sdk_containers", "//oak_session", "//stage0_dice", "@oak_crates_index//:anyhow", @@ -52,7 +53,6 @@ rust_library( "@oak_crates_index//:tokio", "@oak_crates_index//:tokio-stream", "@oak_crates_index//:tonic", - "@oak_crates_index//:tower", ], ) diff --git a/oak_containers_sdk/src/crypto.rs b/oak_containers_sdk/src/crypto.rs index bfb6bebf5c..fb2e4af57e 100644 --- a/oak_containers_sdk/src/crypto.rs +++ b/oak_containers_sdk/src/crypto.rs @@ -23,29 +23,14 @@ use oak_proto_rust::oak::{ containers::v1::{DeriveSessionKeysRequest, KeyOrigin, SignRequest}, crypto::v1::{SessionKeys, Signature}, }; -use tonic::transport::{Endpoint, Uri}; -use tower::service_fn; - -use crate::{IGNORED_ENDPOINT_URI, IPC_SOCKET}; struct OrchestratorCryptoClient { inner: GrpcOrchestratorCryptoClient, } impl OrchestratorCryptoClient { - async fn create() -> anyhow::Result { - let inner: GrpcOrchestratorCryptoClient = { - let channel = Endpoint::try_from(IGNORED_ENDPOINT_URI) - .context("couldn't form endpoint")? - .connect_with_connector(service_fn(move |_: Uri| { - tokio::net::UnixStream::connect(IPC_SOCKET) - })) - .await - .context("couldn't connect to UDS socket")?; - - GrpcOrchestratorCryptoClient::new(channel) - }; - Ok(Self { inner }) + fn create(channel: &tonic::transport::channel::Channel) -> Self { + Self { inner: GrpcOrchestratorCryptoClient::new(channel.clone()) } } async fn derive_session_keys( @@ -85,12 +70,8 @@ pub struct InstanceEncryptionKeyHandle { } impl InstanceEncryptionKeyHandle { - pub async fn create() -> anyhow::Result { - Ok(Self { - orchestrator_crypto_client: OrchestratorCryptoClient::create() - .await - .context("couldn't create Orchestrator crypto client")?, - }) + pub fn create(orchestrator_channel: &tonic::transport::Channel) -> Self { + Self { orchestrator_crypto_client: OrchestratorCryptoClient::create(orchestrator_channel) } } } @@ -129,14 +110,8 @@ pub struct InstanceSigner { } impl InstanceSigner { - pub async fn create() -> anyhow::Result { - Ok(Self { - orchestrator_crypto_client: Arc::new( - OrchestratorCryptoClient::create() - .await - .context("couldn't create Orchestrator crypto client")?, - ), - }) + pub fn create(channel: &tonic::transport::channel::Channel) -> Self { + Self { orchestrator_crypto_client: Arc::new(OrchestratorCryptoClient::create(channel)) } } } diff --git a/oak_containers_sdk/src/lib.rs b/oak_containers_sdk/src/lib.rs index 4a7b0ff9ca..d8b0006ec5 100644 --- a/oak_containers_sdk/src/lib.rs +++ b/oak_containers_sdk/src/lib.rs @@ -18,17 +18,6 @@ pub mod handler; pub mod orchestrator_client; pub mod standalone; -// 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 -// be supplied to create an `Endpoint`. Even though in this case the endpoint -// is technically a file, tonic expects us to provide our own connector, and -// this ignored endpoint. :( -static IGNORED_ENDPOINT_URI: &str = "file://[::]:0"; - -// Path used to facilitate inter-process communication between the orchestrator -// and the trusted application. -const 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 orchestrator_client::OrchestratorClient; diff --git a/oak_containers_sdk/src/orchestrator_client.rs b/oak_containers_sdk/src/orchestrator_client.rs index b6e3629a60..7b338d1cba 100644 --- a/oak_containers_sdk/src/orchestrator_client.rs +++ b/oak_containers_sdk/src/orchestrator_client.rs @@ -13,13 +13,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use anyhow::{Context, Result}; +use anyhow::Result; use oak_grpc::oak::containers::orchestrator_client::OrchestratorClient as GrpcOrchestratorClient; use oak_proto_rust::oak::session::v1::EndorsedEvidence; -use tonic::transport::{Endpoint, Uri}; -use tower::service_fn; - -use crate::{IGNORED_ENDPOINT_URI, IPC_SOCKET}; /// Utility struct used to interface with the Orchestrator. #[derive(Clone)] @@ -28,19 +24,8 @@ pub struct OrchestratorClient { } impl OrchestratorClient { - pub async fn create() -> Result { - let inner: GrpcOrchestratorClient = { - let channel = Endpoint::try_from(IGNORED_ENDPOINT_URI) - .context("couldn't form endpoint")? - .connect_with_connector(service_fn(move |_: Uri| { - tokio::net::UnixStream::connect(IPC_SOCKET) - })) - .await - .context("couldn't connect to UDS socket")?; - - GrpcOrchestratorClient::new(channel) - }; - Ok(Self { inner }) + pub fn create(channel: &tonic::transport::channel::Channel) -> Self { + Self { inner: GrpcOrchestratorClient::new(channel.clone()) } } /// Retrieves the application configuration from the Orchestrator. diff --git a/oak_functions_containers_app/BUILD b/oak_functions_containers_app/BUILD index 245d367fb2..5db2cdd7ad 100644 --- a/oak_functions_containers_app/BUILD +++ b/oak_functions_containers_app/BUILD @@ -115,6 +115,7 @@ rust_binary( deps = APP_DEPS + [ ":lib", "//oak_functions_service:lib_unrestricted", + "//oak_sdk/containers:oak_sdk_containers", ], ) @@ -128,6 +129,7 @@ rust_binary( deps = APP_DEPS + [ ":lib_insecure", "//oak_functions_service:lib_insecure_unrestricted", + "//oak_sdk/containers:oak_sdk_containers", ], ) diff --git a/oak_functions_containers_app/src/main.rs b/oak_functions_containers_app/src/main.rs index 417f3f41a3..29106d16a6 100644 --- a/oak_functions_containers_app/src/main.rs +++ b/oak_functions_containers_app/src/main.rs @@ -20,7 +20,7 @@ use std::{ num::{NonZeroU16, NonZeroU32}, }; -use anyhow::{anyhow, Context}; +use anyhow::Context; use clap::Parser; use oak_containers_agent::{ metrics::{MetricsConfig, OakObserver}, @@ -36,6 +36,7 @@ use oak_proto_rust::oak::functions::config::{ application_config::CommunicationChannel, ApplicationConfig, HandlerType, TcpCommunicationChannel, WasmtimeConfig, }; +use oak_sdk_containers::default_orchestrator_channel; use prost::Message; use tokio::{ io::{AsyncRead, AsyncWrite}, @@ -125,13 +126,12 @@ async fn main() -> Result<(), Box> { let oak_observer = oak_containers_agent::metrics::init_metrics(metrics_config); - let mut client = - OrchestratorClient::create().await.context("couldn't create Orchestrator client")?; - let encryption_key_handle = Box::new( - InstanceEncryptionKeyHandle::create() - .await - .map_err(|error| anyhow!("couldn't create encryption key handle: {:?}", error))?, - ); + let orchestrator_channel = + default_orchestrator_channel().await.context("failed to create channel to orchestrator")?; + + let mut client = OrchestratorClient::create(&orchestrator_channel); + let encryption_key_handle = + Box::new(InstanceEncryptionKeyHandle::create(&orchestrator_channel)); // To be used when connecting trusted app to orchestrator. let application_config = { diff --git a/oak_sdk/containers/BUILD b/oak_sdk/containers/BUILD new file mode 100644 index 0000000000..50c4789070 --- /dev/null +++ b/oak_sdk/containers/BUILD @@ -0,0 +1,41 @@ +# +# 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 = ["//:default_visibility"], + licenses = ["notice"], +) + +rust_library( + name = "oak_sdk_containers", + srcs = ["lib.rs"], + deps = [ + ":orchestrator_channel", + ], +) + +rust_library( + name = "orchestrator_channel", + srcs = ["orchestrator_channel.rs"], + deps = [ + "@oak_crates_index//:anyhow", + "@oak_crates_index//:tokio", + "@oak_crates_index//:tonic", + "@oak_crates_index//:tower", + ], +) diff --git a/oak_sdk/containers/lib.rs b/oak_sdk/containers/lib.rs new file mode 100644 index 0000000000..0d8704ef21 --- /dev/null +++ b/oak_sdk/containers/lib.rs @@ -0,0 +1,17 @@ +// +// Copyright 2025 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 use orchestrator_channel::default_orchestrator_channel; diff --git a/oak_sdk/containers/orchestrator_channel.rs b/oak_sdk/containers/orchestrator_channel.rs new file mode 100644 index 0000000000..b3fea1e1bc --- /dev/null +++ b/oak_sdk/containers/orchestrator_channel.rs @@ -0,0 +1,38 @@ +// +// Copyright 2025 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, Result}; +use tokio::net::UnixStream; +use tonic::transport::{Channel, Endpoint, Uri}; +use tower::service_fn; +// 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 +// be supplied to create an `Endpoint`. Even though in this case the endpoint +// is technically a file, tonic expects us to provide our own connector, and +// this ignored endpoint. :( +static IGNORED_ENDPOINT_URI: &str = "file://[::]:0"; + +// Path used to facilitate inter-process communication between the orchestrator +// and the trusted application. +const IPC_SOCKET: &str = "/oak_utils/orchestrator_ipc"; + +// Connect to the orchestrator instance in a default Oak Containers stack. +pub async fn default_orchestrator_channel() -> Result { + Endpoint::try_from(IGNORED_ENDPOINT_URI) + .context("couldn't form endpoint")? + .connect_with_connector(service_fn(move |_: Uri| UnixStream::connect(IPC_SOCKET))) + .await + .context("couldn't connect to UDS socket") +}