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

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",

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<dyn std::error::Error>> {
     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<dyn std::error::Error>> {
         .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?;

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",
     ],
 )
 

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<tonic::transport::channel::Channel>,
 }
 
 impl OrchestratorCryptoClient {
-    async fn create() -> anyhow::Result<Self> {
-        let inner: GrpcOrchestratorCryptoClient<tonic::transport::channel::Channel> = {
-            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<Self> {
-        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<Self> {
-        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)) }
     }
 }
 

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;

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<Self> {
-        let inner: GrpcOrchestratorClient<tonic::transport::channel::Channel> = {
-            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.

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",
     ],
 )
 

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<dyn std::error::Error>> {
 
     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 = {

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",
+    ],
+)

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;

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<Channel> {
+    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")
+}