feat(cheats): resolveEnv + getChainId + writeXXXUpsert (#11414)

Author: 0xrusowsky

Cargo.toml

@@ -341,6 +341,7 @@ strum = "0.27"
 tempfile = "3.20"
 tokio = "1"
 toml = "0.9"
+toml_edit = "0.23"
 tower = "0.5"
 tower-http = "0.6"
 tracing = "0.1"

crates/cheatcodes/assets/cheatcodes.json

@@ -463,6 +463,13 @@ interface Vm {
 
     // -------- Block and Transaction Properties --------
 
+    /// Gets the current `block.chainid` of the currently selected environment.
+    /// You should use this instead of `block.chainid` if you use `vm.selectFork` or `vm.createSelectFork`, as `block.chainid` could be assumed
+    /// to be constant across a transaction, and as a result will get optimized out by the compiler.
+    /// See https://github.com/foundry-rs/foundry/issues/6180
+    #[cheatcode(group = Evm, safety = Safe)]
+    function getChainId() external view returns (uint256 blockChainId);
+
     /// Sets `block.chainid`.
     #[cheatcode(group = Evm, safety = Unsafe)]
     function chainId(uint256 newChainId) external;
@@ -2027,6 +2034,10 @@ interface Vm {
 
     // ======== Environment Variables ========
 
+    /// Resolves the env variable placeholders of a given input string.
+    #[cheatcode(group = Environment)]
+    function resolveEnv(string calldata input) external returns (string memory);
+
     /// Sets environment variables.
     #[cheatcode(group = Environment)]
     function setEnv(string calldata name, string calldata value) external;
@@ -2542,6 +2553,12 @@ interface Vm {
     #[cheatcode(group = Json)]
     function writeJson(string calldata json, string calldata path, string calldata valueKey) external;
 
+    /// Write a serialized JSON object to an **existing** JSON file, replacing a value with key = <value_key.>
+    /// This is useful to replace a specific value of a JSON file, without having to parse the entire thing.
+    /// Unlike `writeJson`, this cheatcode will create new keys if they didn't previously exist.
+    #[cheatcode(group = Toml)]
+    function writeJsonUpsert(string calldata json, string calldata path, string calldata valueKey) external;
+
     // ======== TOML Parsing and Manipulation ========
 
     // -------- Reading --------
@@ -2647,6 +2664,12 @@ interface Vm {
     #[cheatcode(group = Toml)]
     function writeToml(string calldata json, string calldata path, string calldata valueKey) external;
 
+    /// Takes serialized JSON, converts to TOML and write a serialized TOML table to an **existing** TOML file, replacing a value with key = <value_key.>
+    /// This is useful to replace a specific value of a TOML file, without having to parse the entire thing.
+    /// Unlike `writeToml`, this cheatcode will create new keys if they didn't previously exist.
+    #[cheatcode(group = Toml)]
+    function writeTomlUpsert(string calldata json, string calldata path, string calldata valueKey) external;
+
     // ======== Cryptography ========
 
     // -------- Key Management --------

crates/cheatcodes/spec/src/vm.rs

@@ -6,7 +6,7 @@ use alloy_sol_types::SolValue;
 use std::{env, sync::OnceLock};
 
 /// Stores the forge execution context for the duration of the program.
-static FORGE_CONTEXT: OnceLock<ForgeContext> = OnceLock::new();
+pub static FORGE_CONTEXT: OnceLock<ForgeContext> = OnceLock::new();
 
 impl Cheatcode for setEnvCall {
     fn apply(&self, _state: &mut Cheatcodes) -> Result {
@@ -28,6 +28,15 @@ impl Cheatcode for setEnvCall {
     }
 }
 
+impl Cheatcode for resolveEnvCall {
+    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+        let Self { input } = self;
+        let resolved = foundry_config::resolve::interpolate(input)
+            .map_err(|e| fmt_err!("failed to resolve env var: {e}"))?;
+        Ok(resolved.abi_encode())
+    }
+}
+
 impl Cheatcode for envExistsCall {
     fn apply(&self, _state: &mut Cheatcodes) -> Result {
         let Self { name } = self;

crates/cheatcodes/src/env.rs

@@ -469,6 +469,13 @@ impl Cheatcode for lastCallGasCall {
     }
 }
 
+impl Cheatcode for getChainIdCall {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self {} = self;
+        Ok(U256::from(ccx.ecx.cfg.chain_id).abi_encode())
+    }
+}
+
 impl Cheatcode for chainIdCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { newChainId } = self;

crates/cheatcodes/src/evm.rs

@@ -361,6 +361,23 @@ impl Cheatcode for writeJson_1Call {
     }
 }
 
+impl Cheatcode for writeJsonUpsertCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { json: value, path, valueKey } = self;
+
+        // Read, parse, and update the JSON object
+        let data_path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
+        let data_string = fs::read_to_string(&data_path)?;
+        let mut data =
+            serde_json::from_str(&data_string).unwrap_or_else(|_| Value::String(data_string));
+        upsert_json_value(&mut data, value, valueKey)?;
+
+        // Write the updated content back to the file
+        let json_string = serde_json::to_string_pretty(&data)?;
+        super::fs::write_file(state, path.as_ref(), json_string.as_bytes())
+    }
+}
+
 pub(super) fn check_json_key_exists(json: &str, key: &str) -> Result {
     let json = parse_json_str(json)?;
     let values = select(&json, key)?;
@@ -643,11 +660,68 @@ pub(super) fn resolve_type(type_description: &str) -> Result<DynSolType> {
     bail!("type description should be a valid Solidity type or a EIP712 `encodeType` string")
 }
 
+/// Upserts a value into a JSON object based on a dot-separated key.
+///
+/// This function navigates through a mutable `serde_json::Value` object using a
+/// path-like key. It creates nested JSON objects if they do not exist along the path.
+/// The value is inserted at the final key in the path.
+///
+/// # Arguments
+///
+/// * `data` - A mutable reference to the `serde_json::Value` to be modified.
+/// * `value` - The string representation of the value to upsert. This string is first parsed as
+///   JSON, and if that fails, it's treated as a plain JSON string.
+/// * `key` - A dot-separated string representing the path to the location for upserting.
+pub(super) fn upsert_json_value(data: &mut Value, value: &str, key: &str) -> Result<()> {
+    // Parse the path key into segments.
+    let canonical_key = canonicalize_json_path(key);
+    let parts: Vec<&str> = canonical_key
+        .strip_prefix("$.")
+        .unwrap_or(key)
+        .split('.')
+        .filter(|s| !s.is_empty())
+        .collect();
+
+    if parts.is_empty() {
+        return Err(fmt_err!("'valueKey' cannot be empty or just '$'"));
+    }
+
+    // Separate the final key from the path.
+    // Traverse the objects, creating intermediary ones if necessary.
+    if let Some((key_to_insert, path_to_parent)) = parts.split_last() {
+        let mut current_level = data;
+
+        for segment in path_to_parent {
+            if !current_level.is_object() {
+                return Err(fmt_err!("path segment '{segment}' does not resolve to an object."));
+            }
+            current_level = current_level
+                .as_object_mut()
+                .unwrap()
+                .entry(segment.to_string())
+                .or_insert(Value::Object(Map::new()));
+        }
+
+        // Upsert the new value
+        if let Some(parent_obj) = current_level.as_object_mut() {
+            parent_obj.insert(
+                key_to_insert.to_string(),
+                serde_json::from_str(value).unwrap_or_else(|_| Value::String(value.to_owned())),
+            );
+        } else {
+            return Err(fmt_err!("final destination is not an object, cannot insert key."));
+        }
+    }
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use alloy_primitives::FixedBytes;
     use proptest::strategy::Strategy;
+    use serde_json::json;
 
     fn contains_tuple(value: &DynSolValue) -> bool {
         match value {
@@ -717,4 +791,51 @@ mod tests {
             assert_eq!(value, v);
         }
     }
+
+    #[test]
+    fn test_upsert_json_value() {
+        // Tuples of: (initial_json, key, value_to_upsert, expected)
+        let scenarios = vec![
+            // Simple key-value insert with a plain string
+            (json!({}), "foo", r#""bar""#, json!({"foo": "bar"})),
+            // Overwrite existing value with a number
+            (json!({"foo": "bar"}), "foo", "123", json!({"foo": 123})),
+            // Create nested objects
+            (json!({}), "a.b.c", r#""baz""#, json!({"a": {"b": {"c": "baz"}}})),
+            // Upsert into existing nested object with a boolean
+            (json!({"a": {"b": {}}}), "a.b.c", "true", json!({"a": {"b": {"c": true}}})),
+            // Upsert a JSON object as a value
+            (json!({}), "a.b", r#"{"d": "e"}"#, json!({"a": {"b": {"d": "e"}}})),
+            // Upsert a JSON array as a value
+            (json!({}), "myArray", r#"[1, "test", null]"#, json!({"myArray": [1, "test", null]})),
+        ];
+
+        for (mut initial, key, value_str, expected) in scenarios {
+            upsert_json_value(&mut initial, value_str, key).unwrap();
+            assert_eq!(initial, expected);
+        }
+
+        let error_scenarios = vec![
+            // Path traverses a non-object value
+            (
+                json!({"a": "a string value"}),
+                "a.b",
+                r#""bar""#,
+                "final destination is not an object, cannot insert key.",
+            ),
+            // Empty key should fail
+            (json!({}), "", r#""bar""#, "'valueKey' cannot be empty or just '$'"),
+            // Root path with a trailing dot should fail
+            (json!({}), "$.", r#""bar""#, "'valueKey' cannot be empty or just '$'"),
+        ];
+
+        for (mut initial, key, value_str, error_msg) in error_scenarios {
+            let result = upsert_json_value(&mut initial, value_str, key);
+            assert!(result.is_err(), "Expected an error for key: '{key}' but got Ok");
+            assert!(
+                result.unwrap_err().to_string().contains(error_msg),
+                "Error message for key '{key}' did not contain '{error_msg}'"
+            );
+        }
+    }
 }

crates/cheatcodes/src/json.rs

@@ -5,7 +5,7 @@ use crate::{
     Vm::*,
     json::{
         canonicalize_json_path, check_json_key_exists, parse_json, parse_json_coerce,
-        parse_json_keys, resolve_type,
+        parse_json_keys, resolve_type, upsert_json_value,
     },
 };
 use alloy_dyn_abi::DynSolType;
@@ -194,6 +194,25 @@ impl Cheatcode for writeToml_1Call {
     }
 }
 
+impl Cheatcode for writeTomlUpsertCall {
+    fn apply(&self, state: &mut Cheatcodes) -> Result {
+        let Self { json: value, path, valueKey } = self;
+
+        // Read and parse the TOML file
+        let data_path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
+        let toml_data = fs::read_to_string(&data_path)?;
+
+        // Convert to JSON and update the object
+        let mut json_data: JsonValue =
+            toml::from_str(&toml_data).map_err(|e| fmt_err!("failed parsing TOML: {e}"))?;
+        upsert_json_value(&mut json_data, value, valueKey)?;
+
+        // Serialize back to TOML and write the updated content back to the file
+        let toml_string = format_json_to_toml(json_data)?;
+        super::fs::write_file(state, path.as_ref(), toml_string.as_bytes())
+    }
+}
+
 /// Parse
 fn parse_toml_str(toml: &str) -> Result<TomlValue> {
     toml::from_str(toml).map_err(|e| fmt_err!("failed parsing TOML: {e}"))

crates/cheatcodes/src/toml.rs

@@ -43,7 +43,7 @@ serde.workspace = true
 soldeer-core.workspace = true
 thiserror.workspace = true
 toml = { workspace = true, features = ["preserve_order"] }
-toml_edit = "0.23"
+toml_edit.workspace = true
 tracing.workspace = true
 walkdir.workspace = true
 yansi.workspace = true