Make stage1 pass the launcher addr to orchestrator, and support vsock…

… there We can't just pass random flags around, so this attempts to pass the launcher address around via an environment variable. This requires more testing, though. Change-Id: I14e35be5024a29a18d90e4de67133db84fdba392
project-oak · Jun 25, 2024 · 960d461 · 960d461
1 parent fa2060d
commit 960d461
Show file tree

Hide file tree

Showing 11 changed files with 57 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bazel/crates/oak_crates_index.bzl b/bazel/crates/oak_crates_index.bzl
@@ -103,7 +103,10 @@ def oak_crates_index(cargo_lockfile, lockfile):
                 version = "*",
             ),
             "clap": crate.spec(
-                features = ["derive"],
+                features = [
+                    "derive",
+                    "env",
+                ],
                 version = "*",
             ),
             "command-fds": crate.spec(

diff --git a/bazel/test_workspace/cargo-bazel-test-lock.json b/bazel/test_workspace/cargo-bazel-test-lock.json
@@ -1,5 +1,5 @@
 {
-  "checksum": "b21bd38f017e09deb79be48415d35b2756078e174df6cc5ed486b0b033ff5c5b",
+  "checksum": "2a035addb746c8ca234cb3909b9c267ae8c1865e72c084934f37c1903d734e78",
   "crates": {
     "acpi 5.0.0": {
       "name": "acpi",
@@ -2787,6 +2787,7 @@
             "color",
             "default",
             "derive",
+            "env",
             "error-context",
             "help",
             "std",
@@ -2855,6 +2856,7 @@
         "crate_features": {
           "common": [
             "color",
+            "env",
             "error-context",
             "help",
             "std",

diff --git a/cargo-bazel-lock.json b/cargo-bazel-lock.json
@@ -1,5 +1,5 @@
 {
-  "checksum": "6f60f483c9b3770febf3b2110b3571bfd1a94c5e21aca9cb8b0f1f09ebf05b5d",
+  "checksum": "682754f838a52af6677cda26d9abff71b5539ae2534786da75c150bbb00d9826",
   "crates": {
     "acpi 5.0.0": {
       "name": "acpi",
@@ -2770,6 +2770,7 @@
             "color",
             "default",
             "derive",
+            "env",
             "error-context",
             "help",
             "std",
@@ -2838,6 +2839,7 @@
         "crate_features": {
           "common": [
             "color",
+            "env",
             "error-context",
             "help",
             "std",

diff --git a/oak_containers_orchestrator/BUILD b/oak_containers_orchestrator/BUILD
@@ -65,7 +65,9 @@ rust_library(
         "@oak_crates_index//:tokio",
         "@oak_crates_index//:tokio-stream",
         "@oak_crates_index//:tokio-util",
+        "@oak_crates_index//:tokio-vsock",
         "@oak_crates_index//:tonic",
+        "@oak_crates_index//:tower",
         "@oak_crates_index//:walkdir",
         "@oak_crates_index//:zeroize",
     ],

diff --git a/oak_containers_orchestrator/Cargo.toml b/oak_containers_orchestrator/Cargo.toml
@@ -8,7 +8,7 @@ license = "Apache-2.0"
 [dependencies]
 anyhow = "*"
 ciborium = { version = "*", default-features = false }
-clap = { version = "*", features = ["derive"] }
+clap = { version = "*", features = ["derive", "env"] }
 coset = { version = "*", features = ["std"] }
 hpke = { version = "*", default-features = false, features = [
   "alloc",
@@ -53,7 +53,9 @@ tokio = { version = "*", features = [
 ] }
 tokio-stream = { version = "*", features = ["net"] }
 tokio-util = { version = "*", default-features = false }
+tokio-vsock = { version = "*", features = ["tonic-conn"] }
 tonic = { workspace = true }
+tower = "*"
 walkdir = "*"
 zeroize = "*"
 

diff --git a/oak_containers_orchestrator/src/launcher_client.rs b/oak_containers_orchestrator/src/launcher_client.rs
@@ -15,7 +15,9 @@
 
 use anyhow::Context;
 use opentelemetry_otlp::{TonicExporterBuilder, WithExportConfig};
+use tokio_vsock::{VsockAddr, VsockStream};
 use tonic::transport::Channel;
+use tower::service_fn;
 
 use crate::proto::oak::{
     attestation::v1::Evidence,
@@ -36,7 +38,25 @@ pub struct LauncherClient {
 
 impl LauncherClient {
     pub async fn create(addr: tonic::transport::Uri) -> Result<Self, Box<dyn std::error::Error>> {
-        let channel = Channel::builder(addr.clone()).connect().await?;
+        let channel = if addr.scheme_str() == Some("vsock") {
+            let vsock_addr = VsockAddr::new(
+                addr.host()
+                    .unwrap_or(format!("{}", tokio_vsock::VMADDR_CID_HOST).as_str())
+                    .parse()
+                    .context("invalid vsock CID")?,
+                addr.port_u16().context("invalid vsock port")?.into(),
+            );
+
+            // The C++ gRPC implementations are more particular about the URI scheme; in
+            // particular, they may get confused by the "vsock" scheme. Therfore, create a
+            // fake URI with the "http" scheme to placate the libraries; the actual
+            // connection is made in `connect_with_connector` and that uses the correct URI.
+            Channel::builder(tonic::transport::Uri::from_static("http://0:0"))
+                .connect_with_connector(service_fn(move |_| VsockStream::connect(vsock_addr)))
+                .await?
+        } else {
+            Channel::builder(addr.clone()).connect().await?
+        };
         let inner = GrpcLauncherClient::new(channel.clone());
         let hostlib_key_provisioning_client = HostlibKeyProvisioningClient::new(channel);
         Ok(Self { addr, inner, hostlib_key_provisioning_client })

diff --git a/oak_containers_orchestrator/src/main.rs b/oak_containers_orchestrator/src/main.rs
@@ -28,7 +28,7 @@ static ALLOCATOR: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
 #[derive(Parser, Debug)]
 struct Args {
-    #[arg(default_value = "http://10.0.2.100:8080")]
+    #[arg(env, default_value = "http://10.0.2.100:8080")]
     launcher_addr: String,
 
     #[arg(default_value = "10.0.2.15:4000")]

diff --git a/oak_containers_stage1/src/image.rs b/oak_containers_stage1/src/image.rs
@@ -14,10 +14,14 @@
 // limitations under the License.
 //
 
-use std::{ffi::CString, os::unix::prelude::OsStrExt, path::Path};
+use std::{
+    ffi::{CStr, CString},
+    os::unix::prelude::OsStrExt,
+    path::Path,
+};
 
 use anyhow::{anyhow, Result};
-use nix::unistd::execv;
+use nix::unistd::execve;
 use tar::Archive;
 use xz2::read::XzDecoder;
 
@@ -27,13 +31,16 @@ pub async fn extract(buf: &[u8], dst: &Path) -> Result<()> {
     archive.unpack(dst).map_err(|e| anyhow!(e))
 }
 
-pub fn switch(init: &str) -> Result<!> {
+pub fn switch<SE>(init: &str, env: &[SE]) -> Result<!>
+where
+    SE: AsRef<CStr>,
+{
     // On one hand, I feel like this function should be marked `unsafe` as this will
     // unconditionally switch over to the new executable (if it succeeds) without
     // any more Rust code executing. On the other hand, the return type is `!`,
     // so you shouldn't expect the control to return.
     let args: Vec<CString> =
         std::env::args_os().map(|arg| CString::new(arg.as_bytes()).unwrap()).collect();
-    execv(CString::new(init).unwrap().as_c_str(), &args[..])?;
+    execve(CString::new(init).unwrap().as_c_str(), &args[..], env)?;
     unreachable!()
 }
diff --git a/oak_containers_stage1/src/main.rs b/oak_containers_stage1/src/main.rs
@@ -92,7 +92,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     chroot(".").context("failed to chroot to .")?;
     chdir("/").context("failed to chdir to /")?;
 
-    let mut client = LauncherClient::new(args.launcher_addr)
+    let mut client = LauncherClient::new(args.launcher_addr.clone())
         .await
         .context("error creating the launcher client")?;
 
@@ -150,7 +150,11 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .or_else(|err| if err.kind() == ErrorKind::NotFound { Ok(()) } else { Err(err) })
         .context("error removing `/.dockerenv`")?;
 
-    image::switch(&args.init).context("error switching to the system image")?
+    image::switch(
+        &args.init,
+        &[std::ffi::CString::new(format!("LAUNCHER_ADDR={}", args.launcher_addr)).unwrap()],
+    )
+    .context("error switching to the system image")?
 }
 
 /// Tries to parse a string slice as an address.

diff --git a/oak_containers_system_image/files/etc/systemd/system/oak-orchestrator.service b/oak_containers_system_image/files/etc/systemd/system/oak-orchestrator.service
@@ -12,6 +12,7 @@ ExecStopPost=/usr/bin/journalctl --sync --flush
 ProtectSystem=strict
 RuntimeDirectory=oak_containers_orchestrator
 ReadWritePaths=/oak
+PassEnvironment=LAUNCHER_ADDR
 
 [Install]
 WantedBy=multi-user.target