Refactor to put logic in one place

Signed-off-by: James Sturtevant <jstur@microsoft.com>
containerd · Aug 1, 2023 · 8631d3a · 8631d3a
1 parent d02c2c0
commit 8631d3a
Show file tree

Hide file tree

Showing 5 changed files with 277 additions and 84 deletions.
diff --git a/crates/containerd-shim-wasm/src/sandbox/oci.rs b/crates/containerd-shim-wasm/src/sandbox/oci.rs
@@ -36,6 +36,25 @@ pub fn get_args(spec: &Spec) -> &[String] {
     }
 }
 
+pub fn get_module(spec: &Spec) -> (Option<String>, String) {
+    let args = get_args(spec);
+
+    if !args.is_empty() {
+        let start = args[0].clone();
+        let mut iterator = start.split('#');
+        let mut cmd = iterator.next().unwrap().to_string();
+
+        let stripped = cmd.strip_prefix(std::path::MAIN_SEPARATOR);
+        if let Some(strpd) = stripped {
+            cmd = strpd.to_string();
+        }
+        let method = iterator.next().unwrap_or("_start");
+        return (Some(cmd), method.to_string());
+    }
+
+    (None, "_start".to_string())
+}
+
 pub fn spec_from_file<P: AsRef<Path>>(path: P) -> Result<Spec> {
     let file = File::open(path)?;
     let cfg: Spec = json::from_reader(file)?;
@@ -160,3 +179,125 @@ pub fn setup_prestart_hooks(hooks: &Option<oci_spec::runtime::Hooks>) -> Result<
     }
     Ok(())
 }
+
+#[cfg(test)]
+mod oci_tests {
+    use super::*;
+    use oci_spec::runtime::{ProcessBuilder, RootBuilder, SpecBuilder};
+
+    #[test]
+    fn test_get_args() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(
+                ProcessBuilder::default()
+                    .cwd("/")
+                    .args(vec!["hello.wat".to_string()])
+                    .build()?,
+            )
+            .build()?;
+
+        let args = get_args(&spec);
+        assert_eq!(args.len(), 1);
+        assert_eq!(args[0], "hello.wat");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_args_return_empty() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(ProcessBuilder::default().cwd("/").args(vec![]).build()?)
+            .build()?;
+
+        let args = get_args(&spec);
+        assert_eq!(args.len(), 0);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_args_returns_all() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(
+                ProcessBuilder::default()
+                    .cwd("/")
+                    .args(vec![
+                        "hello.wat".to_string(),
+                        "echo".to_string(),
+                        "hello".to_string(),
+                    ])
+                    .build()?,
+            )
+            .build()?;
+
+        let args = get_args(&spec);
+        assert_eq!(args.len(), 3);
+        assert_eq!(args[0], "hello.wat");
+        assert_eq!(args[1], "echo");
+        assert_eq!(args[2], "hello");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_module_returns_none_when_not_present() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(ProcessBuilder::default().cwd("/").args(vec![]).build()?)
+            .build()?;
+
+        let (module, _) = get_module(&spec);
+        assert_eq!(module, None);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_module_returns_function() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(
+                ProcessBuilder::default()
+                    .cwd("/")
+                    .args(vec![
+                        "hello.wat#foo".to_string(),
+                        "echo".to_string(),
+                        "hello".to_string(),
+                    ])
+                    .build()?,
+            )
+            .build()?;
+
+        let (module, function) = get_module(&spec);
+        assert_eq!(module, Some("hello.wat".to_string()));
+        assert_eq!(function, "foo");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_get_module_returns_start() -> Result<()> {
+        let spec = SpecBuilder::default()
+            .root(RootBuilder::default().path("rootfs").build()?)
+            .process(
+                ProcessBuilder::default()
+                    .cwd("/")
+                    .args(vec![
+                        "hello.wat".to_string(),
+                        "echo".to_string(),
+                        "hello".to_string(),
+                    ])
+                    .build()?,
+            )
+            .build()?;
+
+        let (module, function) = get_module(&spec);
+        assert_eq!(module, Some("hello.wat".to_string()));
+        assert_eq!(function, "_start");
+
+        Ok(())
+    }
+}
diff --git a/crates/containerd-shim-wasmedge/src/executor.rs b/crates/containerd-shim-wasmedge/src/executor.rs
@@ -6,6 +6,7 @@ use oci_spec::runtime::Spec;
 use libc::{STDERR_FILENO, STDIN_FILENO, STDOUT_FILENO};
 use libcontainer::workload::{Executor, ExecutorError};
 use std::os::unix::io::RawFd;
+use log::debug;
 
 use wasmedge_sdk::{
     config::{CommonConfigOptions, ConfigBuilder, HostRegistrationConfigOptions},
@@ -34,7 +35,9 @@ impl Executor for WasmEdgeExecutor {
 
         // TODO: How to get exit code?
         // This was relatively straight forward in go, but wasi and wasmtime are totally separate things in rust
-        match vm.run_func(Some("main"), "_start", params!()) {
+        let (module_name, method) = oci::get_module(spec);
+        debug!("running {:?} with method {}", module_name, method);
+        match vm.run_func(Some("main"), method, params!()) {
             Ok(_) => std::process::exit(0),
             Err(_) => std::process::exit(137),
         };
@@ -51,10 +54,6 @@ impl Executor for WasmEdgeExecutor {
 
 impl WasmEdgeExecutor {
     fn prepare(&self, args: &[String], spec: &Spec) -> anyhow::Result<wasmedge_sdk::Vm> {
-        let mut cmd = args[0].clone();
-        if let Some(stripped) = args[0].strip_prefix(std::path::MAIN_SEPARATOR) {
-            cmd = stripped.to_string();
-        }
         let envs = env_to_wasi(spec);
         let config = ConfigBuilder::new(CommonConfigOptions::default())
             .with_host_registration_config(HostRegistrationConfigOptions::default().wasi(true))
@@ -73,8 +72,18 @@ impl WasmEdgeExecutor {
             Some(envs.iter().map(|s| s as &str).collect()),
             None,
         );
+
+        let (module_name, _) = oci::get_module(spec);
+        let module_name = match module_name {
+            Some(m) => m,
+            None => {
+                return Err(anyhow::Error::msg(
+                    "no module provided cannot load module",
+                ))
+            }
+        };
         let vm = vm
-            .register_module_from_file("main", cmd)
+            .register_module_from_file("main", module_name)
             .map_err(|err| ExecutorError::Execution(err))?;
         if let Some(stdin) = self.stdin {
             dup(STDIN_FILENO)?;

diff --git a/crates/containerd-shim-wasmedge/src/instance.rs b/crates/containerd-shim-wasmedge/src/instance.rs
@@ -430,9 +430,9 @@ mod wasitest {
 
         let dir = tempdir()?;
         let path = dir.path();
-        let wasmbytes = wat2wasm(WASI_HELLO_WAT).unwrap();
+        let wasm_bytes = wat2wasm(WASI_HELLO_WAT).unwrap();
 
-        let res = run_wasi_test(&dir, wasmbytes)?;
+        let res = run_wasi_test(&dir, wasm_bytes)?;
 
         assert_eq!(res.0, 0);
 
@@ -452,9 +452,9 @@ mod wasitest {
         }
 
         let dir = tempdir()?;
-        let wasmbytes = wat2wasm(WASI_RETURN_ERROR).unwrap();
+        let wasm_bytes = wat2wasm(WASI_RETURN_ERROR).unwrap();
 
-        let res = run_wasi_test(&dir, wasmbytes)?;
+        let res = run_wasi_test(&dir, wasm_bytes)?;
 
         // Expect error code from the run.
         assert_eq!(res.0, 137);

diff --git a/crates/containerd-shim-wasmtime/src/executor.rs b/crates/containerd-shim-wasmtime/src/executor.rs
@@ -1,7 +1,7 @@
 use nix::unistd::{dup, dup2};
 use std::{fs::OpenOptions, os::fd::RawFd};
 
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, Result};
 use containerd_shim_wasm::sandbox::oci;
 use libc::{STDERR_FILENO, STDIN_FILENO, STDOUT_FILENO};
 use libcontainer::workload::{Executor, ExecutorError};
@@ -82,20 +82,18 @@ impl WasmtimeExecutor {
         let wctx = wasi_builder.build();
 
         log::info!("wasi context ready");
-        let mut iterator = args
-            .first()
-            .context("args must have at least one argument.")?
-            .split('#');
-        let mut cmd = iterator.next().unwrap().to_string();
-        let stripped = cmd.strip_prefix(std::path::MAIN_SEPARATOR);
-        if let Some(strpd) = stripped {
-            cmd = strpd.to_string();
-        }
-        let method = iterator.next().unwrap_or("_start");
-        let mod_path = cmd;
+        let (module_name, method) = oci::get_module(spec);
+        let module_name = match module_name {
+            Some(m) => m,
+            None => {
+                return Err(anyhow::format_err!(
+                    "no module provided, cannot load module from file within container"
+                ))
+            }
+        };
 
-        log::info!("loading module from file");
-        let module = Module::from_file(&self.engine, mod_path)?;
+        log::info!("loading module from file {} ", module_name);
+        let module = Module::from_file(&self.engine, module_name)?;
         let mut linker = Linker::new(&self.engine);
 
         wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
@@ -106,7 +104,7 @@ impl WasmtimeExecutor {
 
         log::info!("getting start function");
         let start_func = instance
-            .get_func(&mut store, method)
+            .get_func(&mut store, &method)
             .ok_or_else(|| anyhow!("module does not have a WASI start function".to_string()))?;
         Ok((store, start_func))
     }