build.rs

fn main() -> wasmtime::Result<()> {
    component::generate_static_api_tests()?;

    Ok(())
}

mod component {
    use arbitrary::Unstructured;
    use proc_macro2::TokenStream;
    use quote::quote;
    use rand::rngs::StdRng;
    use rand::{Rng, SeedableRng};
    use std::collections::HashMap;
    use std::env;
    use std::fmt::Write;
    use std::fs;
    use std::iter;
    use std::path::PathBuf;
    use std::process::Command;
    use wasmtime::{Error, Result, error::Context as _, format_err};
    use wasmtime_test_util::component_fuzz::{Declarations, MAX_TYPE_DEPTH, TestCase, Type};

    pub fn generate_static_api_tests() -> Result<()> {
        println!("cargo:rerun-if-changed=build.rs");
        let out_dir = PathBuf::from(
            env::var_os("OUT_DIR").expect("The OUT_DIR environment variable must be set"),
        );

        let mut out = String::new();
        write_static_api_tests(&mut out)?;

        let output = out_dir.join("static_component_api.rs");
        fs::write(&output, out)?;

        drop(Command::new("rustfmt").arg(&output).status());

        Ok(())
    }

    fn write_static_api_tests(out: &mut String) -> Result<()> {
        println!("cargo:rerun-if-env-changed=WASMTIME_FUZZ_SEED");
        let seed = if let Ok(seed) = env::var("WASMTIME_FUZZ_SEED") {
            seed.parse::<u64>()
                .with_context(|| format_err!("expected u64 in WASMTIME_FUZZ_SEED"))?
        } else {
            StdRng::from_entropy().r#gen()
        };

        eprintln!(
            "using seed {seed} (set WASMTIME_FUZZ_SEED={seed} in your environment to reproduce)"
        );

        let mut rng = StdRng::seed_from_u64(seed);

        const TYPE_COUNT: usize = 50;
        const TEST_CASE_COUNT: usize = 100;

        let mut type_fuel = 1000;
        let mut types = Vec::new();
        let mut rust_type_names = Vec::new();
        let name_counter = &mut 0;
        let mut declarations = TokenStream::new();
        let mut tests = TokenStream::new();

        // First generate a set of type to select from.
        for _ in 0..TYPE_COUNT {
            let ty = generate(&mut rng, |u| {
                // Only discount fuel if the generation was successful,
                // otherwise we'll get more random data and try again.
                let mut fuel = type_fuel;
                let ret = Type::generate(u, MAX_TYPE_DEPTH, &mut fuel);
                if ret.is_ok() {
                    type_fuel = fuel;
                }
                ret
            })?;

            let rust_ty_name =
                wasmtime_test_util::component_fuzz::rust_type(&ty, name_counter, &mut declarations);
            types.push(ty);
            rust_type_names.push(rust_ty_name);
        }

        fn hash_key(ty: &Type) -> usize {
            let ty: *const Type = ty;
            ty.addr()
        }

        let type_to_name_map = types
            .iter()
            .map(hash_key)
            .zip(rust_type_names.iter().cloned())
            .collect::<HashMap<_, _>>();

        // Next generate a set of static API test cases driven by the above
        // types.
        for index in 0..TEST_CASE_COUNT {
            let (case, rust_params, rust_results) = generate(&mut rng, |u| {
                let mut rust_params = TokenStream::new();
                let mut rust_results = TokenStream::new();
                let case = TestCase::generate(&types, u)?;
                for ty in case.params.iter() {
                    let name = &type_to_name_map[&hash_key(ty)];
                    rust_params.extend(name.clone());
                    rust_params.extend(quote!(,));
                }
                if let Some(ty) = &case.result {
                    let name = &type_to_name_map[&hash_key(ty)];
                    rust_results.extend(name.clone());
                    rust_results.extend(quote!(,));
                }
                Ok((case, rust_params, rust_results))
            })?;

            let Declarations {
                types,
                type_instantiation_args,
                params,
                results,
                caller_module,
                callee_module,
                options,
            } = case.declarations();

            let test = quote!(#index => component_api::static_api_test::<(#rust_params), (#rust_results)>(
                input,
                {
                    static DECLS: Declarations = Declarations {
                        types: Cow::Borrowed(#types),
                        type_instantiation_args: Cow::Borrowed(#type_instantiation_args),
                        params: Cow::Borrowed(#params),
                        results: Cow::Borrowed(#results),
                        caller_module: Cow::Borrowed(#caller_module),
                        callee_module: Cow::Borrowed(#callee_module),
                        options: #options,
                    };
                    &DECLS
                }
            ),);

            tests.extend(test);
        }

        let module = quote! {
            #[allow(unused_imports, reason = "macro-generated code")]
            fn static_component_api_target(input: &mut libfuzzer_sys::arbitrary::Unstructured) -> libfuzzer_sys::arbitrary::Result<()> {
                use wasmtime::Result;
                use wasmtime_test_util::component_fuzz::Declarations;
                use wasmtime_test_util::component::{Float32, Float64};
                use libfuzzer_sys::arbitrary::{self, Arbitrary};
                use std::borrow::Cow;
                use std::sync::{Arc, Once};
                use wasmtime::component::{ComponentType, Lift, Lower};
                use wasmtime_fuzzing::oracles::component_api;

                const SEED: u64 = #seed;

                static ONCE: Once = Once::new();

                ONCE.call_once(|| {
                    eprintln!(
                        "Seed {SEED} was used to generate static component API fuzz tests.\n\
                         Set WASMTIME_FUZZ_SEED={SEED} in your environment at build time to reproduce."
                    );
                });

                #declarations

                match input.int_in_range(0..=(#TEST_CASE_COUNT-1))? {
                    #tests
                    _ => unreachable!()
                }
            }
        };

        write!(out, "{module}")?;

        Ok(())
    }

    fn generate<T>(
        rng: &mut StdRng,
        mut f: impl FnMut(&mut Unstructured<'_>) -> arbitrary::Result<T>,
    ) -> Result<T> {
        let mut bytes = Vec::new();
        loop {
            let count = rng.gen_range(1000..2000);
            bytes.extend(iter::repeat_with(|| rng.r#gen::<u8>()).take(count));

            match f(&mut Unstructured::new(&bytes)) {
                Ok(ret) => break Ok(ret),
                Err(arbitrary::Error::NotEnoughData) => (),
                Err(error) => break Err(Error::from(error)),
            }
        }
    }
}