✨ load-elf subcommand

🧹

bin/Cargo.toml

@@ -12,6 +12,9 @@ clap = { version = "4.4.3", features = ["derive"] }
 tracing = "0.1.37"
 tracing-subscriber = "0.3.17"
 anyhow = "1.0.75"
+serde = { version = "1.0.188", features = ["derive"] }
+serde_json = "1.0.107"
+alloy-primitives = "0.3.3"
 
 # Local
 cannon-mipsevm = { path = "../crates/mipsevm" }

bin/src/cannon.rs

@@ -26,7 +26,7 @@ fn main() -> Result<()> {
     init_tracing_subscriber(v)?;
 
     tracing::debug!(target: "cannon-cli", "Dispatching subcommand");
-    subcommand.dispatch();
+    subcommand.dispatch()?;
 
     Ok(())
 }

bin/src/subcommands/load_elf.rs

@@ -0,0 +1,77 @@
+//! The `load-elf` subcommand for the cannon binary
+
+use super::CannonSubcommandDispatcher;
+use anyhow::Result;
+use cannon_mipsevm::{load_elf, patch_go, patch_stack};
+use clap::{builder::PossibleValue, Args, ValueEnum};
+use std::{fs, path::PathBuf};
+
+/// Command line arguments for `cannon load-elf`
+#[derive(Args, Debug)]
+#[command(author, version, about)]
+pub(crate) struct LoadElfArgs {
+    /// The path to the input 32-bit big-endian MIPS ELF file.
+    #[arg(long)]
+    path: PathBuf,
+
+    /// The type of patch to perform on the ELF file.
+    #[arg(long, value_enum)]
+    patch_kind: PatchKind,
+
+    /// The output path to write the JSON state to. State will be dumped to stdout if set to `-`.
+    /// Not written if not provided.
+    #[arg(long)]
+    output: Option<String>,
+}
+
+#[derive(Clone, Debug)]
+enum PatchKind {
+    Go,
+    Stack,
+}
+
+impl ValueEnum for PatchKind {
+    fn value_variants<'a>() -> &'a [Self] {
+        &[Self::Go, Self::Stack]
+    }
+
+    fn to_possible_value(&self) -> Option<PossibleValue> {
+        Some(match self {
+            Self::Go => PossibleValue::new("go"),
+            Self::Stack => PossibleValue::new("stack"),
+        })
+    }
+}
+
+impl CannonSubcommandDispatcher for LoadElfArgs {
+    fn dispatch(&self) -> Result<()> {
+        tracing::info!(target: "cannon-cli::load-elf", "Loading ELF file @ {}", self.path.display());
+        let elf_raw = fs::read(&self.path)?;
+        let mut state = load_elf(&elf_raw)?;
+        tracing::info!(target: "cannon-cli::load-elf", "Loaded ELF file and constructed the State");
+
+        match self.patch_kind {
+            PatchKind::Go => {
+                tracing::info!(target: "cannon-cli::load-elf", "Patching the ELF file with patch type = Go...");
+                patch_go(&elf_raw, &mut state)
+            }
+            PatchKind::Stack => {
+                tracing::info!(target: "cannon-cli::load-elf", "Patching the ELF file with patch type = Stack...");
+                patch_stack(&mut state)
+            }
+        }?;
+
+        // TODO(clabby): Compression
+        if let Some(ref path_str) = self.output {
+            if path_str == "-" {
+                println!("{}", serde_json::to_string(&state)?);
+            } else {
+                fs::write(path_str, serde_json::to_string(&state)?)?;
+            }
+        }
+
+        tracing::info!(target: "cannon-cli::load-elf", "Patched the ELF file and dumped the State successfully.");
+
+        Ok(())
+    }
+}

bin/src/subcommands/mod.rs

@@ -1,28 +1,32 @@
 //! Subcommands for the `cannon` binary
 
+use anyhow::Result;
 use clap::Subcommand;
 
+mod load_elf;
 mod run;
 mod witness;
 
 pub(crate) trait CannonSubcommandDispatcher {
     /// Dispatches the subcommand
-    fn dispatch(&self);
+    fn dispatch(&self) -> Result<()>;
 }
 
 /// The subcommands for the `cannon` binary
 #[derive(Subcommand, Debug)]
 pub(crate) enum CannonSubcommand {
     Run(run::RunArgs),
     Witness(witness::WitnessArgs),
+    LoadElf(load_elf::LoadElfArgs),
 }
 
 impl CannonSubcommandDispatcher for CannonSubcommand {
-    /// Dispatches the subcommand
-    fn dispatch(&self) {
+    /// TODO(clabby): Dynamic dispatch on Box<dyn CannonSubcommandDispatcher>
+    fn dispatch(&self) -> Result<()> {
         match self {
             CannonSubcommand::Run(args) => args.dispatch(),
             CannonSubcommand::Witness(args) => args.dispatch(),
+            CannonSubcommand::LoadElf(args) => args.dispatch(),
         }
     }
 }

bin/src/subcommands/run.rs

@@ -1,8 +1,8 @@
 //! The `run` subcommand for the cannon binary
 
-use clap::Args;
-
 use super::CannonSubcommandDispatcher;
+use anyhow::Result;
+use clap::Args;
 
 /// Command line arguments for `cannon run`
 #[derive(Args, Debug)]
@@ -51,7 +51,7 @@ pub(crate) struct RunArgs {
 }
 
 impl CannonSubcommandDispatcher for RunArgs {
-    fn dispatch(&self) {
+    fn dispatch(&self) -> Result<()> {
         todo!()
     }
 }

bin/src/subcommands/witness.rs

@@ -1,22 +1,51 @@
 //! The `witness` subcommand for the cannon binary
 
-use clap::Args;
-
 use super::CannonSubcommandDispatcher;
+use alloy_primitives::B256;
+use anyhow::Result;
+use cannon_mipsevm::{State, StateWitnessHasher};
+use clap::Args;
+use std::{fs, path::PathBuf};
 
 /// Command line arguments for `cannon witness`
 #[derive(Args, Debug)]
 #[command(author, version, about)]
 pub(crate) struct WitnessArgs {
     /// The path to the input JSON state.
     #[arg(long)]
-    input: String,
+    input: PathBuf,
 
     /// The path to the output JSON state.
     #[arg(long)]
-    output: String,
+    output: Option<PathBuf>,
 }
 
 impl CannonSubcommandDispatcher for WitnessArgs {
-    fn dispatch(&self) {}
+    fn dispatch(&self) -> Result<()> {
+        tracing::info!(target: "cannon-cli::witness", "Loading state JSON dump from {}", self.input.display());
+
+        // TODO(clabby): State compression on disk. Rn we're reading raw.
+        let state_raw = fs::read_to_string(&self.input)?;
+        let mut state: State = serde_json::from_str(&state_raw)?;
+
+        tracing::info!(target: "cannon-cli::witness", "Loaded state JSON dump and deserialized the State");
+
+        let witness = state.encode_witness()?;
+        let witness_hash = witness.state_hash();
+
+        tracing::info!(target: "cannon-cli::witness", "Encoded witness and computed witness hash: {}", B256::from(witness_hash));
+
+        match self.output {
+            Some(ref output_path) => fs::write(output_path, witness).map_err(|_| {
+                anyhow::anyhow!("Failed to write witness to {}", output_path.display())
+            }),
+            None => {
+                println!("{}", B256::from(witness_hash));
+                Ok(())
+            }
+        }?;
+
+        tracing::info!(target: "cannon-cli::witness", "Wrote witness to {}", self.output.as_ref().map_or("stdout".to_string(), |p| p.display().to_string()));
+        Ok(())
+    }
 }