piecrust 0.32.0

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Copyright (c) DUSK NETWORK. All rights reserved.

use std::collections::HashMap;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{LazyLock, RwLock};
use std::{env, fs, io, mem};

use dusk_wasmtime::Engine;

/// WASM object code belonging to a given contract.
#[derive(Debug, Clone)]
pub struct Module {
    module: dusk_wasmtime::Module,
}

const MODULE_CACHE_META_VERSION: u32 = 1;
const META_VERSION_BYTES: usize = mem::size_of::<u32>();
const META_HASH_BYTES: usize = blake3::OUT_LEN;
// Metadata layout:
// [version:u32][bytecode_hash:blake3][module_hash:blake3][runtime_hash:blake3]
const MODULE_CACHE_META_LEN: usize = META_VERSION_BYTES + (META_HASH_BYTES * 3);
const META_BYTECODE_HASH_OFFSET: usize = META_VERSION_BYTES;
const META_MODULE_HASH_OFFSET: usize =
    META_BYTECODE_HASH_OFFSET + META_HASH_BYTES;
const META_RUNTIME_HASH_OFFSET: usize =
    META_MODULE_HASH_OFFSET + META_HASH_BYTES;

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct ModuleCacheKey {
    path: PathBuf,
    bytecode_hash: [u8; META_HASH_BYTES],
}

static MODULE_CACHE: LazyLock<RwLock<HashMap<ModuleCacheKey, Module>>> =
    LazyLock::new(|| RwLock::new(HashMap::new()));

fn check_single_memory(module: &dusk_wasmtime::Module) -> io::Result<()> {
    // Ensure the module only has one memory
    let n_memories = module
        .exports()
        .filter_map(|exp| exp.ty().memory().map(|_| ()))
        .count();
    if n_memories != 1 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "module has {} memories, but only one is allowed",
                n_memories,
            ),
        ));
    }
    Ok(())
}

fn cache_meta_path(module_path: &Path) -> PathBuf {
    module_path.with_extension(format!("{}.meta", super::OBJECTCODE_EXTENSION))
}

fn cache_runtime_fingerprint() -> blake3::Hash {
    let mut hasher = blake3::Hasher::new();
    // A version bump intentionally invalidates old cache artifacts so we never
    // deserialize object code generated by a potentially incompatible runtime.
    hasher.update(env!("CARGO_PKG_NAME").as_bytes());
    hasher.update(env!("CARGO_PKG_VERSION").as_bytes());
    hasher.update(env::consts::ARCH.as_bytes());
    hasher.update(env::consts::OS.as_bytes());
    hasher.finalize()
}

fn module_cache_key(module_path: &Path, bytecode: &[u8]) -> ModuleCacheKey {
    ModuleCacheKey {
        path: module_path.to_path_buf(),
        bytecode_hash: *blake3::hash(bytecode).as_bytes(),
    }
}

fn insert_cached_module(module_path: &Path, bytecode: &[u8], module: &Module) {
    MODULE_CACHE
        .write()
        .unwrap_or_else(|e| e.into_inner())
        .insert(module_cache_key(module_path, bytecode), module.clone());
}

fn invalidate_cached_modules(module_path: &Path) {
    MODULE_CACHE
        .write()
        .unwrap_or_else(|e| e.into_inner())
        .retain(|key, _| key.path != module_path);
}

fn write_cache_meta(
    module_path: &Path,
    module_bytes: &[u8],
    bytecode: &[u8],
) -> io::Result<()> {
    let mut meta = Vec::with_capacity(MODULE_CACHE_META_LEN);
    meta.extend_from_slice(&MODULE_CACHE_META_VERSION.to_le_bytes());
    meta.extend_from_slice(blake3::hash(bytecode).as_bytes());
    meta.extend_from_slice(blake3::hash(module_bytes).as_bytes());
    meta.extend_from_slice(cache_runtime_fingerprint().as_bytes());
    fs::write(cache_meta_path(module_path), meta)
}

fn validate_cache_meta(
    module_path: &Path,
    module_bytes: &[u8],
    bytecode: &[u8],
) -> io::Result<()> {
    let meta_path = cache_meta_path(module_path);
    let meta = fs::read(&meta_path).map_err(|err| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "failed to read module cache metadata {meta_path:?}: {err}"
            ),
        )
    })?;

    if meta.len() != MODULE_CACHE_META_LEN {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "invalid module cache metadata length for {meta_path:?}: expected {MODULE_CACHE_META_LEN}, got {}",
                meta.len()
            ),
        ));
    }

    let version = u32::from_le_bytes(
        meta[..META_VERSION_BYTES].try_into().expect("slice length"),
    );
    if version != MODULE_CACHE_META_VERSION {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "unsupported module cache metadata version for {meta_path:?}: expected {MODULE_CACHE_META_VERSION}, got {version}",
            ),
        ));
    }

    let expected_bytecode_hash = &meta[META_BYTECODE_HASH_OFFSET
        ..(META_BYTECODE_HASH_OFFSET + META_HASH_BYTES)];
    let expected_module_hash = &meta
        [META_MODULE_HASH_OFFSET..(META_MODULE_HASH_OFFSET + META_HASH_BYTES)];
    let expected_runtime_hash = &meta[META_RUNTIME_HASH_OFFSET
        ..(META_RUNTIME_HASH_OFFSET + META_HASH_BYTES)];

    let bytecode_hash = blake3::hash(bytecode);
    if expected_bytecode_hash != bytecode_hash.as_bytes() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "module cache metadata bytecode hash mismatch for {module_path:?}"
            ),
        ));
    }

    let module_hash = blake3::hash(module_bytes);
    if expected_module_hash != module_hash.as_bytes() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "module cache metadata module hash mismatch for {module_path:?}"
            ),
        ));
    }

    let runtime_fingerprint = cache_runtime_fingerprint();
    if expected_runtime_hash != runtime_fingerprint.as_bytes() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "module cache metadata runtime hash mismatch for {module_path:?}"
            ),
        ));
    }

    Ok(())
}

impl Module {
    pub(crate) fn new<B: AsRef<[u8]>>(
        engine: &Engine,
        bytes: B,
    ) -> io::Result<Self> {
        let module = unsafe {
            dusk_wasmtime::Module::deserialize(engine, bytes).map_err(|e| {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("failed to deserialize module: {}", e),
                )
            })?
        };

        check_single_memory(&module)?;

        Ok(Self { module })
    }

    pub(crate) fn from_cache_file<P: AsRef<Path>>(
        engine: &Engine,
        path: P,
        bytecode: &[u8],
    ) -> io::Result<Self> {
        let path = path.as_ref();
        let module_bytes = fs::read(path).map_err(|err| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("failed to read module cache {path:?}: {err}"),
            )
        })?;

        validate_cache_meta(path, &module_bytes, bytecode)?;
        Self::new(engine, module_bytes)
    }

    pub(crate) fn from_bytecode(
        engine: &Engine,
        bytecode: &[u8],
    ) -> io::Result<Self> {
        let module =
            dusk_wasmtime::Module::new(engine, bytecode).map_err(|e| {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("failed to compile module: {}", e),
                )
            })?;

        check_single_memory(&module)?;

        Ok(Self { module })
    }

    pub(crate) fn serialize(&self) -> Vec<u8> {
        self.module
            .serialize()
            .expect("We don't use WASM components")
    }

    pub(crate) fn write_module_data<P: AsRef<Path>>(
        &self,
        module_path: P,
        bytecode: &[u8],
    ) -> io::Result<()> {
        let module_path = module_path.as_ref();
        let module_bytes = self.serialize();
        fs::write(module_path, &module_bytes)?;
        write_cache_meta(module_path, &module_bytes, bytecode)?;
        insert_cached_module(module_path, bytecode, self);
        Ok(())
    }

    pub(crate) fn load<P: AsRef<Path>>(
        engine: &Engine,
        module_path: P,
        bytecode: &[u8],
    ) -> io::Result<Self> {
        let module_path = module_path.as_ref();
        let cache_key = module_cache_key(module_path, bytecode);

        let cached_module = {
            let cache = MODULE_CACHE.read().unwrap_or_else(|e| e.into_inner());
            cache.get(&cache_key).cloned()
        };

        if let Some(module) = cached_module {
            let mut module_bytes = None;
            validate_cache_meta(
                module_path,
                module_bytes.get_or_insert_with(|| module.serialize()),
                bytecode,
            )?;
            return Ok(module);
        }

        let module = Self::from_cache_file(engine, module_path, bytecode)?;
        insert_cached_module(module_path, bytecode, &module);
        Ok(module)
    }

    pub(crate) fn load_or_recompile<P: AsRef<Path>>(
        engine: &Engine,
        module_path: P,
        bytecode: &[u8],
    ) -> io::Result<Self> {
        let module_path = module_path.as_ref();
        let cache_key = module_cache_key(module_path, bytecode);

        let cached_module = {
            let cache = MODULE_CACHE.read().unwrap_or_else(|e| e.into_inner());
            cache.get(&cache_key).cloned()
        };

        let module = match cached_module {
            Some(module) => {
                let mut module_bytes = None;
                let cache_needs_repair = !module_path.exists()
                    || validate_cache_meta(
                        module_path,
                        module_bytes.get_or_insert_with(|| module.serialize()),
                        bytecode,
                    )
                    .is_err();

                if cache_needs_repair {
                    module.write_module_data(module_path, bytecode)?;
                }

                module
            }
            None => match Self::load(engine, module_path, bytecode) {
                Ok(module) => module,
                Err(err) => {
                    tracing::warn!(
                        "module cache {module_path:?} failed validation; recompiling from bytecode: {err}",
                    );
                    let module = Self::from_bytecode(engine, bytecode)?;
                    module.write_module_data(module_path, bytecode)?;
                    module
                }
            },
        };

        Ok(module)
    }

    pub(crate) fn remove_cache_files<P: AsRef<Path>>(
        module_path: P,
    ) -> io::Result<()> {
        let module_path = module_path.as_ref();
        let meta_path = cache_meta_path(module_path);

        invalidate_cached_modules(module_path);

        if module_path.exists() {
            fs::remove_file(module_path)?;
        }

        if meta_path.exists() {
            fs::remove_file(meta_path)?;
        }

        Ok(())
    }

    pub(crate) fn is_64(&self) -> bool {
        self.module
            .exports()
            .filter_map(|exp| exp.ty().memory().map(|mem_ty| mem_ty.is_64()))
            .next()
            .expect("We guarantee the module has one memory")
    }
}

impl Deref for Module {
    type Target = dusk_wasmtime::Module;

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