Merge pull request #2332 from fermyon/backport-2318-to-v2.3

[Backport v2.3] feat(loader): support loading AOT compiled components

crates/trigger/Cargo.toml

@@ -8,6 +8,10 @@ edition = { workspace = true }
 llm = ["spin-llm-local"]
 llm-metal = ["llm", "spin-llm-local/metal"]
 llm-cublas = ["llm", "spin-llm-local/cublas"]
+# Enables loading AOT compiled components, a potentially unsafe operation. See
+# `<TriggerLoader as Loader>::::enable_loading_aot_compiled_components`
+# documentation for more information about the safety risks.
+unsafe-aot-compilation = []
 
 [dependencies]
 anyhow = "1.0"
@@ -54,4 +58,4 @@ wasmtime-wasi = { workspace = true }
 wasmtime-wasi-http = { workspace = true }
 
 [dev-dependencies]
-tempfile = "3.8.0"
+tempfile = "3.8.0"

crates/trigger/src/loader.rs

@@ -1,5 +1,3 @@
-#![allow(dead_code)] // Refactor WIP
-
 use std::path::PathBuf;
 
 use anyhow::{ensure, Context, Result};
@@ -13,18 +11,65 @@ use tokio::fs;
 
 use spin_common::{ui::quoted_path, url::parse_file_url};
 
+/// Compilation status of all components of a Spin application
+pub enum CompilationStatus {
+    #[cfg(feature = "unsafe-aot-compilation")]
+    /// All components are componentized and ahead of time (AOT) compiled to cwasm.
+    AllAotComponents,
+    /// No components are ahead of time (AOT) compiled.
+    NoneAot,
+}
+
+/// Loader for the Spin runtime
 pub struct TriggerLoader {
+    /// Working directory where files for mounting exist.
     working_dir: PathBuf,
+    /// Set the static assets of the components in the temporary directory as writable.
     allow_transient_write: bool,
+    /// Declares the compilation status of all components of a Spin application.
+    compilation_status: CompilationStatus,
 }
 
 impl TriggerLoader {
     pub fn new(working_dir: impl Into<PathBuf>, allow_transient_write: bool) -> Self {
         Self {
             working_dir: working_dir.into(),
             allow_transient_write,
+            compilation_status: CompilationStatus::NoneAot,
         }
     }
+
+    /// Updates the TriggerLoader to load AOT precompiled components
+    ///
+    /// **Warning: This feature may bypass important security guarantees of the
+    /// Wasmtime
+    // security sandbox if used incorrectly! Read this documentation
+    // carefully.**
+    ///
+    /// Usually, components are compiled just-in-time from portable Wasm
+    /// sources. This method causes components to instead be loaded
+    /// ahead-of-time as Wasmtime-precompiled native executable binaries.
+    /// Precompiled binaries must be produced with a compatible Wasmtime engine
+    /// using the same Wasmtime version and compiler target settings - typically
+    /// by a host with the same processor that will be executing them. See the
+    /// Wasmtime documentation for more information:
+    /// https://docs.rs/wasmtime/latest/wasmtime/struct.Module.html#method.deserialize
+    ///
+    /// # Safety
+    ///
+    /// This method is marked as `unsafe` because it enables potentially unsafe
+    /// behavior if
+    // used to load malformed or malicious precompiled binaries. Loading sources
+    // from an
+    /// incompatible Wasmtime engine will fail but is otherwise safe. This
+    /// method is safe if it can be guaranteed that `<TriggerLoader as
+    /// Loader>::load_component` will only ever be called with a trusted
+    /// `LockedComponentSource`. **Precompiled binaries must never be loaded
+    /// from untrusted sources.**
+    #[cfg(feature = "unsafe-aot-compilation")]
+    pub unsafe fn enable_loading_aot_compiled_components(&mut self) {
+        self.compilation_status = CompilationStatus::AllAotComponents;
+    }
 }
 
 #[async_trait]
@@ -49,15 +94,34 @@ impl Loader for TriggerLoader {
             .as_ref()
             .context("LockedComponentSource missing source field")?;
         let path = parse_file_url(source)?;
-        let bytes = fs::read(&path).await.with_context(|| {
-            format!(
-                "failed to read component source from disk at path '{}'",
-                path.display()
-            )
-        })?;
-        let component = spin_componentize::componentize_if_necessary(&bytes)?;
-        spin_core::Component::new(engine, component.as_ref())
-            .with_context(|| format!("loading module {}", quoted_path(&path)))
+        match self.compilation_status {
+            #[cfg(feature = "unsafe-aot-compilation")]
+            CompilationStatus::AllAotComponents => {
+                match engine.detect_precompiled_file(&path)?{
+                    Some(wasmtime::Precompiled::Component) => {
+                        unsafe {
+                            spin_core::Component::deserialize_file(engine, &path)
+                                .with_context(|| format!("deserializing module {}", quoted_path(&path)))
+                        }
+                    },
+                    Some(wasmtime::Precompiled::Module) => anyhow::bail!("Spin loader is configured to load only AOT compiled components but an AOT compiled module provided at {}", quoted_path(&path)),
+                    None => {
+                        anyhow::bail!("Spin loader is configured to load only AOT compiled components, but {} is not precompiled", quoted_path(&path))
+                    }
+                }
+            },
+            CompilationStatus::NoneAot => {
+                let bytes = fs::read(&path).await.with_context(|| {
+                format!(
+                    "failed to read component source from disk at path {}",
+                    quoted_path(&path)
+                )
+            })?;
+            let component = spin_componentize::componentize_if_necessary(&bytes)?;
+            spin_core::Component::new(engine, component.as_ref())
+                .with_context(|| format!("loading module {}", quoted_path(&path)))
+            }
+        }
     }
 
     async fn load_module(
@@ -102,3 +166,54 @@ impl Loader for TriggerLoader {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use spin_app::locked::ContentRef;
+    use std::io::Write;
+    use tempfile::NamedTempFile;
+
+    fn precompiled_component(file: &mut NamedTempFile) -> LockedComponentSource {
+        let wasmtime_engine = wasmtime::Engine::default();
+        let component = wasmtime::component::Component::new(&wasmtime_engine, "(component)")
+            .unwrap()
+            .serialize()
+            .unwrap();
+        let file_uri = format!("file://{}", file.path().to_str().unwrap());
+        file.write_all(&component).unwrap();
+        LockedComponentSource {
+            content: ContentRef {
+                source: Some(file_uri),
+                ..Default::default()
+            },
+            content_type: "application/wasm".to_string(),
+        }
+    }
+
+    #[cfg(feature = "unsafe-aot-compilation")]
+    #[tokio::test]
+    async fn load_component_succeeds_for_precompiled_component() {
+        let mut file = NamedTempFile::new().unwrap();
+        let source = precompiled_component(&mut file);
+        let mut loader = super::TriggerLoader::new("/unreferenced", false);
+        unsafe {
+            loader.enable_loading_aot_compiled_components();
+        }
+        loader
+            .load_component(&spin_core::wasmtime::Engine::default(), &source)
+            .await
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn load_component_fails_for_precompiled_component() {
+        let mut file = NamedTempFile::new().unwrap();
+        let source = precompiled_component(&mut file);
+        let loader = super::TriggerLoader::new("/unreferenced", false);
+        let result = loader
+            .load_component(&spin_core::wasmtime::Engine::default(), &source)
+            .await;
+        assert!(result.is_err());
+    }
+}