miden-vm 0.36.0

use std::{fs, path::PathBuf, sync::Arc, time::Instant};

use clap::Parser;
use miden_assembly::{
    Assembler, DefaultSourceManager,
    diagnostics::{IntoDiagnostic, Report, Result, WrapErr},
};
use miden_mast_package::Package;
use miden_vm::{
    ExecutionClaim, KernelDescriptor, ProgramInfo, internal::InputFile, serde::Deserializable,
};

use super::data::{OutputFile, ProgramHash, ProofFile};

#[derive(Debug, Clone, Parser)]
#[command(about = "Verify a Miden program")]
pub struct VerifyCmd {
    /// Path to input file
    #[arg(short = 'i', long = "input", value_parser)]
    input_file: Option<PathBuf>,
    /// Path to output file
    #[arg(short = 'o', long = "output", value_parser)]
    output_file: Option<PathBuf>,
    /// Path to proof file
    #[arg(short = 'p', long = "proof", value_parser)]
    proof_file: PathBuf,
    /// Program hash (hex)
    #[arg(short = 'x', long = "program-hash")]
    program_hash: String,

    /// Path to a file (.masm or .masp) containing the kernel to be loaded with the program
    #[arg(long = "kernel", value_parser)]
    kernel_file: Option<PathBuf>,
}

impl VerifyCmd {
    pub fn execute(&self) -> Result<(), Report> {
        // Validate the kernel file's extension before doing any other file I/O (mirrors the
        // same ordering issue already fixed for `prove` in #3587).
        if let Some(ref kernel_path) = self.kernel_file {
            validate_kernel_extension(kernel_path)?;
        }

        let (input_file, output_file) = self.infer_defaults()?;

        println!("===============================================================================");
        println!("Verifying proof: {}", self.proof_file.display());
        println!("-------------------------------------------------------------------------------");

        // read program hash from input
        let program_hash = ProgramHash::read(&self.program_hash).map_err(Report::msg)?;

        // load input data from file
        let input_data = InputFile::read(&Some(input_file), self.proof_file.as_ref())?;

        // fetch the stack inputs from the arguments
        let stack_inputs = input_data.parse_stack_inputs().map_err(Report::msg)?;

        // load outputs data from file
        let outputs_data =
            OutputFile::read(&Some(output_file), self.proof_file.as_ref()).map_err(Report::msg)?;

        // load proof from file
        let proof = ProofFile::read(&Some(self.proof_file.clone()), self.proof_file.as_ref())
            .map_err(Report::msg)?;

        let now = Instant::now();

        // Load kernel if provided, otherwise use default
        let kernel = if let Some(ref kernel_path) = self.kernel_file {
            if !kernel_path.is_file() {
                return Err(Report::msg(format!(
                    "Kernel file `{}` must be a file.",
                    kernel_path.display()
                )));
            }
            load_kernel_descriptor(kernel_path)?
        } else {
            KernelDescriptor::default()
        };
        let program_info = ProgramInfo::new(program_hash, kernel);

        // verify proof
        let stack_outputs = outputs_data.stack_outputs().map_err(Report::msg)?;
        let claim = ExecutionClaim::from_program_info(program_info, stack_inputs, stack_outputs);
        let outcome = miden_vm::Verifier::new()
            .verify(&claim, &proof)
            .into_diagnostic()
            .wrap_err("Program failed verification!")?;
        if let Some(root) = outcome.outstanding_precompile_root() {
            return Err(Report::msg(format!(
                "Program proof is valid but incomplete; outstanding precompile root: {root}"
            )));
        }
        let vm_security_level = outcome.vm_security_parameters().conjectured_security_level();
        let security_level = outcome.precompile_security_parameters().map_or(
            vm_security_level,
            |precompile_parameters| {
                vm_security_level.min(precompile_parameters.conjectured_security_level())
            },
        );

        println!(
            "Verification complete in {} ms. Security level: {} bits",
            now.elapsed().as_millis(),
            security_level
        );

        Ok(())
    }

    fn infer_defaults(&self) -> Result<(PathBuf, PathBuf), Report> {
        if !self.proof_file.exists() {
            return Err(Report::msg("Proof file does not exist"));
        }
        let default_path = |ext: &str| self.proof_file.with_extension(ext);

        let input_file =
            self.input_file.as_ref().map_or_else(|| default_path("inputs"), PathBuf::clone);
        let output_file = self
            .output_file
            .as_ref()
            .map_or_else(|| default_path("outputs"), PathBuf::clone);

        Ok((input_file, output_file))
    }
}

/// Validates that `kernel_path` has a `.masm` or `.masp` extension, without touching the
/// filesystem.
fn validate_kernel_extension(kernel_path: &std::path::Path) -> Result<(), Report> {
    let ext = kernel_path.extension().and_then(|s| s.to_str()).unwrap_or("").to_lowercase();
    match ext.as_str() {
        "masm" | "masp" => Ok(()),
        _ => Err(Report::msg(format!(
            "Kernel file `{}` must have a .masm or .masp extension",
            kernel_path.display()
        ))),
    }
}

/// Loads a kernel descriptor from a file (.masm or .masp).
///
/// Callers are expected to have already validated the extension via
/// [`validate_kernel_extension`]; this only re-derives it to pick the right loading path.
fn load_kernel_descriptor(kernel_path: &PathBuf) -> Result<KernelDescriptor, Report> {
    // Determine file type based on extension
    let ext = kernel_path.extension().and_then(|s| s.to_str()).unwrap_or("").to_lowercase();

    // Load kernel from .masp package or compile from .masm source
    let kernel_pkg = match ext.as_str() {
        "masp" => {
            // Load kernel from package file
            let bytes = fs::read(kernel_path).into_diagnostic().wrap_err_with(|| {
                format!("Failed to read kernel package `{}`", kernel_path.display())
            })?;
            let package =
                Package::read_from_bytes(&bytes).into_diagnostic().wrap_err_with(|| {
                    format!("Failed to deserialize kernel package `{}`", kernel_path.display())
                })?;
            if !package.is_kernel() {
                return Err(Report::msg(format!(
                    "invalid kernel package, package is of type {}",
                    package.kind,
                )));
            }
            Arc::new(package)
        },
        "masm" => {
            // Compile kernel from assembly source
            let source_manager = Arc::new(DefaultSourceManager::default());
            Assembler::new(source_manager)
                .assemble_kernel_from_root("kernel", kernel_path)
                .map(Arc::<Package>::from)
                .wrap_err_with(|| {
                    format!("Failed to compile kernel from `{}`", kernel_path.display())
                })?
        },
        _ => {
            return Err(Report::msg(format!(
                "Kernel file `{}` must have a .masm or .masp extension",
                kernel_path.display()
            )));
        },
    };

    // Extract the kernel descriptor from the kernel package.
    kernel_pkg.to_kernel_descriptor()
}

#[cfg(test)]
mod tests {
    use std::{fs, fs::File};

    use super::*;

    #[test]
    fn infer_defaults_uses_proof_file_basename_for_defaults() {
        // prepare a unique temp directory
        let base =
            std::env::temp_dir().join(format!("miden_vm_verify_test_{}", std::process::id()));
        fs::create_dir_all(&base).expect("create temp test dir");

        // create a dummy proof file
        let proof_path = base.join("proof_file");
        File::create(&proof_path).expect("create proof file");

        // build command with no explicit input/output
        let cmd = VerifyCmd {
            input_file: None,
            output_file: None,
            proof_file: proof_path.clone(),
            program_hash: "00".to_string(),
            kernel_file: None,
        };

        // exercise
        let (input, output) = cmd.infer_defaults().expect("infer defaults");

        // verify: defaults are proof file with replaced extensions
        assert_eq!(input, proof_path.with_extension("inputs"));
        assert_eq!(output, proof_path.with_extension("outputs"));

        // cleanup best-effort
        let _ = fs::remove_file(&proof_path);
        let _ = fs::remove_dir_all(&base);
    }

    #[test]
    fn execute_rejects_bad_kernel_extension_before_touching_other_files() {
        // None of these paths need to exist: the bad kernel extension must be rejected before
        // any of that file I/O runs.
        let cmd = VerifyCmd {
            input_file: None,
            output_file: None,
            proof_file: PathBuf::from("/nonexistent/does-not-exist.proof"),
            program_hash: "00".to_string(),
            kernel_file: Some(PathBuf::from("kernel.txt")),
        };

        let err = cmd.execute().expect_err("expected the bad kernel extension to be rejected");
        let message = format!("{err}");
        assert!(
            message.contains("must have a .masm or .masp extension"),
            "expected a kernel-extension error, got: {message}"
        );
    }
}