Split Oak Containers Orchestrator

This is part of the refactoring needed to support multiple Oak
Containers Orchestrator binaries for different attesters (e.g. DICE and
Intel TDX RTMRs). This splits the binary into a separate package from
the library that implements the orchestrator logic.

BUG: 380442628

Change-Id: I2b24e495b6923c6bd47530e6a2af075b81f8e248

justfile

@@ -482,7 +482,7 @@ containers_placer_artifacts:
 
     cp --force --preserve=timestamps bazel-bin/oak_containers/kernel/bzImage artifacts/oak_containers_kernel
     cp --force --preserve=timestamps bazel-bin/oak_containers/agent/bin/oak_containers_agent artifacts
-    cp --force --preserve=timestamps bazel-bin/oak_containers/orchestrator/bin/oak_containers_orchestrator artifacts
+    cp --force --preserve=timestamps bazel-bin/oak_containers/orchestrator_bin/bin/oak_containers_orchestrator artifacts
     cp --force --preserve=timestamps bazel-bin/oak_containers/syslogd/oak_containers_syslogd artifacts
 
     cp --force --preserve=timestamps --no-preserve=mode \
@@ -494,7 +494,7 @@ bazel_build_copy package target:
     cp --force --preserve=timestamps "bazel-bin/{{package}}/{{target}}" artifacts
 
 oak_containers_agent: (bazel_build_copy "oak_containers/agent" "bin/oak_containers_agent")
-oak_containers_orchestrator: (bazel_build_copy "oak_containers/orchestrator" "bin/oak_containers_orchestrator")
+oak_containers_orchestrator: (bazel_build_copy "oak_containers/orchestrator_bin" "bin/oak_containers_orchestrator")
 oak_containers_syslogd: (bazel_build_copy "oak_containers/syslogd" "oak_containers_syslogd")
 
 oak_containers_system_image:

oak_containers/orchestrator/BUILD

@@ -14,7 +14,7 @@
 # limitations under the License.
 #
 
-load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_library", "rust_test")
+load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test")
 
 package(
     default_visibility = ["//:default_visibility"],
@@ -38,6 +38,7 @@ rust_library(
         "//oak_attestation",
         "//oak_attestation_types",
         "//oak_attestation_verification",
+        "//oak_containers/agent",
         "//oak_containers/attestation:oak_containers_attestation",
         "//oak_containers/channel",
         "//oak_crypto",
@@ -47,6 +48,7 @@ rust_library(
         "@oak_crates_index//:anyhow",
         "@oak_crates_index//:async-stream",
         "@oak_crates_index//:ciborium",
+        "@oak_crates_index//:clap",
         "@oak_crates_index//:coset",
         "@oak_crates_index//:hpke",
         "@oak_crates_index//:log",
@@ -84,26 +86,6 @@ rust_test(
     ],
 )
 
-rust_binary(
-    name = "bin/oak_containers_orchestrator",
-    srcs = ["src/main.rs"],
-    deps = [
-        ":oak_containers_orchestrator",
-        "//oak_attestation_types",
-        "//oak_containers/agent",
-        "//oak_containers/attestation:oak_containers_attestation",
-        "//oak_proto_rust",
-        "//oak_proto_rust/grpc",
-        "@oak_crates_index//:anyhow",
-        "@oak_crates_index//:clap",
-        "@oak_crates_index//:nix",
-        "@oak_crates_index//:prost",
-        "@oak_crates_index//:tikv-jemallocator",
-        "@oak_crates_index//:tokio",
-        "@oak_crates_index//:tokio-util",
-    ],
-)
-
 alias(
     name = "orchestrator",
     actual = ":oak_containers_orchestrator",

oak_containers/orchestrator/README.md

@@ -6,6 +6,6 @@
 
 # Orchestrator
 
-Implementation of the orchestrator, a binary responsible for loading a container
-and running it in a container runtime. It also exposes configuration and remote
-attestation logic to the instantiated container.
+Implementation of the orchestrator which is responsible for loading a container
+image and running it in a container runtime. It also exposes configuration and
+remote attestation logic to the instantiated container.

oak_containers/orchestrator/src/lib.rs

@@ -13,10 +13,178 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::{path::PathBuf, sync::Arc};
+
+use anyhow::{anyhow, Context};
+use clap::Parser;
+use launcher_client::LauncherClient;
+use oak_attestation_types::attester::Attester;
+use oak_containers_agent::{metrics::MetricsConfig, set_error_handler};
+use oak_containers_attestation::generate_instance_keys;
+use oak_proto_rust::oak::containers::v1::KeyProvisioningRole;
+use prost::Message;
+use tokio_util::sync::CancellationToken;
+
 mod cdi;
 pub mod container_runtime;
 pub mod dice;
 pub mod ipc_server;
 pub mod key_provisioning;
 pub mod launcher_client;
 pub mod logging;
+
+#[derive(Parser, Debug)]
+struct Args {
+    #[arg(env, default_value = "http://10.0.2.100:8080")]
+    launcher_addr: String,
+
+    #[arg(default_value = "10.0.2.15:4000")]
+    orchestrator_addr: String,
+
+    #[arg(long, default_value = "/oak_container")]
+    container_dir: PathBuf,
+
+    #[arg(long, default_value = "/oak_utils/orchestrator_ipc")]
+    ipc_socket_path: PathBuf,
+
+    #[arg(long, default_value = "oakc")]
+    runtime_user: String,
+}
+
+pub async fn main() -> anyhow::Result<()> {
+    crate::logging::setup()?;
+
+    let args = Args::parse();
+
+    let launcher_client = Arc::new(
+        LauncherClient::create(args.launcher_addr.parse()?)
+            .await
+            .map_err(|error| anyhow!("couldn't create client: {:?}", error))?,
+    );
+
+    set_error_handler(|err| eprintln!("oak_containers_orchestrator: OTLP error: {}", err))?;
+
+    let metrics_config = MetricsConfig {
+        launcher_addr: args.launcher_addr,
+        scope: String::from("orchestrator"),
+        excluded_metrics: None,
+    };
+
+    let _oak_observer = oak_containers_agent::metrics::init_metrics(metrics_config);
+
+    // Get key provisioning role.
+    let key_provisioning_role = launcher_client
+        .get_key_provisioning_role()
+        .await
+        .map_err(|error| anyhow!("couldn't get key provisioning role: {:?}", error))?;
+
+    // Generate application keys.
+    let (instance_keys, instance_public_keys) = generate_instance_keys();
+    let (mut group_keys, group_public_keys) =
+        if key_provisioning_role == KeyProvisioningRole::Leader {
+            let (group_keys, group_public_keys) = instance_keys.generate_group_keys();
+            (Some(Arc::new(group_keys)), Some(group_public_keys))
+        } else {
+            (None, None)
+        };
+
+    // Load application.
+    let container_bundle = launcher_client
+        .get_container_bundle()
+        .await
+        .map_err(|error| anyhow!("couldn't get container bundle: {:?}", error))?;
+    let application_config = launcher_client
+        .get_application_config()
+        .await
+        .map_err(|error| anyhow!("couldn't get application config: {:?}", error))?;
+
+    // Create a container event and add it to the event log.
+    let mut attester = crate::dice::load_stage1_dice_data()?;
+    let container_event = oak_containers_attestation::create_container_event(
+        &container_bundle,
+        &application_config,
+        &instance_public_keys,
+    );
+    attester
+        .extend(&container_event.encode_to_vec())
+        .context("couldn't add container event to the evidence")?;
+
+    // Add the container event to the DICE chain.
+    let container_layer = oak_containers_attestation::create_container_dice_layer(&container_event);
+    let evidence = attester.add_application_keys(
+        container_layer,
+        &instance_public_keys.encryption_public_key,
+        &instance_public_keys.signing_public_key,
+        if let Some(ref group_public_keys) = group_public_keys {
+            Some(&group_public_keys.encryption_public_key)
+        } else {
+            None
+        },
+        None,
+    )?;
+
+    // Send the attestation evidence to the Hostlib.
+    launcher_client
+        .send_attestation_evidence(evidence.clone())
+        .await
+        .map_err(|error| anyhow!("couldn't send attestation evidence: {:?}", error))?;
+
+    // Request group keys.
+    if key_provisioning_role == KeyProvisioningRole::Follower {
+        let get_group_keys_response = launcher_client
+            .get_group_keys()
+            .await
+            .map_err(|error| anyhow!("couldn't get group keys: {:?}", error))?;
+        let provisioned_group_keys = instance_keys
+            .provide_group_keys(get_group_keys_response)
+            .context("couldn't provide group keys")?;
+        group_keys = Some(Arc::new(provisioned_group_keys));
+    }
+
+    if let Some(path) = args.ipc_socket_path.parent() {
+        tokio::fs::create_dir_all(path).await?;
+    }
+
+    let endorsements = launcher_client
+        .get_endorsements()
+        .await
+        .map_err(|e| anyhow!("coudln't get endorsements from launcher: {e:?}"))?;
+
+    let (orchestrator_server, crypto_server) = crate::ipc_server::create_services(
+        evidence,
+        endorsements,
+        instance_keys,
+        group_keys.clone().context("group keys were not provisioned")?,
+        application_config,
+        launcher_client,
+    );
+
+    // Start application and gRPC servers.
+    let user = nix::unistd::User::from_name(&args.runtime_user)
+        .context(format!("error resolving user {}", args.runtime_user))?
+        .context(format!("user `{}` not found", args.runtime_user))?;
+    let cancellation_token = CancellationToken::new();
+    tokio::try_join!(
+        crate::ipc_server::server(
+            &args.ipc_socket_path,
+            orchestrator_server,
+            crypto_server,
+            cancellation_token.clone(),
+        ),
+        crate::key_provisioning::create(
+            &args.orchestrator_addr,
+            group_keys.context("group keys were not provisioned")?,
+            cancellation_token.clone(),
+        ),
+        crate::container_runtime::run(
+            &container_bundle,
+            &args.container_dir,
+            user.uid,
+            user.gid,
+            &args.ipc_socket_path,
+            cancellation_token,
+        ),
+    )?;
+
+    Ok(())
+}