fvm 2.12.0

use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::ops::Deref;
use std::sync::{Arc, Mutex};

use anyhow::{Context, anyhow};
use cid::Cid;
use fvm_ipld_blockstore::Blockstore;
use fvm_wasm_instrument::gas_metering::GAS_COUNTER_NAME;
use fvm_wasm_instrument::parity_wasm::elements;
use wasmtime::OptLevel::Speed;
use wasmtime::{
    Global, GlobalType, Linker, Memory, MemoryType, Module, Mutability, Val, ValType,
    WasmBacktraceDetails,
};

use crate::Kernel;
use crate::gas::WasmGasPrices;
use crate::machine::NetworkConfig;
use crate::syscalls::error::Abort;
use crate::syscalls::{InvocationData, bind_syscalls};

/// A caching wasmtime engine.
#[derive(Clone)]
pub struct Engine(Arc<EngineInner>);

/// Container managing engines with different consensus-affecting configurations.
#[derive(Clone)]
pub struct MultiEngine(Arc<Mutex<HashMap<EngineConfig, Engine>>>);

/// The proper way of getting this struct is to convert from `NetworkConfig`
#[derive(Clone, Eq, PartialEq, Hash)]
pub struct EngineConfig {
    pub max_wasm_stack: u32,
    pub wasm_prices: &'static WasmGasPrices,
    pub actor_redirect: Vec<(Cid, Cid)>,
}

impl From<&NetworkConfig> for EngineConfig {
    fn from(nc: &NetworkConfig) -> Self {
        EngineConfig {
            max_wasm_stack: nc.max_wasm_stack,
            wasm_prices: &nc.price_list.wasm_rules,
            actor_redirect: nc.actor_redirect.clone(),
        }
    }
}

impl MultiEngine {
    pub fn new() -> MultiEngine {
        MultiEngine(Arc::new(Mutex::new(HashMap::new())))
    }

    pub fn get(&self, nc: &NetworkConfig) -> anyhow::Result<Engine> {
        let mut engines = self
            .0
            .lock()
            .map_err(|_| anyhow::Error::msg("multiengine lock is poisoned"))?;

        let ec: EngineConfig = nc.into();

        let engine = match engines.entry(ec.clone()) {
            Occupied(entry) => entry.into_mut(),
            Vacant(entry) => entry.insert(Engine::new_default(ec)?),
        };

        Ok(engine.clone())
    }
}

impl Default for MultiEngine {
    fn default() -> Self {
        Self::new()
    }
}

pub fn default_wasmtime_config() -> wasmtime::Config {
    let mut c = wasmtime::Config::default();

    // Explicitly disable custom page sizes, we always assume 64KiB.
    c.wasm_custom_page_sizes(false);

    // wasmtime default: true
    // We disable this as we always charge for memory regardless and `memory_init_cow` can balloon
    // compiled wasm modules.
    c.memory_init_cow(false);

    // Note: Threads are disabled by default.
    // If we add the "wasmtime/threads" feature in the future,
    // we would explicitly set c.wasm_threads(false) here.

    // wasmtime default: true
    // simd isn't supported in wasm-instrument, but if we add support there, we can probably enable
    // this.
    // Note: stack limits may need adjusting after this is enabled
    c.wasm_simd(false);
    c.wasm_relaxed_simd(false);
    c.relaxed_simd_deterministic(true);

    // wasmtime default: false
    // Supports wide instructions (https://github.com/WebAssembly/wide-arithmetic).
    // We'd like this, but we'll need to (a) update our gas instrumentation logic to support it and
    // (b) make sure our build pipeline can take advantage of it. We should probably wait for it to
    // be enabled by default.
    c.wasm_wide_arithmetic(false);

    // wasmtime default: true
    // We don't support the return_call_* functions.
    c.wasm_tail_call(false);

    // wasmtime default: true
    c.wasm_multi_memory(false);

    // wasmtime default: true
    c.wasm_memory64(false);

    // wasmtime default: true
    // Note: wasm-instrument only supports this at a basic level, for M2 we will
    // need to add more advanced support
    c.wasm_bulk_memory(true);

    // wasmtime default: true
    // we should be able to enable this for M2, just need to make sure that it's
    // handled correctly in wasm-instrument
    c.wasm_multi_value(false);

    // Note: GC is disabled by default.
    // If we add the "wasmtime/gc" feature in the future,
    // we would explicitly set c.wasm_gc(false) and c.wasm_function_references(false) here.

    // wasmtime default: false
    //
    // from wasmtime docs:
    // > When Cranelift is used as a code generation backend this will
    // > configure it to replace NaNs with a single canonical value. This
    // > is useful for users requiring entirely deterministic WebAssembly
    // > computation. This is not required by the WebAssembly spec, so it is
    // > not enabled by default.
    c.cranelift_nan_canonicalization(true);

    // wasmtime default: 512KiB
    // Set to something much higher than the instrumented limiter.
    // Note: This is in bytes, while the instrumented limit is in stack elements
    c.max_wasm_stack(4 << 20);

    // Execution cost accouting is done through wasm instrumentation,
    c.consume_fuel(false);
    c.epoch_interruption(false);

    // Disable debug-related things, wasm-instrument doesn't fix debug info
    // yet, so those aren't useful, just add overhead
    c.debug_info(false);
    c.generate_address_map(false);
    c.cranelift_debug_verifier(false);
    c.native_unwind_info(false);
    c.wasm_backtrace(false);
    c.wasm_backtrace_details(WasmBacktraceDetails::Disable);

    // Reiterate some defaults
    c.guard_before_linear_memory(true);
    c.parallel_compilation(true);

    // Note: Caching is disabled by default.
    // If we add the "wasmtime/cache" feature in the future,
    // we would explicitly set c.disable_cache() here.

    // Note: Async is disabled by default.
    // If we add the "wasmtime/async" feature in the future,
    // we would explicitly set c.async_support(false) here.

    // Doesn't seem to have significant impact on the time it takes to load code
    // todo(M2): make sure this is guaranteed to run in linear time.
    c.cranelift_opt_level(Speed);

    // Use traditional unix signal handlers instead of Mach ports on MacOS. The Mach ports signal
    // handlers don't appear to be capturing all SIGILL signals (when run under lotus) and we're not
    // entirely sure why. Upstream documentation indicates that this could be due to the use of
    // `fork`, but we're only using threads, not subprocesses.
    //
    // The downside to using traditional signal handlers is that this may interfere with some
    // debugging tools. But we'll just have to live with that.
    c.macos_use_mach_ports(false);

    // wasmtime default: true
    // Disable extended const support. We'll probably enable this in the future but that requires a
    // FIP.
    c.wasm_extended_const(false);

    // Note: Component model is disabled by default.
    // If we add the "wasmtime/component-model" feature in the future,
    // we would explicitly set c.wasm_component_model(false) here.

    // wasmtime default: true
    // TODO: Consider disabling this to make performance more deterministic. But we benchmarked with
    // it on, so we leave it on for now.
    // https://github.com/filecoin-project/ref-fvm/issues/2129
    // c.table_lazy_init(false);

    c
}

struct EngineInner {
    engine: wasmtime::Engine,

    /// These two fields are used used in the store constructor to avoid resolve a chicken & egg
    /// situation: We need the store before we can get the real values, but we need to create the
    /// `InvocationData` before we can make the store.
    ///
    /// Alternatively, we could use `Option`s. But then we need to unwrap everywhere.
    dummy_gas_global: Global,
    dummy_memory: Memory,

    module_cache: Mutex<HashMap<Cid, Module>>,
    instance_cache: Mutex<HashMap<TypeId, Box<dyn Any + Send>>>,
    config: EngineConfig,

    actor_redirect: HashMap<Cid, Cid>,
}

impl Deref for Engine {
    type Target = wasmtime::Engine;

    fn deref(&self) -> &Self::Target {
        &self.0.engine
    }
}

impl Engine {
    pub fn new_default(ec: EngineConfig) -> anyhow::Result<Self> {
        Engine::new(&default_wasmtime_config(), ec)
    }

    /// Create a new Engine from a wasmtime config.
    pub fn new(c: &wasmtime::Config, ec: EngineConfig) -> anyhow::Result<Self> {
        let engine = wasmtime::Engine::new(c)?;

        let mut dummy_store = wasmtime::Store::new(&engine, ());
        let gg_type = GlobalType::new(ValType::I64, Mutability::Var);
        let dummy_gg = Global::new(&mut dummy_store, gg_type, Val::I64(0))
            .expect("failed to create dummy gas global");

        let dummy_memory = Memory::new(&mut dummy_store, MemoryType::new(0, Some(0)))
            .expect("failed to create dummy memory");

        let actor_redirect = ec.actor_redirect.iter().cloned().collect();

        Ok(Engine(Arc::new(EngineInner {
            engine,
            dummy_memory,
            dummy_gas_global: dummy_gg,
            module_cache: Default::default(),
            instance_cache: Mutex::new(HashMap::new()),
            config: ec,
            actor_redirect,
        })))
    }
}

struct Cache<K> {
    linker: wasmtime::Linker<InvocationData<K>>,
}

impl Engine {
    /// Instantiates and caches the Wasm modules for the bytecodes addressed by
    /// the supplied CIDs. Only uncached entries are actually fetched and
    /// instantiated. Blockstore failures and entry inexistence shortcircuit
    /// make this method return an Err immediately.
    pub fn preload<'a, BS, I>(&self, blockstore: BS, cids: I) -> anyhow::Result<()>
    where
        BS: Blockstore,
        I: IntoIterator<Item = &'a Cid>,
    {
        let mut cache = self.0.module_cache.lock().expect("module_cache poisoned");
        for cid in cids {
            let cid = self.with_redirect(cid);
            if cache.contains_key(cid) {
                continue;
            }
            let wasm = blockstore.get(cid)?.ok_or_else(|| {
                anyhow!(
                    "no wasm bytecode in blockstore for CID {}",
                    &cid.to_string()
                )
            })?;
            let module = self.load_raw(wasm.as_slice())?;
            cache.insert(*cid, module);
        }
        Ok(())
    }

    fn with_redirect<'a>(&'a self, k: &'a Cid) -> &'a Cid {
        match &self.0.actor_redirect.get(k) {
            Some(cid) => cid,
            None => k,
        }
    }

    /// Load some wasm code into the engine.
    pub fn load_bytecode(&self, k: &Cid, wasm: &[u8]) -> anyhow::Result<Module> {
        let k = self.with_redirect(k);
        let mut cache = self.0.module_cache.lock().expect("module_cache poisoned");
        let module = match cache.get(k) {
            Some(module) => module.clone(),
            None => {
                let module = self.load_raw(wasm)?;
                cache.insert(*k, module.clone());
                module
            }
        };
        Ok(module)
    }

    fn load_raw(&self, raw_wasm: &[u8]) -> anyhow::Result<Module> {
        // First make sure that non-instrumented wasm is valid
        Module::validate(&self.0.engine, raw_wasm)
            .map_err(anyhow::Error::msg)
            .with_context(|| "failed to validate actor wasm")?;

        // Note: when adding debug mode support (with recorded syscall replay) don't instrument to
        // avoid breaking debug info

        use fvm_wasm_instrument::gas_metering::inject;
        use fvm_wasm_instrument::inject_stack_limiter;
        use fvm_wasm_instrument::parity_wasm::deserialize_buffer;

        let m = deserialize_buffer(raw_wasm)?;

        // stack limiter adds post/pre-ambles to call instructions; We want to do that
        // before injecting gas accounting calls to avoid this overhead in every single
        // block of code.
        let m =
            inject_stack_limiter(m, self.0.config.max_wasm_stack).map_err(anyhow::Error::msg)?;

        // inject gas metering based on a price list. This function will
        // * add a new mutable i64 global import, gas.gas_counter
        // * push a gas counter function which deduces gas from the global, and
        //   traps when gas.gas_counter is less than zero
        // * optionally push a function which wraps memory.grow instruction
        //   making it charge gas based on memory requested
        // * divide code into metered blocks, and add a call to the gas counter
        //   function before entering each metered block
        let mut m = inject(m, self.0.config.wasm_prices, "gas")
            .map_err(|_| anyhow::Error::msg("injecting gas counter failed"))?;

        // Work around #602. Remove this once paritytech/parity-wasm#331 is merged and bubbled.
        fix_wasm_sections(&mut m);

        let wasm = m.to_bytes()?;
        let module = Module::from_binary(&self.0.engine, wasm.as_slice())?;

        Ok(module)
    }

    /// Load compiled wasm code into the engine.
    ///
    /// # Safety
    ///
    /// See [`wasmtime::Module::deserialize`] for safety information.
    pub unsafe fn load_compiled(&self, k: &Cid, compiled: &[u8]) -> anyhow::Result<Module> {
        let k = self.with_redirect(k);
        let mut cache = self.0.module_cache.lock().expect("module_cache poisoned");
        let module = match cache.get(k) {
            Some(module) => module.clone(),
            None => {
                let module = unsafe { Module::deserialize(&self.0.engine, compiled)? };
                cache.insert(*k, module.clone());
                module
            }
        };
        Ok(module)
    }

    /// Lookup a loaded wasmtime module.
    pub fn get_module(&self, k: &Cid) -> Option<Module> {
        let k = self.with_redirect(k);
        self.0
            .module_cache
            .lock()
            .expect("module_cache poisoned")
            .get(k)
            .cloned()
    }

    /// Lookup and instantiate a loaded wasmtime module with the given store. This will cache the
    /// linker, syscalls, "pre" isntance, etc.
    pub fn get_instance<K: Kernel>(
        &self,
        store: &mut wasmtime::Store<InvocationData<K>>,
        k: &Cid,
    ) -> anyhow::Result<Option<wasmtime::Instance>> {
        let k = self.with_redirect(k);
        let mut instance_cache = self.0.instance_cache.lock().expect("cache poisoned");

        let type_id = TypeId::of::<K>();
        let cache: &mut Cache<K> = match instance_cache.entry(type_id) {
            Occupied(e) => &mut *e
                .into_mut()
                .downcast_mut()
                .expect("invalid instance cache entry"),
            Vacant(e) => &mut *e
                .insert({
                    let mut linker: Linker<InvocationData<K>> = Linker::new(&self.0.engine);
                    linker.allow_shadowing(true);

                    bind_syscalls(&mut linker)?;
                    Box::new(Cache { linker })
                })
                .downcast_mut()
                .expect("invalid instance cache entry"),
        };
        let gas_global = store.data_mut().avail_gas_global;
        cache
            .linker
            .define(&store, "gas", GAS_COUNTER_NAME, gas_global)
            .context("failed to define gas counter")
            .map_err(Abort::Fatal)?;

        let module_cache = self.0.module_cache.lock().expect("module_cache poisoned");
        let module = match module_cache.get(k) {
            Some(module) => module,
            None => return Ok(None),
        };

        let pre_instance = cache
            .linker
            .instantiate_pre(module)
            .context("failed to link actor module")?;
        let instance = pre_instance.instantiate(&mut *store)?;

        Ok(Some(instance))
    }

    /// Construct a new wasmtime "store" from the given kernel.
    pub fn new_store<K: Kernel>(&self, kernel: K) -> wasmtime::Store<InvocationData<K>> {
        let id = InvocationData {
            kernel,
            last_error: None,
            avail_gas_global: self.0.dummy_gas_global,
            last_milligas_available: 0,
            memory: self.0.dummy_memory,
        };

        let mut store = wasmtime::Store::new(&self.0.engine, id);
        let ggtype = GlobalType::new(ValType::I64, Mutability::Var);
        let gg = Global::new(&mut store, ggtype, Val::I64(0))
            .expect("failed to create available_gas global");
        store.data_mut().avail_gas_global = gg;

        store
    }
}

// Workaround for https://github.com/filecoin-project/ref-fvm/issues/602
//
// This removes the out-of-order data count section, if it exists, and re-inserts it (with the
// correct data-count).
fn fix_wasm_sections(module: &mut elements::Module) {
    module
        .sections_mut()
        .retain(|sec| !matches!(sec, elements::Section::DataCount(_)));
    if let Some(data_count) = module.data_section().map(|data| data.entries().len()) {
        module
            .insert_section(elements::Section::DataCount(data_count as u32))
            .expect("section wasn't deleted");
    }
}