Improve the API for loading shim from file or OCI

Signed-off-by: James Sturtevant <[email protected]>
containerd · Nov 16, 2023 · 3a96d65 · 3a96d65
1 parent 46b6f5f
commit 3a96d65
Show file tree

Hide file tree

Showing 9 changed files with 139 additions and 62 deletions.
diff --git a/crates/containerd-shim-wasm/src/container/context.rs b/crates/containerd-shim-wasm/src/container/context.rs
@@ -10,10 +10,6 @@ pub trait RuntimeContext {
     // path to the entrypoint executable.
     fn args(&self) -> &[String];
 
-    // ctx.entrypoint() returns the entrypoint path from arguments on the runtime
-    // spec process field.
-    fn entrypoint(&self) -> Option<&Path>;
-
     // ctx.wasi_entrypoint() returns a `WasiEntrypoint` with the path to the module to use
     // as an entrypoint and the name of the exported function to call, obtained from the
     // arguments on process OCI spec.
@@ -22,16 +18,22 @@ pub trait RuntimeContext {
     //   "/app/app.wasm#entry" -> { path: "/app/app.wasm", func: "entry" }
     //   "my_module.wat" -> { path: "my_module.wat", func: "_start" }
     //   "#init" -> { path: "", func: "init" }
-    fn wasi_entrypoint(&self) -> WasiEntrypoint;
+    fn entrypoint(&self) -> WasiEntrypoint;
 
-    fn wasm_layers(&self) -> &[WasmLayer];
+    fn wasi_loading_strategy(&self) -> WasiLoadingStrategy;
 
     fn platform(&self) -> &Platform;
 }
 
-pub struct WasiEntrypoint {
+pub enum WasiLoadingStrategy<'a> {
+    File(PathBuf),
+    Oci(&'a [WasmLayer]),
+}
+
+pub struct WasiEntrypoint<'a> {
     pub path: PathBuf,
     pub func: String,
+    pub arg0: Option<&'a Path>,
 }
 
 pub(crate) struct WasiContext<'a> {
@@ -50,21 +52,26 @@ impl RuntimeContext for WasiContext<'_> {
             .unwrap_or_default()
     }
 
-    fn entrypoint(&self) -> Option<&Path> {
-        self.args().first().map(Path::new)
-    }
+    fn entrypoint(&self) -> WasiEntrypoint {
+        let arg0 = self.args().first();
 
-    fn wasi_entrypoint(&self) -> WasiEntrypoint {
-        let arg0 = self.args().first().map(String::as_str).unwrap_or("");
-        let (path, func) = arg0.split_once('#').unwrap_or((arg0, "_start"));
+        let entry_point = arg0.map(String::as_str).unwrap_or("");
+        let (path, func) = entry_point
+            .split_once('#')
+            .unwrap_or((entry_point, "_start"));
         WasiEntrypoint {
             path: PathBuf::from(path),
             func: func.to_string(),
+            arg0: arg0.map(Path::new),
         }
     }
 
-    fn wasm_layers(&self) -> &[WasmLayer] {
-        self.wasm_layers
+    fn wasi_loading_strategy(&self) -> WasiLoadingStrategy {
+        if self.wasm_layers.is_empty() {
+            WasiLoadingStrategy::File(self.entrypoint().path.clone())
+        } else {
+            WasiLoadingStrategy::Oci(self.wasm_layers)
+        }
     }
 
     fn platform(&self) -> &Platform {
@@ -75,6 +82,7 @@ impl RuntimeContext for WasiContext<'_> {
 #[cfg(test)]
 mod tests {
     use anyhow::Result;
+    use oci_spec::image::Descriptor;
     use oci_spec::runtime::{ProcessBuilder, RootBuilder, SpecBuilder};
 
     use super::*;
@@ -167,7 +175,7 @@ mod tests {
             platform: &Platform::default(),
         };
 
-        let path = ctx.wasi_entrypoint().path;
+        let path = ctx.entrypoint().path;
         assert!(path.as_os_str().is_empty());
 
         Ok(())
@@ -195,9 +203,10 @@ mod tests {
             platform: &Platform::default(),
         };
 
-        let WasiEntrypoint { path, func } = ctx.wasi_entrypoint();
+        let WasiEntrypoint { path, func, arg0 } = ctx.entrypoint();
         assert_eq!(path, Path::new("hello.wat"));
         assert_eq!(func, "foo");
+        assert_eq!(arg0, Some(Path::new("hello.wat#foo")));
 
         Ok(())
     }
@@ -224,9 +233,74 @@ mod tests {
             platform: &Platform::default(),
         };
 
-        let WasiEntrypoint { path, func } = ctx.wasi_entrypoint();
+        let WasiEntrypoint { path, func, arg0 } = ctx.entrypoint();
         assert_eq!(path, Path::new("/root/hello.wat"));
         assert_eq!(func, "_start");
+        assert_eq!(arg0, Some(Path::new("/root/hello.wat")));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_loading_strategy_is_file_when_no_layers() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(
+                ProcessBuilder::default()
+                    .cwd("/")
+                    .args(vec![
+                        "/root/hello.wat#foo".to_string(),
+                        "echo".to_string(),
+                        "hello".to_string(),
+                    ])
+                    .build()?,
+            )
+            .build()?;
+
+        let ctx = WasiContext {
+            spec: &spec,
+            wasm_layers: &[],
+            platform: &Platform::default(),
+        };
+
+        let expected_path = PathBuf::from("/root/hello.wat");
+        assert!(matches!(
+            ctx.wasi_loading_strategy(),
+            WasiLoadingStrategy::File(p) if p == expected_path
+        ));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_loading_strategy_is_oci_when_layers_present() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(
+                ProcessBuilder::default()
+                    .cwd("/")
+                    .args(vec![
+                        "/root/hello.wat".to_string(),
+                        "echo".to_string(),
+                        "hello".to_string(),
+                    ])
+                    .build()?,
+            )
+            .build()?;
+
+        let ctx = WasiContext {
+            spec: &spec,
+            wasm_layers: &[WasmLayer {
+                layer: vec![],
+                config: Descriptor::new(oci_spec::image::MediaType::Other("".to_string()), 10, ""),
+            }],
+            platform: &Platform::default(),
+        };
+
+        assert!(matches!(
+            ctx.wasi_loading_strategy(),
+            WasiLoadingStrategy::Oci(_)
+        ));
 
         Ok(())
     }

diff --git a/crates/containerd-shim-wasm/src/container/engine.rs b/crates/containerd-shim-wasm/src/container/engine.rs
@@ -20,7 +20,7 @@ pub trait Engine: Clone + Send + Sync + 'static {
     /// * a parsable `wat` file.
     fn can_handle(&self, ctx: &impl RuntimeContext) -> Result<()> {
         let path = ctx
-            .wasi_entrypoint()
+            .entrypoint()
             .path
             .resolve_in_path_or_cwd()
             .next()
@@ -36,4 +36,8 @@ pub trait Engine: Clone + Send + Sync + 'static {
 
         Ok(())
     }
+
+    fn supported_layers_types() -> &'static [&'static str] {
+        &["application/vnd.bytecodealliance.wasm.component.layer.v0+wasm"]
+    }
 }
diff --git a/crates/containerd-shim-wasm/src/container/mod.rs b/crates/containerd-shim-wasm/src/container/mod.rs
@@ -15,7 +15,7 @@ mod engine;
 mod path;
 
 pub(crate) use context::WasiContext;
-pub use context::{RuntimeContext, WasiEntrypoint};
+pub use context::{RuntimeContext, WasiEntrypoint, WasiLoadingStrategy};
 pub use engine::Engine;
 pub use instance::Instance;
 pub use path::PathResolve;

diff --git a/crates/containerd-shim-wasm/src/sandbox/containerd.rs b/crates/containerd-shim-wasm/src/sandbox/containerd.rs
@@ -119,6 +119,7 @@ impl Client {
     pub fn load_modules(
         &self,
         containerd_id: impl ToString,
+        supported_layer_types: &[&str],
     ) -> Result<(Vec<oci::WasmLayer>, Platform)> {
         let image_name = self.get_image(containerd_id.to_string())?;
         let digest = self.get_image_content_sha(image_name)?;
@@ -141,7 +142,7 @@ impl Client {
         let layers = manifest
             .layers()
             .iter()
-            .filter(|x| !is_image_layer_type(x.media_type()))
+            .filter(|x| is_wasm_layer(x.media_type(), supported_layer_types))
             .map(|config| {
                 self.read_content(config.digest()).map(|module| WasmLayer {
                     config: config.clone(),
@@ -153,20 +154,6 @@ impl Client {
     }
 }
 
-fn is_image_layer_type(media_type: &MediaType) -> bool {
-    match media_type {
-        MediaType::ImageLayer
-        | MediaType::ImageLayerGzip
-        | MediaType::ImageLayerNonDistributable
-        | MediaType::ImageLayerNonDistributableGzip
-        | MediaType::ImageLayerNonDistributableZstd
-        | MediaType::ImageLayerZstd => true,
-        MediaType::Other(s)
-            if s.as_str()
-                .starts_with("application/vnd.docker.image.rootfs.") =>
-        {
-            true
-        }
-        _ => false,
-    }
+fn is_wasm_layer(media_type: &MediaType, supported_layer_types: &[&str]) -> bool {
+    supported_layer_types.contains(&media_type.to_string().as_str())
 }
diff --git a/crates/containerd-shim-wasm/src/sys/unix/container/executor.rs b/crates/containerd-shim-wasm/src/sys/unix/container/executor.rs
@@ -105,6 +105,7 @@ impl<E: Engine> Executor<E> {
 fn is_linux_container(ctx: &impl RuntimeContext) -> Result<()> {
     let executable = ctx
         .entrypoint()
+        .arg0
         .context("no entrypoint provided")?
         .resolve_in_path()
         .find_map(|p| -> Option<PathBuf> {

diff --git a/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs b/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs
@@ -45,7 +45,7 @@ impl<E: Engine> SandboxInstance for Instance<E> {
 
         // check if container is OCI image with wasm layers and attempt to read the module
         let (modules, platform) = containerd::Client::connect(cfg.get_containerd_address(), &namespace)?
-            .load_modules(&id)
+            .load_modules(&id, E::supported_layers_types())
             .unwrap_or_else(|e| {
                 log::warn!("Error obtaining wasm layers for container {id}.  Will attempt to use files inside container image. Error: {e}");
                 (vec![], Platform::default())

diff --git a/crates/containerd-shim-wasmedge/src/instance.rs b/crates/containerd-shim-wasmedge/src/instance.rs
@@ -1,6 +1,6 @@
 use anyhow::{bail, Context, Result};
 use containerd_shim_wasm::container::{
-    Engine, Instance, PathResolve, RuntimeContext, Stdio, WasiEntrypoint,
+    Engine, Instance, PathResolve, RuntimeContext, Stdio, WasiEntrypoint, WasiLoadingStrategy,
 };
 use log::debug;
 use wasmedge_sdk::config::{ConfigBuilder, HostRegistrationConfigOptions};
@@ -38,7 +38,8 @@ impl Engine for WasmEdgeEngine {
         let WasiEntrypoint {
             path: entrypoint_path,
             func,
-        } = ctx.wasi_entrypoint();
+            arg0: _,
+        } = ctx.entrypoint();
 
         let mut vm = self.vm.clone();
         vm.wasi_module_mut()
@@ -57,28 +58,30 @@ impl Engine for WasmEdgeEngine {
         PluginManager::load(None)?;
         let vm = vm.auto_detect_plugins()?;
 
-        let vm = match ctx.wasm_layers() {
-            [] => {
-                debug!("loading module from file");
-                let path = entrypoint_path
+        let vm = match ctx.wasi_loading_strategy() {
+            WasiLoadingStrategy::File(path) => {
+                debug!("loading module from file {path:?}");
+                let path = path
                     .resolve_in_path_or_cwd()
                     .next()
                     .context("module not found")?;
 
                 vm.register_module_from_file(&mod_name, path)
                     .context("registering module")?
             }
-            [module] => {
+            WasiLoadingStrategy::Oci([module]) => {
                 log::info!("loading module from wasm OCI layers");
                 vm.register_module_from_bytes(&mod_name, &module.layer)
                     .context("registering module")?
             }
-            [..] => bail!("only a single module is supported when using images with OCI layers"),
+            WasiLoadingStrategy::Oci(_modules) => {
+                bail!("only a single module is supported when using images with OCI layers")
+            }
         };
 
         stdio.redirect()?;
 
-        log::debug!("running {entrypoint_path:?} with method {func:?}");
+        log::debug!("running with method {func:?}");
         vm.run_func(Some(&mod_name), func, vec![])?;
 
         let status = vm

diff --git a/crates/containerd-shim-wasmer/src/instance.rs b/crates/containerd-shim-wasmer/src/instance.rs
@@ -1,6 +1,6 @@
 use anyhow::{bail, Context, Result};
 use containerd_shim_wasm::container::{
-    Engine, Instance, PathResolve, RuntimeContext, Stdio, WasiEntrypoint,
+    Engine, Instance, PathResolve, RuntimeContext, Stdio, WasiEntrypoint, WasiLoadingStrategy,
 };
 use wasmer::{Module, Store};
 use wasmer_wasix::virtual_fs::host_fs::FileSystem;
@@ -21,7 +21,11 @@ impl Engine for WasmerEngine {
     fn run_wasi(&self, ctx: &impl RuntimeContext, stdio: Stdio) -> Result<i32> {
         let args = ctx.args();
         let envs = std::env::vars();
-        let WasiEntrypoint { path, func } = ctx.wasi_entrypoint();
+        let WasiEntrypoint {
+            path,
+            func,
+            arg0: _,
+        } = ctx.entrypoint();
 
         let mod_name = match path.file_stem() {
             Some(name) => name.to_string_lossy().to_string(),
@@ -31,8 +35,8 @@ impl Engine for WasmerEngine {
         log::info!("Create a Store");
         let mut store = Store::new(self.engine.clone());
 
-        let module = match ctx.wasm_layers() {
-            [] => {
+        let module = match ctx.wasi_loading_strategy() {
+            WasiLoadingStrategy::File(path) => {
                 log::info!("loading module from file {path:?}");
                 let path = path
                     .resolve_in_path_or_cwd()
@@ -41,11 +45,14 @@ impl Engine for WasmerEngine {
 
                 Module::from_file(&store, path)?
             }
-            [module] => {
+            WasiLoadingStrategy::Oci([module]) => {
+                log::info!("loading module wasm OCI layers");
                 log::info!("loading module wasm OCI layers");
                 Module::from_binary(&store, &module.layer)?
             }
-            [..] => bail!("only a single module is supported when using images with OCI layers"),
+            WasiLoadingStrategy::Oci(_modules) => {
+                bail!("only a single module is supported when using images with OCI layers")
+            }
         };
 
         let runtime = tokio::runtime::Builder::new_multi_thread()