piecrust 0.30.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::fs;
use std::io;
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::{env, 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;

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 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)
    }

    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();

        match Self::from_cache_file(engine, module_path, bytecode) {
            Ok(module) => Ok(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)?;
                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);

        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
    }
}