spin_componentize/
lib.rs

1#![deny(warnings)]
2
3use {
4    anyhow::{anyhow, Context, Result},
5    module_info::ModuleInfo,
6    std::{borrow::Cow, collections::HashSet},
7    wasm_encoder::reencode::{Reencode, RoundtripReencoder},
8    wasm_encoder::{CustomSection, ExportSection, ImportSection, Module, RawSection},
9    wasmparser::{Encoding, Parser, Payload},
10    wit_component::{metadata, ComponentEncoder},
11};
12
13pub mod bugs;
14
15#[cfg(test)]
16mod abi_conformance;
17mod module_info;
18
19const SPIN_ADAPTER: &[u8] = include_bytes!(concat!(
20    env!("OUT_DIR"),
21    "/wasm32-unknown-unknown/release/wasi_snapshot_preview1_spin.wasm"
22));
23const PREVIEW1_ADAPTER: &[u8] = include_bytes!(concat!(
24    env!("OUT_DIR"),
25    "/wasm32-unknown-unknown/release/wasi_snapshot_preview1_upstream.wasm"
26));
27
28const COMMAND_ADAPTER: &[u8] = include_bytes!(concat!(
29    env!("OUT_DIR"),
30    "/wasm32-unknown-unknown/release/wasi_snapshot_preview1_command.wasm"
31));
32
33static ADAPTER_NAME: &str = "wasi_snapshot_preview1";
34static CUSTOM_SECTION_NAME: &str = "component-type:reactor";
35static WORLD_NAME: &str = "reactor";
36
37static EXPORT_INTERFACES: &[(&str, &str)] = &[
38    ("handle-redis-message", "inbound-redis"),
39    ("handle-http-request", "inbound-http"),
40];
41
42pub fn componentize_if_necessary(module_or_component: &[u8]) -> Result<Cow<[u8]>> {
43    for payload in Parser::new(0).parse_all(module_or_component) {
44        if let Payload::Version { encoding, .. } = payload.context("unable to parse binary")? {
45            return match encoding {
46                Encoding::Component => Ok(Cow::Borrowed(module_or_component)),
47                Encoding::Module => componentize(module_or_component).map(Cow::Owned),
48            };
49        }
50    }
51    Err(anyhow!("unable to determine wasm binary encoding"))
52}
53
54pub fn componentize(module: &[u8]) -> Result<Vec<u8>> {
55    let module_info = ModuleInfo::from_module(module)?;
56    match WitBindgenVersion::detect(&module_info)? {
57        WitBindgenVersion::V0_2OrNone => componentize_old_module(module, &module_info),
58        WitBindgenVersion::GreaterThanV0_4 => componentize_new_bindgen(module),
59        WitBindgenVersion::Other(other) => Err(anyhow::anyhow!(
60            "cannot adapt modules created with wit-bindgen version {other}"
61        )),
62    }
63}
64
65/// In order to properly componentize modules, we need to know which
66/// version of wit-bindgen was used
67#[derive(Debug)]
68enum WitBindgenVersion {
69    GreaterThanV0_4,
70    V0_2OrNone,
71    Other(String),
72}
73
74impl WitBindgenVersion {
75    fn detect(module_info: &ModuleInfo) -> Result<Self> {
76        if let Some(processors) = module_info.bindgen_processors() {
77            let bindgen_version = processors
78                .iter()
79                .find_map(|(key, value)| key.starts_with("wit-bindgen").then_some(value.as_str()));
80            if let Some(v) = bindgen_version {
81                let mut parts = v.split('.');
82                let Some(major) = parts.next().and_then(|p| p.parse::<u8>().ok()) else {
83                    return Ok(Self::Other(v.to_owned()));
84                };
85                let Some(minor) = parts.next().and_then(|p| p.parse::<u8>().ok()) else {
86                    return Ok(Self::Other(v.to_owned()));
87                };
88                if (major == 0 && minor < 5) || major >= 1 {
89                    return Ok(Self::Other(v.to_owned()));
90                }
91                // Either there should be no patch version or nothing after patch
92                if parts.next().is_none() || parts.next().is_none() {
93                    return Ok(Self::GreaterThanV0_4);
94                } else {
95                    return Ok(Self::Other(v.to_owned()));
96                }
97            }
98        }
99        Ok(Self::V0_2OrNone)
100    }
101}
102
103/// Modules produced with wit-bindgen 0.5 and newer only need wasi preview 1 to preview 2 adapter
104pub fn componentize_new_bindgen(module: &[u8]) -> Result<Vec<u8>> {
105    ComponentEncoder::default()
106        .validate(true)
107        .module(module)?
108        .adapter("wasi_snapshot_preview1", PREVIEW1_ADAPTER)?
109        .encode()
110}
111
112/// Modules *not* produced with wit-bindgen >= 0.5 could be old wit-bindgen or no wit-bindgen
113pub fn componentize_old_module(module: &[u8], module_info: &ModuleInfo) -> Result<Vec<u8>> {
114    // If the module has a _start export and doesn't obviously use wit-bindgen
115    // it is likely an old p1 command module.
116    if module_info.has_start_export && !module_info.probably_uses_wit_bindgen() {
117        bugs::WasiLibc377Bug::check(module_info)?;
118        componentize_command(module)
119    } else {
120        componentize_old_bindgen(module)
121    }
122}
123
124/// Modules produced with wit-bindgen 0.2 need more extensive adaption
125pub fn componentize_old_bindgen(module: &[u8]) -> Result<Vec<u8>> {
126    let (module, exports) = retarget_imports_and_get_exports(ADAPTER_NAME, module)?;
127    let allowed = exports
128        .into_iter()
129        .filter_map(|export| {
130            EXPORT_INTERFACES
131                .iter()
132                .find_map(|(k, v)| (*k == export).then_some(*v))
133        })
134        .collect::<HashSet<&str>>();
135
136    let (adapter, mut bindgen) = metadata::decode(SPIN_ADAPTER)?;
137    let adapter = adapter.expect(
138        "adapter module was malformed, and did not contain a 'component-type' custom section",
139    );
140
141    let world = bindgen
142        .resolve
143        .worlds
144        .iter()
145        .find_map(|(k, v)| (v.name == WORLD_NAME).then_some(k))
146        .ok_or_else(|| anyhow!("world not found: {WORLD_NAME}"))?;
147
148    bindgen.resolve.worlds[world].exports.retain(|k, _| {
149        let k = match &k {
150            wit_parser::WorldKey::Name(n) => n,
151            wit_parser::WorldKey::Interface(i) => match &bindgen.resolve.interfaces[*i].name {
152                Some(n) => n,
153                None => return true,
154            },
155        };
156        allowed.contains(k.as_str())
157    });
158
159    let body = metadata::encode(
160        &bindgen.resolve,
161        world,
162        wit_component::StringEncoding::UTF8,
163        None,
164    )?;
165
166    let adapter = add_custom_section(CUSTOM_SECTION_NAME, &body, &adapter)?;
167
168    ComponentEncoder::default()
169        .validate(true)
170        .module(&module)?
171        .adapter(ADAPTER_NAME, &adapter)?
172        .encode()
173}
174
175pub fn componentize_command(module: &[u8]) -> Result<Vec<u8>> {
176    ComponentEncoder::default()
177        .validate(true)
178        .module(module)?
179        .adapter(ADAPTER_NAME, COMMAND_ADAPTER)?
180        .encode()
181}
182
183fn retarget_imports_and_get_exports(target: &str, module: &[u8]) -> Result<(Vec<u8>, Vec<String>)> {
184    let mut result = Module::new();
185    let mut exports_result = Vec::new();
186
187    for payload in Parser::new(0).parse_all(module) {
188        match payload? {
189            Payload::ImportSection(reader) => {
190                let mut imports = ImportSection::new();
191                for import in reader {
192                    let import = import?;
193                    let (module, field) = if import.module == target {
194                        (Cow::Borrowed(import.module), Cow::Borrowed(import.name))
195                    } else {
196                        (
197                            Cow::Borrowed(target),
198                            Cow::Owned(format!("{}:{}", import.module, import.name)),
199                        )
200                    };
201                    let ty = RoundtripReencoder.entity_type(import.ty)?;
202                    imports.import(&module, &field, ty);
203                }
204                result.section(&imports);
205            }
206
207            Payload::ExportSection(reader) => {
208                let mut exports = ExportSection::new();
209                for export in reader {
210                    let export = export?;
211                    exports_result.push(export.name.to_owned());
212                    let kind = RoundtripReencoder.export_kind(export.kind);
213                    exports.export(export.name, kind, export.index);
214                }
215                result.section(&exports);
216            }
217
218            payload => {
219                if let Some((id, range)) = payload.as_section() {
220                    result.section(&RawSection {
221                        id,
222                        data: &module[range],
223                    });
224                }
225            }
226        }
227    }
228
229    Ok((result.finish(), exports_result))
230}
231
232fn add_custom_section(name: &str, data: &[u8], module: &[u8]) -> Result<Vec<u8>> {
233    let mut result = Module::new();
234
235    for payload in Parser::new(0).parse_all(module) {
236        if let Some((id, range)) = payload?.as_section() {
237            result.section(&RawSection {
238                id,
239                data: &module[range],
240            });
241        }
242    }
243
244    result.section(&CustomSection {
245        name: Cow::Borrowed(name),
246        data: Cow::Borrowed(data),
247    });
248
249    Ok(result.finish())
250}
251
252#[cfg(test)]
253mod tests {
254    use std::{path::PathBuf, process};
255
256    use anyhow::Context;
257    use wasmtime_wasi::pipe::MemoryOutputPipe;
258
259    use {
260        super::abi_conformance::{
261            InvocationStyle, KeyValueReport, LlmReport, MysqlReport, PostgresReport, RedisReport,
262            Report, TestConfig, WasiReport,
263        },
264        anyhow::{anyhow, Result},
265        tokio::fs,
266        wasmtime::{
267            component::{Component, Linker},
268            Config, Engine, Store,
269        },
270        wasmtime_wasi::{
271            bindings::Command, pipe::MemoryInputPipe, IoView, ResourceTable, WasiView,
272        },
273        wasmtime_wasi::{WasiCtx, WasiCtxBuilder},
274    };
275
276    async fn run_spin(module: &[u8]) -> Result<()> {
277        let mut config = Config::new();
278        config.wasm_component_model(true);
279        config.async_support(true);
280
281        let engine = Engine::new(&config)?;
282
283        let component = Component::new(
284            &engine,
285            crate::componentize(module).context("could not componentize")?,
286        )
287        .context("failed to instantiate componentized bytes")?;
288
289        let report = super::abi_conformance::test(
290            &component,
291            &engine,
292            TestConfig {
293                invocation_style: InvocationStyle::InboundHttp,
294            },
295        )
296        .await
297        .context("abi conformance test failed")?;
298
299        let expected = Report {
300            inbound_http: Ok(()),
301            inbound_redis: Ok(()),
302            config: Ok(()),
303            http: Ok(()),
304            redis: RedisReport {
305                publish: Ok(()),
306                set: Ok(()),
307                get: Ok(()),
308                incr: Ok(()),
309                del: Ok(()),
310                sadd: Ok(()),
311                srem: Ok(()),
312                smembers: Ok(()),
313                execute: Ok(()),
314            },
315            postgres: PostgresReport {
316                execute: Ok(()),
317                query: Ok(()),
318            },
319            mysql: MysqlReport {
320                execute: Ok(()),
321                query: Ok(()),
322            },
323            key_value: KeyValueReport {
324                open: Ok(()),
325                get: Ok(()),
326                set: Ok(()),
327                delete: Ok(()),
328                exists: Ok(()),
329                get_keys: Ok(()),
330                close: Ok(()),
331            },
332            llm: LlmReport { infer: Ok(()) },
333            wasi: WasiReport {
334                env: Ok(()),
335                epoch: Ok(()),
336                random: Ok(()),
337                stdio: Ok(()),
338                read: Ok(()),
339                readdir: Ok(()),
340                stat: Ok(()),
341            },
342        };
343
344        if report == expected {
345            Ok(())
346        } else {
347            Err(anyhow!("{report:#?}"))
348        }
349    }
350
351    async fn run_command(module: &[u8]) -> Result<()> {
352        let mut config = Config::new();
353        config.wasm_component_model(true);
354        config.async_support(true);
355
356        let engine = Engine::new(&config)?;
357
358        struct Wasi {
359            ctx: WasiCtx,
360            table: ResourceTable,
361        }
362        impl IoView for Wasi {
363            fn table(&mut self) -> &mut ResourceTable {
364                &mut self.table
365            }
366        }
367        impl WasiView for Wasi {
368            fn ctx(&mut self) -> &mut WasiCtx {
369                &mut self.ctx
370            }
371        }
372
373        let mut linker = Linker::<Wasi>::new(&engine);
374        wasmtime_wasi::add_to_linker_async(&mut linker)?;
375        let mut ctx = WasiCtxBuilder::new();
376        let stdout = MemoryOutputPipe::new(1024);
377        ctx.stdin(MemoryInputPipe::new("So rested he by the Tumtum tree"))
378            .stdout(stdout.clone())
379            .args(&["Jabberwocky"]);
380
381        let table = ResourceTable::new();
382        let wasi = Wasi {
383            ctx: ctx.build(),
384            table,
385        };
386
387        let mut store = Store::new(&engine, wasi);
388
389        let component = Component::new(&engine, crate::componentize_command(module)?)?;
390
391        let wasi = Command::instantiate_async(&mut store, &component, &linker).await?;
392
393        wasi.wasi_cli_run()
394            .call_run(&mut store)
395            .await?
396            .map_err(|()| anyhow!("command returned with failing exit status"))?;
397
398        drop(store);
399
400        let stdout = stdout.try_into_inner().unwrap().to_vec();
401
402        assert_eq!(
403            b"Jabberwocky\nSo rested he by the Tumtum tree" as &[_],
404            &stdout
405        );
406
407        Ok(())
408    }
409
410    #[tokio::test]
411    async fn rust_wit_bindgen_02() -> Result<()> {
412        build_rust_test_case("rust-case-0.2");
413        run_spin(
414            &fs::read(concat!(
415                env!("OUT_DIR"),
416                "/wasm32-wasip1/release/rust_case_02.wasm"
417            ))
418            .await?,
419        )
420        .await
421    }
422
423    #[tokio::test]
424    async fn rust_wit_bindgen_08() -> Result<()> {
425        build_rust_test_case("rust-case-0.8");
426        run_spin(
427            &fs::read(concat!(
428                env!("OUT_DIR"),
429                "/wasm32-wasip1/release/rust_case_08.wasm"
430            ))
431            .await?,
432        )
433        .await
434    }
435
436    #[ignore]
437    #[tokio::test]
438    async fn go() -> Result<()> {
439        let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").unwrap());
440        let mut cmd = process::Command::new("tinygo");
441        cmd.arg("build")
442            .current_dir("tests/go-case")
443            .arg("-target=wasip1")
444            .arg("-gc=leaking")
445            .arg("-buildmode=c-shared")
446            .arg("-no-debug")
447            .arg("-o")
448            .arg(out_dir.join("go_case.wasm"))
449            .arg("main.go");
450
451        // If just skip this if TinyGo is not installed
452        _ = cmd.status();
453        run_spin(&fs::read(concat!(env!("OUT_DIR"), "/go_case.wasm")).await?).await
454    }
455
456    #[tokio::test]
457    async fn rust_command() -> Result<()> {
458        build_rust_test_case("rust-command");
459        run_command(
460            &fs::read(concat!(
461                env!("OUT_DIR"),
462                "/wasm32-wasip1/release/rust-command.wasm"
463            ))
464            .await?,
465        )
466        .await
467    }
468
469    fn build_rust_test_case(name: &str) {
470        let out_dir = PathBuf::from(std::env::var_os("OUT_DIR").unwrap());
471        let mut cmd = process::Command::new("cargo");
472        cmd.arg("build")
473            .current_dir(format!("tests/{name}"))
474            .arg("--release")
475            .arg("--target=wasm32-wasip1")
476            .env("CARGO_TARGET_DIR", out_dir);
477
478        let status = cmd.status().unwrap();
479        assert!(status.success());
480    }
481}