fix(cast storage): respect --json for layout (#9332)

* feat(cast storage): allow ugly printing of layout Prior to this change, `cast storage $ADDRESS --rpc-url $RPC_URL --etherscan-api-key $ETHERSCAN_API_KEY` always provided a prettified output. This change adds a `--pretty` flag to `cast storage` which defaults to `true` thus retaining backwards compatibility. Passing `--pretty=false` to `cast storage` results in the json output of the storage layout being produced instead. * fix: remove default value from help text The default value is accessible via `cast storage --help` * fix(cast storage): provide output json path * test(cast): add storage_layout_simple_json test * fix(cast storage): use `--json` flag to ugly print * fix(cast storage): include values in json mode * fix(cast-storage): quiet compilation in all cases * chore: cargo clippy * use fixtures, assert JSON * only quiet if JSON mode, avoid unnecessary warning (if you pass an API key you already expect to fetch remote, very likely default) --------- Co-authored-by: zerosnacks <[email protected]> Co-authored-by: zerosnacks <[email protected]>
foundry-rs · Nov 18, 2024 · d275a49 · d275a49
1 parent 44c86e7
commit d275a49
Show file tree

Hide file tree

Showing 4 changed files with 503 additions and 9 deletions.
diff --git a/crates/cast/bin/cmd/storage.rs b/crates/cast/bin/cmd/storage.rs
@@ -17,6 +17,7 @@ use foundry_common::{
     abi::find_source,
     compile::{etherscan_project, ProjectCompiler},
     ens::NameOrAddress,
+    shell,
 };
 use foundry_compilers::{
     artifacts::{ConfigurableContractArtifact, StorageLayout},
@@ -31,6 +32,7 @@ use foundry_config::{
     impl_figment_convert_cast, Config,
 };
 use semver::Version;
+use serde::{Deserialize, Serialize};
 use std::str::FromStr;
 
 /// The minimum Solc version for outputting storage layouts.
@@ -45,7 +47,7 @@ pub struct StorageArgs {
     #[arg(value_parser = NameOrAddress::from_str)]
     address: NameOrAddress,
 
-    /// The storage slot number.
+    /// The storage slot number. If not provided, it gets the full storage layout.
     #[arg(value_parser = parse_slot)]
     slot: Option<B256>,
 
@@ -109,19 +111,22 @@ impl StorageArgs {
         if project.paths.has_input_files() {
             // Find in artifacts and pretty print
             add_storage_layout_output(&mut project);
-            let out = ProjectCompiler::new().compile(&project)?;
+            let out = ProjectCompiler::new().quiet(shell::is_json()).compile(&project)?;
             let artifact = out.artifacts().find(|(_, artifact)| {
                 artifact.get_deployed_bytecode_bytes().is_some_and(|b| *b == address_code)
             });
             if let Some((_, artifact)) = artifact {
-                return fetch_and_print_storage(provider, address, block, artifact, true).await;
+                return fetch_and_print_storage(
+                    provider,
+                    address,
+                    block,
+                    artifact,
+                    !shell::is_json(),
+                )
+                .await;
             }
         }
 
-        // Not a forge project or artifact not found
-        // Get code from Etherscan
-        sh_warn!("No matching artifacts found, fetching source code from Etherscan...")?;
-
         if !self.etherscan.has_key() {
             eyre::bail!("You must provide an Etherscan API key if you're fetching a remote contract's storage.");
         }
@@ -180,7 +185,7 @@ impl StorageArgs {
         // Clear temp directory
         root.close()?;
 
-        fetch_and_print_storage(provider, address, block, artifact, true).await
+        fetch_and_print_storage(provider, address, block, artifact, !shell::is_json()).await
     }
 }
 
@@ -215,6 +220,14 @@ impl StorageValue {
     }
 }
 
+/// Represents the storage layout of a contract and its values.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+struct StorageReport {
+    #[serde(flatten)]
+    layout: StorageLayout,
+    values: Vec<B256>,
+}
+
 async fn fetch_and_print_storage<P: Provider<T, AnyNetwork>, T: Transport + Clone>(
     provider: P,
     address: Address,
@@ -255,7 +268,22 @@ async fn fetch_storage_slots<P: Provider<T, AnyNetwork>, T: Transport + Clone>(
 
 fn print_storage(layout: StorageLayout, values: Vec<StorageValue>, pretty: bool) -> Result<()> {
     if !pretty {
-        sh_println!("{}", serde_json::to_string_pretty(&serde_json::to_value(layout)?)?)?;
+        let values: Vec<_> = layout
+            .storage
+            .iter()
+            .zip(&values)
+            .map(|(slot, storage_value)| {
+                let storage_type = layout.types.get(&slot.storage_type);
+                storage_value.value(
+                    slot.offset,
+                    storage_type.and_then(|t| t.number_of_bytes.parse::<usize>().ok()),
+                )
+            })
+            .collect();
+        sh_println!(
+            "{}",
+            serde_json::to_string_pretty(&serde_json::to_value(StorageReport { layout, values })?)?
+        )?;
         return Ok(())
     }
 

diff --git a/crates/cast/tests/cli/main.rs b/crates/cast/tests/cli/main.rs
@@ -1131,6 +1131,23 @@ casttest!(storage_layout_simple, |_prj, cmd| {
 "#]]);
 });
 
+// <https://github.com/foundry-rs/foundry/pull/9332>
+casttest!(storage_layout_simple_json, |_prj, cmd| {
+    cmd.args([
+        "storage",
+        "--rpc-url",
+        next_rpc_endpoint(NamedChain::Mainnet).as_str(),
+        "--block",
+        "21034138",
+        "--etherscan-api-key",
+        next_mainnet_etherscan_api_key().as_str(),
+        "0x13b0D85CcB8bf860b6b79AF3029fCA081AE9beF2",
+        "--json",
+    ])
+    .assert_success()
+    .stdout_eq(file!["../fixtures/storage_layout_simple.json": Json]);
+});
+
 // <https://github.com/foundry-rs/foundry/issues/6319>
 casttest!(storage_layout_complex, |_prj, cmd| {
     cmd.args([
@@ -1164,6 +1181,22 @@ casttest!(storage_layout_complex, |_prj, cmd| {
 "#]]);
 });
 
+casttest!(storage_layout_complex_json, |_prj, cmd| {
+    cmd.args([
+        "storage",
+        "--rpc-url",
+        next_rpc_endpoint(NamedChain::Mainnet).as_str(),
+        "--block",
+        "21034138",
+        "--etherscan-api-key",
+        next_mainnet_etherscan_api_key().as_str(),
+        "0xBA12222222228d8Ba445958a75a0704d566BF2C8",
+        "--json",
+    ])
+    .assert_success()
+    .stdout_eq(file!["../fixtures/storage_layout_complex.json": Json]);
+});
+
 casttest!(balance, |_prj, cmd| {
     let rpc = next_http_rpc_endpoint();
     let usdt = "0xdac17f958d2ee523a2206206994597c13d831ec7";