Change the extraction functions in orhestrator/stage1 to take Buf

This make it more generic and we can pass in non-contiguous buffers in
the future by going via that trait.

Bug: 396664122
Change-Id: If29db8aa6edeaaace33d614044bb9f1c064f9546

oak_containers/attestation/BUILD

@@ -32,6 +32,7 @@ rust_library(
         "//oak_dice",
         "//oak_proto_rust",
         "@oak_crates_index//:anyhow",
+        "@oak_crates_index//:bytes",
         "@oak_crates_index//:ciborium",
         "@oak_crates_index//:coset",
         "@oak_crates_index//:p256",

oak_containers/attestation/src/lib.rs

@@ -18,6 +18,7 @@ extern crate alloc;
 use alloc::{string::ToString, vec};
 
 use anyhow::Context;
+use bytes::Buf;
 use ciborium::Value;
 use coset::cwt::ClaimName;
 use oak_crypto::{
@@ -33,7 +34,7 @@ use oak_proto_rust::oak::{
 use prost::Message;
 
 /// Measures the system image and returns a corresponding event log entry.
-pub fn create_system_layer_event(system_image: &[u8]) -> Event {
+pub fn create_system_layer_event<B: Buf>(system_image: B) -> Event {
     let digest = oak_attestation::MeasureDigest::measure_digest(system_image);
     Event {
         tag: "stage1".to_string(),
@@ -46,9 +47,9 @@ pub fn create_system_layer_event(system_image: &[u8]) -> Event {
 
 /// Creates a container event that includes image bytes and configuration
 /// measurements and public keys used by the container.
-pub fn create_container_event(
-    container_bytes: &[u8],
-    config_bytes: &[u8],
+pub fn create_container_event<A: Buf, B: Buf>(
+    container_bytes: A,
+    config_bytes: B,
     instance_public_keys: &InstancePublicKeys,
 ) -> Event {
     let container_digest = oak_attestation::MeasureDigest::measure_digest(container_bytes);

oak_containers/orchestrator/BUILD

@@ -51,6 +51,7 @@ rust_library(
         "//oak_proto_rust/grpc",
         "@oak_crates_index//:anyhow",
         "@oak_crates_index//:async-stream",
+        "@oak_crates_index//:bytes",
         "@oak_crates_index//:ciborium",
         "@oak_crates_index//:clap",
         "@oak_crates_index//:coset",

oak_containers/orchestrator/src/container_runtime.rs

@@ -20,12 +20,13 @@ use std::{
 };
 
 use anyhow::Context;
+use bytes::Buf;
 use nix::unistd::{Gid, Uid};
 use oci_spec::runtime::{LinuxIdMapping, LinuxIdMappingBuilder, Mount, Spec};
 use tokio_util::sync::CancellationToken;
 
-pub async fn run(
-    container_bundle: &[u8],
+pub async fn run<B: Buf>(
+    container_bundle: B,
     container_dir: &Path,
     runtime_uid: Uid,
     runtime_gid: Gid,
@@ -34,7 +35,7 @@ pub async fn run(
 ) -> Result<(), anyhow::Error> {
     tokio::fs::create_dir_all(container_dir).await?;
     log::info!("Unpacking container bundle");
-    let mut archive = tar::Archive::new(container_bundle);
+    let mut archive = tar::Archive::new(container_bundle.reader());
     archive.unpack(container_dir)?;
 
     for entry in walkdir::WalkDir::new(container_dir) {

oak_containers/orchestrator/src/lib.rs

@@ -104,8 +104,8 @@ pub async fn main<A: Attester + ApplicationKeysAttester + Serializable>() -> any
     // Create a container event and add it to the event log.
     let mut attester: A = crate::dice::load_stage1_dice_data()?;
     let container_event = oak_containers_attestation::create_container_event(
-        &container_bundle,
-        &application_config,
+        &container_bundle[..],
+        &application_config[..],
         &instance_public_keys,
     );
     attester
@@ -188,7 +188,7 @@ pub async fn main<A: Attester + ApplicationKeysAttester + Serializable>() -> any
             cancellation_token.clone(),
         ),
         crate::container_runtime::run(
-            &container_bundle,
+            &container_bundle[..],
             &args.container_dir,
             user.uid,
             user.gid,

oak_containers/stage1/BUILD

@@ -33,6 +33,7 @@ rust_library(
         "//oak_proto_rust",
         "//oak_proto_rust/grpc",
         "@oak_crates_index//:anyhow",
+        "@oak_crates_index//:bytes",
         "@oak_crates_index//:clap",
         "@oak_crates_index//:futures-util",
         "@oak_crates_index//:nix",

oak_containers/stage1/src/image.rs

@@ -21,12 +21,13 @@ use std::{
 };
 
 use anyhow::{anyhow, Result};
+use bytes::Buf;
 use nix::unistd::execve;
 use tar::Archive;
 use xz2::read::XzDecoder;
 
-pub async fn extract(buf: &[u8], dst: &Path) -> Result<()> {
-    let decoder = XzDecoder::new(buf);
+pub async fn extract<B: Buf>(buf: B, dst: &Path) -> Result<()> {
+    let decoder = XzDecoder::new(buf.reader());
     let mut archive = Archive::new(decoder);
     archive.unpack(dst).map_err(|e| anyhow!(e))
 }

oak_containers/stage1/src/lib.rs

@@ -119,11 +119,11 @@ pub async fn main<A: Attester + Serializable>(args: &Args) -> Result<(), Box<dyn
 
     // For safety we generate the DICE data for the next layer before processing the
     // compressed system image.
-    let system_image_event = oak_containers_attestation::create_system_layer_event(&buf);
+    let system_image_event = oak_containers_attestation::create_system_layer_event(&buf[..]);
     attester.extend(&system_image_event.encode_to_vec())?;
     let dice_data = attester.serialize();
 
-    image::extract(&buf, Path::new("/")).await.context("error loading the system image")?;
+    image::extract(&buf[..], Path::new("/")).await.context("error loading the system image")?;
 
     // If the image didn't contain a `/etc/machine-id` file, create a placeholder
     // one that systemd will replace during startup. If you don't have that file

oak_sdk/standalone/standalone.rs

@@ -192,7 +192,7 @@ impl<'a> StandaloneBuilder<'a> {
         // Add container event and add it to the event log.
         let container_event = oak_containers_attestation::create_container_event(
             application_image,
-            &application_config,
+            &application_config[..],
             &instance_public_keys,
         );
         attester