Merge pull request #2818 from lann/reimplement-aot-loader

Reimplement unsafe-aot-compilation feature

crates/trigger/Cargo.toml

@@ -8,6 +8,12 @@ homepage.workspace = true
 repository.workspace = true
 rust-version.workspace = true
 
+[features]
+# Enables loading AOT compiled components, a potentially unsafe operation. See
+# `ComponentLoader::enable_loading_aot_compiled_components`
+# documentation for more information about the safety risks.
+unsafe-aot-compilation = []
+
 [dependencies]
 anyhow = "1"
 clap = { version = "3.1.18", features = ["derive", "env"] }

crates/trigger/src/cli.rs

@@ -13,10 +13,10 @@ use spin_app::App;
 use spin_common::sloth;
 use spin_common::ui::quoted_path;
 use spin_common::url::parse_file_url;
-use spin_core::async_trait;
 use spin_factors::RuntimeFactors;
-use spin_factors_executor::{ComponentLoader, FactorsExecutor};
+use spin_factors_executor::FactorsExecutor;
 
+use crate::loader::ComponentLoader;
 use crate::{Trigger, TriggerApp};
 pub use initial_kv_setter::InitialKvSetterHook;
 pub use launch_metadata::LaunchMetadata;
@@ -312,73 +312,14 @@ impl<T: Trigger<B::Factors>, B: RuntimeFactorsBuilder> TriggerAppBuilder<T, B> {
 
         let (factors, runtime_config) = B::build(&common_options, &options)?;
 
-        // TODO: port the rest of the component loader logic
-        struct SimpleComponentLoader;
-
-        #[async_trait]
-        impl spin_compose::ComponentSourceLoader for SimpleComponentLoader {
-            async fn load_component_source(
-                &self,
-                source: &spin_app::locked::LockedComponentSource,
-            ) -> anyhow::Result<Vec<u8>> {
-                let source = source
-                    .content
-                    .source
-                    .as_ref()
-                    .context("LockedComponentSource missing source field")?;
-
-                let path = parse_file_url(source)?;
-
-                let bytes: Vec<u8> = tokio::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)?;
-
-                Ok(component.into())
-            }
-        }
-
-        #[async_trait]
-        impl ComponentLoader for SimpleComponentLoader {
-            async fn load_component(
-                &self,
-                engine: &spin_core::wasmtime::Engine,
-                component: &spin_factors::AppComponent,
-            ) -> anyhow::Result<spin_core::Component> {
-                let source = component
-                    .source()
-                    .content
-                    .source
-                    .as_ref()
-                    .context("LockedComponentSource missing source field")?;
-                let path = parse_file_url(source)?;
-
-                let composed = spin_compose::compose(self, component.locked)
-                    .await
-                    .with_context(|| {
-                        format!(
-                            "failed to resolve dependencies for component {:?}",
-                            component.locked.id
-                        )
-                    })?;
-
-                spin_core::Component::new(engine, composed)
-                    .with_context(|| format!("compiling wasm {}", quoted_path(&path)))
-            }
-        }
-
         let mut executor = FactorsExecutor::new(core_engine_builder, factors)?;
         B::configure_app(&mut executor, &runtime_config, &common_options, &options)?;
         let executor = Arc::new(executor);
 
         let configured_app = {
             let _sloth_guard = warn_if_wasm_build_slothful();
             executor
-                .load_app(app, runtime_config.into(), &SimpleComponentLoader)
+                .load_app(app, runtime_config.into(), &ComponentLoader::default())
                 .await?
         };
 

crates/trigger/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod cli;
+pub mod loader;
 
 use std::future::Future;
 

crates/trigger/src/loader.rs

@@ -0,0 +1,126 @@
+use anyhow::Context as _;
+use spin_common::{ui::quoted_path, url::parse_file_url};
+use spin_core::{async_trait, wasmtime, Component};
+use spin_factors::AppComponent;
+
+#[derive(Default)]
+pub struct ComponentLoader {
+    #[cfg(feature = "unsafe-aot-compilation")]
+    aot_compilation_enabled: bool,
+}
+
+impl ComponentLoader {
+    /// 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.aot_compilation_enabled = true;
+    }
+
+    #[cfg(feature = "unsafe-aot-compilation")]
+    fn load_precompiled_component(
+        &self,
+        engine: &wasmtime::Engine,
+        path: &std::path::Path,
+    ) -> anyhow::Result<Component> {
+        assert!(self.aot_compilation_enabled);
+        match engine.detect_precompiled_file(path)? {
+            Some(wasmtime::Precompiled::Component) => unsafe {
+                Component::deserialize_file(engine, path)
+            },
+            Some(wasmtime::Precompiled::Module) => {
+                anyhow::bail!("expected AOT compiled component but found module");
+            }
+            None => {
+                anyhow::bail!("expected AOT compiled component but found other data");
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl spin_factors_executor::ComponentLoader for ComponentLoader {
+    async fn load_component(
+        &self,
+        engine: &wasmtime::Engine,
+        component: &AppComponent,
+    ) -> anyhow::Result<Component> {
+        let source = component
+            .source()
+            .content
+            .source
+            .as_ref()
+            .context("LockedComponentSource missing source field")?;
+        let path = parse_file_url(source)?;
+
+        #[cfg(feature = "unsafe-aot-compilation")]
+        if self.aot_compilation_enabled {
+            return self
+                .load_precompiled_component(engine, &path)
+                .with_context(|| format!("error deserializing component from {path:?}"));
+        }
+
+        let composed = spin_compose::compose(&ComponentSourceLoader, component.locked)
+            .await
+            .with_context(|| {
+                format!(
+                    "failed to resolve dependencies for component {:?}",
+                    component.locked.id
+                )
+            })?;
+
+        spin_core::Component::new(engine, composed)
+            .with_context(|| format!("failed to compile component from {}", quoted_path(&path)))
+    }
+}
+
+struct ComponentSourceLoader;
+
+#[async_trait]
+impl spin_compose::ComponentSourceLoader for ComponentSourceLoader {
+    async fn load_component_source(
+        &self,
+        source: &spin_app::locked::LockedComponentSource,
+    ) -> anyhow::Result<Vec<u8>> {
+        let source = source
+            .content
+            .source
+            .as_ref()
+            .context("LockedComponentSource missing source field")?;
+
+        let path = parse_file_url(source)?;
+
+        let bytes: Vec<u8> = tokio::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)?;
+
+        Ok(component.into())
+    }
+}