foundry_cheatcodes/
toml.rs

1//! Implementations of [`Toml`](spec::Group::Toml) cheatcodes.
2
3use crate::{
4    Cheatcode, CheatcodesCore, Result,
5    Vm::*,
6    json::{
7        check_json_key_exists, parse_json, parse_json_coerce, parse_json_coerce_default,
8        parse_json_keys, resolve_type, split_value_key,
9    },
10};
11use alloy_dyn_abi::DynSolType;
12use alloy_sol_types::SolValue;
13use foundry_common::{fmt::StructDefinitions, fs};
14use foundry_config::fs_permissions::FsAccessKind;
15use serde_json::Value as JsonValue;
16use toml::Value as TomlValue;
17use toml_edit::{DocumentMut, Item, Table, TableLike};
18
19impl Cheatcode for keyExistsTomlCall {
20    fn apply_context_free(&self) -> Result {
21        let Self { toml, key } = self;
22        check_json_key_exists(&toml_to_json_string(toml)?, key)
23    }
24}
25
26impl Cheatcode for parseToml_0Call {
27    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
28        let Self { toml } = self;
29        parse_toml(
30            toml,
31            "$",
32            state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
33        )
34    }
35}
36
37impl Cheatcode for parseToml_1Call {
38    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
39        let Self { toml, key } = self;
40        parse_toml(
41            toml,
42            key,
43            state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
44        )
45    }
46}
47
48macro_rules! impl_parse_toml {
49    ($call:ident, $call_with_default:ident, $ty:expr) => {
50        impl Cheatcode for $call {
51            fn apply_context_free(&self) -> Result {
52                let Self { toml, key } = self;
53                parse_toml_coerce(toml, key, &$ty)
54            }
55        }
56
57        impl Cheatcode for $call_with_default {
58            fn apply_context_free(&self) -> Result {
59                let Self { toml, key, defaultValue } = self;
60                parse_toml_coerce_default(toml, key, &$ty, defaultValue)
61            }
62        }
63    };
64}
65
66impl_parse_toml!(parseTomlUint_0Call, parseTomlUint_1Call, DynSolType::Uint(256));
67impl_parse_toml!(
68    parseTomlUintArray_0Call,
69    parseTomlUintArray_1Call,
70    DynSolType::Array(Box::new(DynSolType::Uint(256)))
71);
72impl_parse_toml!(parseTomlInt_0Call, parseTomlInt_1Call, DynSolType::Int(256));
73impl_parse_toml!(
74    parseTomlIntArray_0Call,
75    parseTomlIntArray_1Call,
76    DynSolType::Array(Box::new(DynSolType::Int(256)))
77);
78impl_parse_toml!(parseTomlBool_0Call, parseTomlBool_1Call, DynSolType::Bool);
79impl_parse_toml!(
80    parseTomlBoolArray_0Call,
81    parseTomlBoolArray_1Call,
82    DynSolType::Array(Box::new(DynSolType::Bool))
83);
84impl_parse_toml!(parseTomlAddress_0Call, parseTomlAddress_1Call, DynSolType::Address);
85impl_parse_toml!(
86    parseTomlAddressArray_0Call,
87    parseTomlAddressArray_1Call,
88    DynSolType::Array(Box::new(DynSolType::Address))
89);
90impl_parse_toml!(parseTomlString_0Call, parseTomlString_1Call, DynSolType::String);
91impl_parse_toml!(
92    parseTomlStringArray_0Call,
93    parseTomlStringArray_1Call,
94    DynSolType::Array(Box::new(DynSolType::String))
95);
96impl_parse_toml!(parseTomlBytes_0Call, parseTomlBytes_1Call, DynSolType::Bytes);
97impl_parse_toml!(
98    parseTomlBytesArray_0Call,
99    parseTomlBytesArray_1Call,
100    DynSolType::Array(Box::new(DynSolType::Bytes))
101);
102impl_parse_toml!(parseTomlBytes32_0Call, parseTomlBytes32_1Call, DynSolType::FixedBytes(32));
103impl_parse_toml!(
104    parseTomlBytes32Array_0Call,
105    parseTomlBytes32Array_1Call,
106    DynSolType::Array(Box::new(DynSolType::FixedBytes(32)))
107);
108
109impl Cheatcode for parseTomlType_0Call {
110    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
111        let Self { toml, typeDescription } = self;
112        parse_toml_coerce(
113            toml,
114            "$",
115            &resolve_type(
116                typeDescription,
117                state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
118            )?,
119        )
120        .map(|v| v.abi_encode())
121    }
122}
123
124impl Cheatcode for parseTomlType_1Call {
125    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
126        let Self { toml, key, typeDescription } = self;
127        parse_toml_coerce(
128            toml,
129            key,
130            &resolve_type(
131                typeDescription,
132                state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
133            )?,
134        )
135        .map(|v| v.abi_encode())
136    }
137}
138
139impl Cheatcode for parseTomlTypeArrayCall {
140    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
141        let Self { toml, key, typeDescription } = self;
142        let ty = resolve_type(
143            typeDescription,
144            state.analysis.as_ref().and_then(|analysis| analysis.struct_defs().ok()),
145        )?;
146        parse_toml_coerce(toml, key, &DynSolType::Array(Box::new(ty))).map(|v| v.abi_encode())
147    }
148}
149
150impl Cheatcode for parseTomlKeysCall {
151    fn apply_context_free(&self) -> Result {
152        let Self { toml, key } = self;
153        parse_toml_keys(toml, key)
154    }
155}
156
157impl Cheatcode for writeToml_0Call {
158    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
159        let Self { json, path } = self;
160        let value =
161            serde_json::from_str(json).unwrap_or_else(|_| JsonValue::String(json.to_owned()));
162
163        let toml_string = format_json_to_toml(value)?;
164        super::fs::write_file(state, path.as_ref(), toml_string.as_bytes())
165    }
166}
167
168impl Cheatcode for writeToml_1Call {
169    fn apply_context(&self, state: &mut CheatcodesCore) -> Result {
170        let Self { json: value, path, valueKey } = self;
171
172        // Read and parse the TOML file, keeping its formatting and comments.
173        // If the file doesn't exist, start with an empty document so the file is created.
174        let data_path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
175        let mut document = if data_path.exists() {
176            parse_toml_document(&fs::locked_read_to_string(&data_path)?)?
177        } else {
178            DocumentMut::new()
179        };
180        upsert_toml_value(&mut document, value, valueKey)?;
181
182        super::fs::write_file(state, path.as_ref(), document.to_string().as_bytes())
183    }
184}
185
186/// Parse
187fn parse_toml_str(toml: &str) -> Result<TomlValue> {
188    toml::from_str(toml).map_err(|e| fmt_err!("failed parsing TOML: {e}"))
189}
190
191/// Parse a TOML string and return the value at the given path.
192fn parse_toml(toml: &str, key: &str, struct_defs: Option<&StructDefinitions>) -> Result {
193    parse_json(&toml_to_json_string(toml)?, key, struct_defs)
194}
195
196/// Parse a TOML string and return the value at the given path, coercing it to the given type.
197fn parse_toml_coerce(toml: &str, key: &str, ty: &DynSolType) -> Result {
198    parse_json_coerce(&toml_to_json_string(toml)?, key, ty)
199}
200
201/// Parse a TOML string and return the value at the given path, coercing it to the given type, or
202/// return the default if the path does not exist.
203fn parse_toml_coerce_default<T: SolValue>(
204    toml: &str,
205    key: &str,
206    ty: &DynSolType,
207    default: &T,
208) -> Result {
209    parse_json_coerce_default(&toml_to_json_string(toml)?, key, ty, default)
210}
211
212/// Parse a TOML string and return an array of all keys at the given path.
213fn parse_toml_keys(toml: &str, key: &str) -> Result {
214    parse_json_keys(&toml_to_json_string(toml)?, key)
215}
216
217/// Convert a TOML string to a JSON string.
218fn toml_to_json_string(toml: &str) -> Result<String> {
219    let toml = parse_toml_str(toml)?;
220    let json = toml_to_json_value(toml);
221    serde_json::to_string(&json).map_err(|e| fmt_err!("failed to serialize JSON: {e}"))
222}
223
224/// Format a JSON value to a TOML pretty string.
225fn format_json_to_toml(json: JsonValue) -> Result<String> {
226    let toml = json_to_toml_value(json);
227    toml::to_string_pretty(&toml).map_err(|e| fmt_err!("failed to serialize TOML: {e}"))
228}
229
230/// Convert a TOML value to a JSON value.
231pub(super) fn toml_to_json_value(toml: TomlValue) -> JsonValue {
232    match toml {
233        TomlValue::String(s) => match s.as_str() {
234            "null" => JsonValue::Null,
235            _ => JsonValue::String(s),
236        },
237        TomlValue::Integer(i) => JsonValue::Number(i.into()),
238        TomlValue::Float(f) => match serde_json::Number::from_f64(f) {
239            Some(n) => JsonValue::Number(n),
240            None => JsonValue::String(f.to_string()),
241        },
242        TomlValue::Boolean(b) => JsonValue::Bool(b),
243        TomlValue::Array(a) => JsonValue::Array(a.into_iter().map(toml_to_json_value).collect()),
244        TomlValue::Table(t) => {
245            JsonValue::Object(t.into_iter().map(|(k, v)| (k, toml_to_json_value(v))).collect())
246        }
247        TomlValue::Datetime(d) => JsonValue::String(d.to_string()),
248    }
249}
250
251/// Convert a JSON value to a TOML value.
252fn json_to_toml_value(json: JsonValue) -> TomlValue {
253    match json {
254        JsonValue::String(s) => TomlValue::String(s),
255        JsonValue::Number(n) => match n.as_i64() {
256            Some(i) => TomlValue::Integer(i),
257            None => match n.as_f64() {
258                Some(f) => TomlValue::Float(f),
259                None => TomlValue::String(n.to_string()),
260            },
261        },
262        JsonValue::Bool(b) => TomlValue::Boolean(b),
263        JsonValue::Array(a) => TomlValue::Array(a.into_iter().map(json_to_toml_value).collect()),
264        JsonValue::Object(o) => {
265            TomlValue::Table(o.into_iter().map(|(k, v)| (k, json_to_toml_value(v))).collect())
266        }
267        JsonValue::Null => TomlValue::String("null".to_string()),
268    }
269}
270
271/// Parses a TOML string into a document that keeps its formatting and comments.
272fn parse_toml_document(toml: &str) -> Result<DocumentMut> {
273    toml.parse().map_err(|e| fmt_err!("failed parsing TOML: {e}"))
274}
275
276/// Inserts or replaces the value at `key` in a TOML document, creating intermediate tables if
277/// necessary.
278///
279/// Only the item at `key` is rewritten, so comments and formatting elsewhere in the document are
280/// kept.
281fn upsert_toml_value(document: &mut DocumentMut, value: &str, key: &str) -> Result<()> {
282    let parts = split_value_key(key)?;
283
284    // Separate the final key from the path.
285    // Traverse the tables, creating implicit intermediary ones if necessary.
286    if let Some((key_to_insert, path_to_parent)) = parts.split_last() {
287        let mut current_level = document.as_item_mut();
288
289        for segment in path_to_parent {
290            let is_inline = current_level.is_inline_table();
291            let Some(table) = current_level.as_table_like_mut() else {
292                return Err(fmt_err!("path segment '{segment}' does not resolve to an object."));
293            };
294            if !table.contains_key(segment) {
295                let mut intermediary = Table::new();
296                intermediary.set_implicit(true);
297                insert_toml_item(table, is_inline, segment, Item::Table(intermediary));
298            }
299            current_level = table.get_mut(segment).unwrap();
300        }
301
302        let is_inline = current_level.is_inline_table();
303        let Some(parent) = current_level.as_table_like_mut() else {
304            return Err(fmt_err!("final destination is not an object, cannot insert key."));
305        };
306
307        let value =
308            serde_json::from_str(value).unwrap_or_else(|_| JsonValue::String(value.to_owned()));
309        let mut item = json_to_toml_item(value)?;
310
311        // Replace an existing item in place: `insert` would reset the key's formatting, which holds
312        // the comments above it.
313        match parent.get_mut(key_to_insert) {
314            Some(existing) if !existing.is_none() => {
315                if existing.is_value() {
316                    // Keep inline values inline instead of turning them into table sections.
317                    item = item.into_value().map_or_else(|item| item, Item::Value);
318                }
319                let shape_changed =
320                    std::mem::discriminant(existing) != std::mem::discriminant(&item);
321                // The comments above a table header are part of the table's decor.
322                let header_prefix =
323                    tables(existing).first().and_then(|table| table.decor().prefix()).cloned();
324                match (&*existing, &mut item) {
325                    (Item::Value(old), Item::Value(new)) => {
326                        *new.decor_mut() = old.decor().clone();
327                    }
328                    (old, new) => {
329                        for (old, new) in tables(old).into_iter().zip(tables_mut(new)) {
330                            *new.decor_mut() = old.decor().clone();
331                            new.set_position(old.position());
332                            if !old.is_implicit() {
333                                new.set_implicit(false);
334                            }
335                        }
336                    }
337                }
338                let new_is_value = item.is_value();
339                *existing = item;
340                if shape_changed {
341                    let mut key = parent.key_mut(key_to_insert).expect("replaced key must exist");
342                    let decor = key.leaf_decor_mut();
343                    decor.clear();
344                    if new_is_value {
345                        decor.set_suffix(" ");
346                        if let Some(prefix) = header_prefix {
347                            decor.set_prefix(prefix);
348                        }
349                    }
350                }
351            }
352            _ => insert_toml_item(parent, is_inline, key_to_insert, item),
353        }
354    }
355
356    Ok(())
357}
358
359/// Inserts a new `item` at `key` into `parent`.
360///
361/// Inline tables can only hold values. There, the new value takes over the whitespace or comment
362/// that followed the previous last value, so the inline table keeps its layout.
363fn insert_toml_item(parent: &mut dyn TableLike, is_inline: bool, key: &str, mut item: Item) {
364    if is_inline {
365        let trailing =
366            parent.iter_mut().filter_map(|(_, item)| item.as_value_mut()).last().and_then(|last| {
367                let suffix = last.decor().suffix().cloned();
368                last.decor_mut().set_suffix("");
369                suffix
370            });
371        item = item.into_value().map_or_else(|item| item, Item::Value);
372        if let Some(value) = item.as_value_mut() {
373            value.decor_mut().clear();
374            if let Some(trailing) = trailing {
375                value.decor_mut().set_suffix(trailing);
376            }
377        }
378    }
379    parent.insert(key, item);
380}
381
382/// Converts a JSON value to a TOML item, formatted the same way as [`format_json_to_toml`].
383fn json_to_toml_item(value: JsonValue) -> Result<Item> {
384    const KEY: &str = "value";
385
386    let wrapper = JsonValue::Object([(KEY.to_string(), value)].into_iter().collect());
387    let mut item = parse_toml_document(&format_json_to_toml(wrapper)?)?
388        .remove(KEY)
389        .ok_or_else(|| fmt_err!("failed to serialize TOML value"))?;
390    // Drop the layout of the temporary document so new tables are placed after their parent.
391    reset_table_layout(&mut item);
392    Ok(item)
393}
394
395/// Recursively clears the document position and header whitespace of all tables in `item`.
396fn reset_table_layout(item: &mut Item) {
397    for table in tables_mut(item) {
398        table.set_position(None);
399        table.decor_mut().clear();
400        for (_, item) in table.iter_mut() {
401            reset_table_layout(item);
402        }
403    }
404}
405
406/// Returns the tables of `item` that are written with a header.
407fn tables(item: &Item) -> Vec<&Table> {
408    match item {
409        Item::Table(table) => vec![table],
410        Item::ArrayOfTables(array) => array.iter().collect(),
411        _ => Vec::new(),
412    }
413}
414
415/// Returns the tables of `item` that are written with a header.
416fn tables_mut(item: &mut Item) -> Vec<&mut Table> {
417    match item {
418        Item::Table(table) => vec![table],
419        Item::ArrayOfTables(array) => array.iter_mut().collect(),
420        _ => Vec::new(),
421    }
422}
423
424#[cfg(test)]
425mod tests {
426    use super::*;
427
428    const CONFIG: &str = r#"# Deployment config.
429
430owner = "0x000000000000000000000000000000000000dEaD" # multisig
431deployed_at = 2024-04-27T11:57:21Z
432max_supply = 0xffff_ffff
433salt = 'literal-string'
434
435# Mainnet settings.
436[mainnet]
437token = "0x0000000000000000000000000000000000000000" # filled by script
438limits = { daily = 1_000, weekly = 5_000 }
439
440# Base settings.
441[base]
442chain_id = 8453
443"#;
444
445    fn upsert(toml: &str, value: &str, key: &str) -> Result<String> {
446        let mut document = parse_toml_document(toml)?;
447        upsert_toml_value(&mut document, value, key)?;
448        Ok(document.to_string())
449    }
450
451    #[test]
452    fn upsert_toml_keeps_comment_above_key() {
453        let toml = "key1 = \"1\"\n\n# this is key2\nkey2 = \"2\"\n";
454        assert_eq!(
455            upsert(toml, "abcd", ".key2").unwrap(),
456            "key1 = \"1\"\n\n# this is key2\nkey2 = \"abcd\"\n"
457        );
458    }
459
460    #[test]
461    fn upsert_toml_replaces_only_the_value() {
462        let address = "0x000000000000000000000000000000000000bEEF";
463        assert_eq!(
464            upsert(CONFIG, address, ".mainnet.token").unwrap(),
465            CONFIG.replace(
466                r#"token = "0x0000000000000000000000000000000000000000""#,
467                &format!(r#"token = "{address}""#)
468            )
469        );
470        assert_eq!(
471            upsert(CONFIG, "2000", ".mainnet.limits.daily").unwrap(),
472            CONFIG.replace("daily = 1_000", "daily = 2000")
473        );
474        assert_eq!(
475            upsert(CONFIG, "{\"chain_id\": 10}", "base").unwrap(),
476            CONFIG.replace("chain_id = 8453", "chain_id = 10")
477        );
478        assert_eq!(
479            upsert(CONFIG, "{\"daily\": 2000}", ".mainnet.limits").unwrap(),
480            CONFIG.replace("daily = 1_000, weekly = 5_000", "daily = 2000")
481        );
482    }
483
484    #[test]
485    fn upsert_toml_adds_keys() {
486        assert_eq!(
487            upsert(CONFIG, "30000000", ".base.gas_limit").unwrap(),
488            format!("{CONFIG}gas_limit = 30000000\n")
489        );
490        assert_eq!(
491            upsert(CONFIG, "{\"block\": 123}", ".optimism.contracts").unwrap(),
492            format!("{CONFIG}\n[optimism.contracts]\nblock = 123\n")
493        );
494        assert_eq!(
495            upsert(CONFIG, "{\"monthly\": 9}", ".mainnet.limits.extra").unwrap(),
496            CONFIG.replace("weekly = 5_000 }", "weekly = 5_000, extra = { monthly = 9 } }")
497        );
498    }
499
500    #[test]
501    fn upsert_toml_preserves_inline_table_formatting_when_adding_keys() {
502        let toml = "limits = { daily = 1_000, weekly = 5_000 }\n";
503        assert_eq!(
504            upsert(toml, "9000", ".limits.monthly").unwrap(),
505            "limits = { daily = 1_000, weekly = 5_000, monthly = 9000 }\n"
506        );
507        assert_eq!(
508            upsert(toml, "9000", ".limits.extra.monthly").unwrap(),
509            "limits = { daily = 1_000, weekly = 5_000, extra = { monthly = 9000 } }\n"
510        );
511
512        let toml = "limits={daily=1_000,weekly = 5_000}\n";
513        assert_eq!(
514            upsert(toml, "9000", ".limits.monthly").unwrap(),
515            "limits={daily=1_000,weekly = 5_000, monthly = 9000}\n"
516        );
517
518        let toml = "limits = {\n    # Keep this cap conservative.\n    daily  = 1_000,\n    weekly = 5_000,\n}\n";
519        assert_eq!(
520            upsert(toml, "9000", ".limits.monthly").unwrap(),
521            "limits = {\n    # Keep this cap conservative.\n    daily  = 1_000,\n    weekly = 5_000, monthly = 9000,\n}\n"
522        );
523    }
524
525    #[test]
526    fn upsert_toml_replaces_arrays_of_tables_in_place() {
527        let toml = "# Production RPC endpoints.\n[[rpc]]\nurl = \"old\"\n";
528        assert_eq!(
529            upsert(toml, r#"[{"url":"new"}]"#, ".rpc").unwrap(),
530            "# Production RPC endpoints.\n[[rpc]]\nurl = \"new\"\n"
531        );
532    }
533
534    #[test]
535    fn upsert_toml_keeps_comments_when_the_item_shape_changes() {
536        let toml = "# Deployment settings.\n[deployment] # maintained by the deploy script\nchain_id = 1\n";
537        assert_eq!(
538            upsert(toml, "disabled", ".deployment").unwrap(),
539            "# Deployment settings.\ndeployment = \"disabled\"\n"
540        );
541
542        let toml = "# Production RPC endpoint.\n[rpc]\nurl = \"a\"\n";
543        assert_eq!(
544            upsert(toml, r#"[{"url":"a"},{"url":"b"}]"#, ".rpc").unwrap(),
545            "# Production RPC endpoint.\n[[rpc]]\nurl = \"a\"\n\n[[rpc]]\nurl = \"b\"\n"
546        );
547
548        let toml = "# Production RPC endpoints.\nrpc = []\n";
549        assert_eq!(
550            upsert(toml, r#"[{"url":"a"},{"url":"b"}]"#, ".rpc").unwrap(),
551            "# Production RPC endpoints.\nrpc = [{ url = \"a\" }, { url = \"b\" }]\n"
552        );
553    }
554
555    #[test]
556    fn upsert_toml_formats_new_values_like_write_toml() {
557        let value = r#"{"list": ["0x01", "0x02"], "empty": {}, "nested": {"a": {"b": 1}}}"#;
558        let expected = format_json_to_toml(
559            serde_json::json!({ "new": serde_json::from_str::<JsonValue>(value).unwrap() }),
560        )
561        .unwrap();
562        assert_eq!(upsert("", value, ".new").unwrap(), expected);
563    }
564
565    #[test]
566    fn upsert_toml_errors() {
567        assert_eq!(
568            upsert(CONFIG, "1", ".owner.x").unwrap_err().to_string(),
569            "final destination is not an object, cannot insert key."
570        );
571        assert_eq!(
572            upsert(CONFIG, "1", ".owner.x.y").unwrap_err().to_string(),
573            "path segment 'x' does not resolve to an object."
574        );
575        assert_eq!(
576            upsert(CONFIG, "1", "$.").unwrap_err().to_string(),
577            "'valueKey' cannot be empty or just '$'"
578        );
579    }
580}
foundry_cheatcodes/toml.rs

foundry_cheatcodes/
toml.rs
