Add an API to detect precompiled modules/components (bytecodealliance…

…#6832)

This commit adds a new `Engine::detect_precompiled` API to inspect some
bytes and determine if they look like a precompiled artifact of either a
core wasm module or component. This is something I'll be using soon in
an upcoming refactor of the Wasmtime CLI to support components, but it's
something we've also talked about before which can be useful for systems
storing both precompiled modules and components.

Implementation-wise this looks at the ELF header of the input and
determines if it's got all the right flags that Wasmtime sets for the
various bits and bobs of our object format.

crates/wasmtime/src/engine.rs

@@ -629,6 +629,26 @@ impl Engine {
         code.publish()?;
         Ok(Arc::new(code))
     }
+
+    /// Detects whether the bytes provided are a precompiled object produced by
+    /// Wasmtime.
+    ///
+    /// This function will inspect the header of `bytes` to determine if it
+    /// looks like a precompiled core wasm module or a precompiled component.
+    /// This does not validate the full structure or guarantee that
+    /// deserialization will succeed, instead it helps higher-levels of the
+    /// stack make a decision about what to do next when presented with the
+    /// `bytes` as an input module.
+    ///
+    /// If the `bytes` looks like a precompiled object previously produced by
+    /// [`Module::serialize`](crate::Module::serialize),
+    /// [`Component::serialize`](crate::component::Component::serialize),
+    /// [`Engine::precompile_module`], or [`Engine::precompile_component`], then
+    /// this will return `Some(...)` indicating so. Otherwise `None` is
+    /// returned.
+    pub fn detect_precompiled(&self, bytes: &[u8]) -> Option<Precompiled> {
+        serialization::detect_precompiled(bytes)
+    }
 }
 
 impl Default for Engine {
@@ -637,6 +657,15 @@ impl Default for Engine {
     }
 }
 
+/// Return value from the [`Engine::detect_precompiled`] API.
+#[derive(PartialEq, Eq, Copy, Clone, Debug)]
+pub enum Precompiled {
+    /// The input bytes look like a precompiled core wasm module.
+    Module,
+    /// The input bytes look like a precompiled wasm component.
+    Component,
+}
+
 #[cfg(test)]
 mod tests {
     use std::{

crates/wasmtime/src/engine/serialization.rs

@@ -21,7 +21,7 @@
 //! other random ELF files, as well as provide better error messages for
 //! using wasmtime artifacts across versions.
 
-use crate::{Engine, ModuleVersionStrategy};
+use crate::{Engine, ModuleVersionStrategy, Precompiled};
 use anyhow::{anyhow, bail, Context, Result};
 use object::write::{Object, StandardSegment};
 use object::{File, FileFlags, Object as _, ObjectSection, SectionKind};
@@ -145,6 +145,23 @@ pub fn check_compatible(engine: &Engine, mmap: &MmapVec, expected: ObjectKind) -
     bincode::deserialize::<Metadata>(data)?.check_compatible(engine)
 }
 
+pub fn detect_precompiled(bytes: &[u8]) -> Option<Precompiled> {
+    let obj = File::parse(bytes).ok()?;
+    match obj.flags() {
+        FileFlags::Elf {
+            os_abi: obj::ELFOSABI_WASMTIME,
+            abi_version: 0,
+            e_flags: obj::EF_WASMTIME_MODULE,
+        } => Some(Precompiled::Module),
+        FileFlags::Elf {
+            os_abi: obj::ELFOSABI_WASMTIME,
+            abi_version: 0,
+            e_flags: obj::EF_WASMTIME_COMPONENT,
+        } => Some(Precompiled::Component),
+        _ => None,
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 struct Metadata {
     target: String,

tests/all/component_model/aot.rs

@@ -1,6 +1,6 @@
 use anyhow::Result;
 use wasmtime::component::{Component, Linker};
-use wasmtime::{Module, Store};
+use wasmtime::{Module, Precompiled, Store};
 
 #[test]
 fn module_component_mismatch() -> Result<()> {
@@ -119,3 +119,17 @@ fn usable_exported_modules() -> Result<()> {
     core_linker.instantiate(&mut store, &module)?;
     Ok(())
 }
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn detect_precompiled() -> Result<()> {
+    let engine = super::engine();
+    let buffer = Component::new(&engine, "(component)")?.serialize()?;
+    assert_eq!(engine.detect_precompiled(&[]), None);
+    assert_eq!(engine.detect_precompiled(&buffer[..5]), None);
+    assert_eq!(
+        engine.detect_precompiled(&buffer),
+        Some(Precompiled::Component)
+    );
+    Ok(())
+}

tests/all/module_serialize.rs

@@ -119,3 +119,20 @@ fn deserialize_from_serialized() -> Result<()> {
     assert!(buffer1 == buffer2);
     Ok(())
 }
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn detect_precompiled() -> Result<()> {
+    let engine = Engine::default();
+    let buffer = serialize(
+        &engine,
+        "(module (func (export \"run\") (result i32) i32.const 42))",
+    )?;
+    assert_eq!(engine.detect_precompiled(&[]), None);
+    assert_eq!(engine.detect_precompiled(&buffer[..5]), None);
+    assert_eq!(
+        engine.detect_precompiled(&buffer),
+        Some(Precompiled::Module)
+    );
+    Ok(())
+}