feat: oci compliance work (#85)

* chore: rework oci crate to be more composable

* feat: image pull is now internally explicit

* feat: utilize vfs for assembling oci images

* feat: rework oci to preserve permissions via a vfs

crates/ctl/src/cli/destroy.rs

@@ -52,31 +52,30 @@ impl DestroyCommand {
 async fn wait_guest_destroyed(id: &str, events: EventStream) -> Result<()> {
     let mut stream = events.subscribe();
     while let Ok(event) = stream.recv().await {
-        if let Event::GuestChanged(changed) = event {
-            let Some(guest) = changed.guest else {
-                continue;
-            };
+        let Event::GuestChanged(changed) = event;
+        let Some(guest) = changed.guest else {
+            continue;
+        };
 
-            if guest.id != id {
-                continue;
+        if guest.id != id {
+            continue;
+        }
+
+        let Some(state) = guest.state else {
+            continue;
+        };
+
+        if let Some(ref error) = state.error_info {
+            if state.status() == GuestStatus::Failed {
+                error!("destroy failed: {}", error.message);
+                std::process::exit(1);
+            } else {
+                error!("guest error: {}", error.message);
             }
+        }
 
-            let Some(state) = guest.state else {
-                continue;
-            };
-
-            if let Some(ref error) = state.error_info {
-                if state.status() == GuestStatus::Failed {
-                    error!("destroy failed: {}", error.message);
-                    std::process::exit(1);
-                } else {
-                    error!("guest error: {}", error.message);
-                }
-            }
-
-            if state.status() == GuestStatus::Destroyed {
-                std::process::exit(0);
-            }
+        if state.status() == GuestStatus::Destroyed {
+            std::process::exit(0);
         }
     }
     Ok(())

crates/ctl/src/cli/launch.rs

@@ -2,17 +2,16 @@ use std::collections::HashMap;
 
 use anyhow::Result;
 use clap::Parser;
-use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
 use krata::{
     events::EventStream,
     v1::{
         common::{
-            guest_image_spec::Image, GuestImageSpec, GuestOciImageSpec, GuestSpec, GuestStatus,
-            GuestTaskSpec, GuestTaskSpecEnvVar,
+            guest_image_spec::Image, GuestImageSpec, GuestOciImageFormat, GuestOciImageSpec,
+            GuestSpec, GuestStatus, GuestTaskSpec, GuestTaskSpecEnvVar,
         },
         control::{
             control_service_client::ControlServiceClient, watch_events_reply::Event,
-            CreateGuestRequest, OciProgressEventLayerPhase, OciProgressEventPhase,
+            CreateGuestRequest, PullImageRequest,
         },
     },
 };
@@ -20,11 +19,15 @@ use log::error;
 use tokio::select;
 use tonic::{transport::Channel, Request};
 
-use crate::console::StdioConsoleStream;
+use crate::{console::StdioConsoleStream, pull::pull_interactive_progress};
+
+use super::pull::PullImageFormat;
 
 #[derive(Parser)]
 #[command(about = "Launch a new guest")]
 pub struct LauchCommand {
+    #[arg(short = 'S', long, default_value = "squashfs", help = "Image format")]
+    image_format: PullImageFormat,
     #[arg(short, long, help = "Name of the guest")]
     name: Option<String>,
     #[arg(
@@ -71,11 +74,25 @@ impl LauchCommand {
         mut client: ControlServiceClient<Channel>,
         events: EventStream,
     ) -> Result<()> {
+        let response = client
+            .pull_image(PullImageRequest {
+                image: self.oci.clone(),
+                format: match self.image_format {
+                    PullImageFormat::Squashfs => GuestOciImageFormat::Squashfs.into(),
+                    PullImageFormat::Erofs => GuestOciImageFormat::Erofs.into(),
+                },
+            })
+            .await?;
+        let reply = pull_interactive_progress(response.into_inner()).await?;
+
         let request = CreateGuestRequest {
             spec: Some(GuestSpec {
                 name: self.name.unwrap_or_default(),
                 image: Some(GuestImageSpec {
-                    image: Some(Image::Oci(GuestOciImageSpec { image: self.oci })),
+                    image: Some(Image::Oci(GuestOciImageSpec {
+                        digest: reply.digest,
+                        format: reply.format,
+                    })),
                 }),
                 vcpus: self.cpus,
                 mem: self.mem,
@@ -126,14 +143,9 @@ impl LauchCommand {
 
 async fn wait_guest_started(id: &str, events: EventStream) -> Result<()> {
     let mut stream = events.subscribe();
-    let mut multi_progress: Option<(MultiProgress, HashMap<String, ProgressBar>)> = None;
     while let Ok(event) = stream.recv().await {
         match event {
             Event::GuestChanged(changed) => {
-                if let Some((multi_progress, _)) = multi_progress.as_mut() {
-                    let _ = multi_progress.clear();
-                }
-
                 let Some(guest) = changed.guest else {
                     continue;
                 };
@@ -164,102 +176,6 @@ async fn wait_guest_started(id: &str, events: EventStream) -> Result<()> {
                     break;
                 }
             }
-
-            Event::OciProgress(oci) => {
-                if multi_progress.is_none() {
-                    multi_progress = Some((MultiProgress::new(), HashMap::new()));
-                }
-
-                let Some((multi_progress, progresses)) = multi_progress.as_mut() else {
-                    continue;
-                };
-
-                match oci.phase() {
-                    OciProgressEventPhase::Resolved
-                    | OciProgressEventPhase::ConfigAcquire
-                    | OciProgressEventPhase::LayerAcquire => {
-                        if progresses.is_empty() && !oci.layers.is_empty() {
-                            for layer in &oci.layers {
-                                let bar = ProgressBar::new(layer.total);
-                                bar.set_style(
-                                    ProgressStyle::with_template("{msg} {wide_bar}").unwrap(),
-                                );
-                                progresses.insert(layer.id.clone(), bar.clone());
-                                multi_progress.add(bar);
-                            }
-                        }
-
-                        for layer in oci.layers {
-                            let Some(progress) = progresses.get_mut(&layer.id) else {
-                                continue;
-                            };
-
-                            let phase = match layer.phase() {
-                                OciProgressEventLayerPhase::Waiting => "waiting",
-                                OciProgressEventLayerPhase::Downloading => "downloading",
-                                OciProgressEventLayerPhase::Downloaded => "downloaded",
-                                OciProgressEventLayerPhase::Extracting => "extracting",
-                                OciProgressEventLayerPhase::Extracted => "extracted",
-                                _ => "unknown",
-                            };
-
-                            let simple = if let Some((_, hash)) = layer.id.split_once(':') {
-                                hash
-                            } else {
-                                id
-                            };
-                            let simple = if simple.len() > 10 {
-                                &simple[0..10]
-                            } else {
-                                simple
-                            };
-                            let message = format!("{:width$} {}", simple, phase, width = 10);
-
-                            if message != progress.message() {
-                                progress.set_message(message);
-                            }
-
-                            progress.update(|state| {
-                                state.set_len(layer.total);
-                                state.set_pos(layer.value);
-                            });
-                        }
-                    }
-
-                    OciProgressEventPhase::Packing => {
-                        for (key, bar) in &mut *progresses {
-                            if key == "packing" {
-                                continue;
-                            }
-                            bar.finish_and_clear();
-                            multi_progress.remove(bar);
-                        }
-                        progresses.retain(|k, _| k == "packing");
-                        if progresses.is_empty() {
-                            let progress = ProgressBar::new(100);
-                            progress.set_message("packing");
-                            progress.set_style(
-                                ProgressStyle::with_template("{msg} {wide_bar}").unwrap(),
-                            );
-                            progresses.insert("packing".to_string(), progress);
-                        }
-                        let Some(progress) = progresses.get("packing") else {
-                            continue;
-                        };
-
-                        progress.update(|state| {
-                            state.set_len(oci.total);
-                            state.set_pos(oci.value);
-                        });
-                    }
-
-                    _ => {}
-                }
-
-                for progress in progresses {
-                    progress.1.tick();
-                }
-            }
         }
     }
     Ok(())

crates/ctl/src/cli/mod.rs

@@ -5,6 +5,7 @@ pub mod launch;
 pub mod list;
 pub mod logs;
 pub mod metrics;
+pub mod pull;
 pub mod resolve;
 pub mod top;
 pub mod watch;
@@ -21,7 +22,7 @@ use tonic::{transport::Channel, Request};
 use self::{
     attach::AttachCommand, destroy::DestroyCommand, idm_snoop::IdmSnoopCommand,
     launch::LauchCommand, list::ListCommand, logs::LogsCommand, metrics::MetricsCommand,
-    resolve::ResolveCommand, top::TopCommand, watch::WatchCommand,
+    pull::PullCommand, resolve::ResolveCommand, top::TopCommand, watch::WatchCommand,
 };
 
 #[derive(Parser)]
@@ -48,6 +49,7 @@ pub enum Commands {
     Destroy(DestroyCommand),
     List(ListCommand),
     Attach(AttachCommand),
+    Pull(PullCommand),
     Logs(LogsCommand),
     Watch(WatchCommand),
     Resolve(ResolveCommand),
@@ -101,6 +103,10 @@ impl ControlCommand {
             Commands::Top(top) => {
                 top.run(client, events).await?;
             }
+
+            Commands::Pull(pull) => {
+                pull.run(client).await?;
+            }
         }
         Ok(())
     }

crates/ctl/src/cli/pull.rs

@@ -0,0 +1,42 @@
+use anyhow::Result;
+use clap::{Parser, ValueEnum};
+use krata::v1::{
+    common::GuestOciImageFormat,
+    control::{control_service_client::ControlServiceClient, PullImageRequest},
+};
+
+use tonic::transport::Channel;
+
+use crate::pull::pull_interactive_progress;
+
+#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
+pub enum PullImageFormat {
+    Squashfs,
+    Erofs,
+}
+
+#[derive(Parser)]
+#[command(about = "Pull an image into the cache")]
+pub struct PullCommand {
+    #[arg(help = "Image name")]
+    image: String,
+    #[arg(short = 's', long, default_value = "squashfs", help = "Image format")]
+    image_format: PullImageFormat,
+}
+
+impl PullCommand {
+    pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
+        let response = client
+            .pull_image(PullImageRequest {
+                image: self.image.clone(),
+                format: match self.image_format {
+                    PullImageFormat::Squashfs => GuestOciImageFormat::Squashfs.into(),
+                    PullImageFormat::Erofs => GuestOciImageFormat::Erofs.into(),
+                },
+            })
+            .await?;
+        let reply = pull_interactive_progress(response.into_inner()).await?;
+        println!("{}", reply.digest);
+        Ok(())
+    }
+}

crates/ctl/src/cli/watch.rs

@@ -29,10 +29,9 @@ impl WatchCommand {
         loop {
             let event = stream.recv().await?;
 
-            if let Event::GuestChanged(changed) = event {
-                let guest = changed.guest.clone();
-                self.print_event("guest.changed", changed, guest)?;
-            }
+            let Event::GuestChanged(changed) = event;
+            let guest = changed.guest.clone();
+            self.print_event("guest.changed", changed, guest)?;
         }
     }
 

crates/ctl/src/console.rs

@@ -69,27 +69,26 @@ impl StdioConsoleStream {
         Ok(tokio::task::spawn(async move {
             let mut stream = events.subscribe();
             while let Ok(event) = stream.recv().await {
-                if let Event::GuestChanged(changed) = event {
-                    let Some(guest) = changed.guest else {
-                        continue;
-                    };
+                let Event::GuestChanged(changed) = event;
+                let Some(guest) = changed.guest else {
+                    continue;
+                };
 
-                    let Some(state) = guest.state else {
-                        continue;
-                    };
+                let Some(state) = guest.state else {
+                    continue;
+                };
 
-                    if guest.id != id {
-                        continue;
-                    }
+                if guest.id != id {
+                    continue;
+                }
 
-                    if let Some(exit_info) = state.exit_info {
-                        return Some(exit_info.code);
-                    }
+                if let Some(exit_info) = state.exit_info {
+                    return Some(exit_info.code);
+                }
 
-                    let status = state.status();
-                    if status == GuestStatus::Destroying || status == GuestStatus::Destroyed {
-                        return Some(10);
-                    }
+                let status = state.status();
+                if status == GuestStatus::Destroying || status == GuestStatus::Destroyed {
+                    return Some(10);
                 }
             }
             None

crates/ctl/src/lib.rs

@@ -2,3 +2,4 @@ pub mod cli;
 pub mod console;
 pub mod format;
 pub mod metrics;
+pub mod pull;

crates/ctl/src/metrics.rs

@@ -82,7 +82,7 @@ impl MultiMetricCollector {
             let collect = select! {
                 x = events.recv() => match x {
                     Ok(event) => {
-                        if let Event::GuestChanged(changed) = event {
+                        let Event::GuestChanged(changed) = event;
                             let Some(guest) = changed.guest else {
                                 continue;
                             };
@@ -93,7 +93,6 @@ impl MultiMetricCollector {
                             if state.status() != GuestStatus::Destroying {
                                 guests.push(guest);
                             }
-                        }
                         false
                     },
 

crates/ctl/src/pull.rs

@@ -0,0 +1,118 @@
+use std::collections::HashMap;
+
+use anyhow::{anyhow, Result};
+use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
+use krata::v1::control::{PullImageProgressLayerPhase, PullImageProgressPhase, PullImageReply};
+use tokio_stream::StreamExt;
+use tonic::Streaming;
+
+pub async fn pull_interactive_progress(
+    mut stream: Streaming<PullImageReply>,
+) -> Result<PullImageReply> {
+    let mut multi_progress: Option<(MultiProgress, HashMap<String, ProgressBar>)> = None;
+
+    while let Some(reply) = stream.next().await {
+        let reply = reply?;
+
+        if reply.progress.is_none() && !reply.digest.is_empty() {
+            return Ok(reply);
+        }
+
+        let Some(oci) = reply.progress else {
+            continue;
+        };
+
+        if multi_progress.is_none() {
+            multi_progress = Some((MultiProgress::new(), HashMap::new()));
+        }
+
+        let Some((multi_progress, progresses)) = multi_progress.as_mut() else {
+            continue;
+        };
+
+        match oci.phase() {
+            PullImageProgressPhase::Resolved
+            | PullImageProgressPhase::ConfigAcquire
+            | PullImageProgressPhase::LayerAcquire => {
+                if progresses.is_empty() && !oci.layers.is_empty() {
+                    for layer in &oci.layers {
+                        let bar = ProgressBar::new(layer.total);
+                        bar.set_style(ProgressStyle::with_template("{msg} {bar}").unwrap());
+                        progresses.insert(layer.id.clone(), bar.clone());
+                        multi_progress.add(bar);
+                    }
+                }
+
+                for layer in oci.layers {
+                    let Some(progress) = progresses.get_mut(&layer.id) else {
+                        continue;
+                    };
+
+                    let phase = match layer.phase() {
+                        PullImageProgressLayerPhase::Waiting => "waiting",
+                        PullImageProgressLayerPhase::Downloading => "downloading",
+                        PullImageProgressLayerPhase::Downloaded => "downloaded",
+                        PullImageProgressLayerPhase::Extracting => "extracting",
+                        PullImageProgressLayerPhase::Extracted => "extracted",
+                        _ => "unknown",
+                    };
+
+                    let simple = if let Some((_, hash)) = layer.id.split_once(':') {
+                        hash
+                    } else {
+                        "unknown"
+                    };
+                    let simple = if simple.len() > 10 {
+                        &simple[0..10]
+                    } else {
+                        simple
+                    };
+                    let message = format!(
+                        "{:width$} {:phwidth$}",
+                        simple,
+                        phase,
+                        width = 10,
+                        phwidth = 11
+                    );
+
+                    if message != progress.message() {
+                        progress.set_message(message);
+                    }
+
+                    progress.update(|state| {
+                        state.set_len(layer.total);
+                        state.set_pos(layer.value);
+                    });
+                }
+            }
+
+            PullImageProgressPhase::Packing => {
+                for (key, bar) in &mut *progresses {
+                    if key == "packing" {
+                        continue;
+                    }
+                    bar.finish_and_clear();
+                    multi_progress.remove(bar);
+                }
+                progresses.retain(|k, _| k == "packing");
+                if progresses.is_empty() {
+                    let progress = ProgressBar::new(100);
+                    progress.set_message("packing    ");
+                    progress.set_style(ProgressStyle::with_template("{msg} {bar}").unwrap());
+                    progresses.insert("packing".to_string(), progress);
+                }
+                let Some(progress) = progresses.get("packing") else {
+                    continue;
+                };
+
+                progress.update(|state| {
+                    state.set_len(oci.total);
+                    state.set_pos(oci.value);
+                });
+            }
+
+            _ => {}
+        }
+    }
+    Err(anyhow!("never received final reply for image pull"))
+}