Skip to main content

forge_script_sequence/
sequence.rs

1use crate::transaction::TransactionWithMetadata;
2use alloy_network::{Network, ReceiptResponse};
3use alloy_primitives::{TxHash, hex, map::HashMap};
4use eyre::{ContextCompat, Result, WrapErr};
5use foundry_common::{SELECTOR_LEN, TransactionMaybeSigned, fs, shell};
6use foundry_compilers::ArtifactId;
7use foundry_config::Config;
8use serde::{Deserialize, Serialize};
9use std::{
10    collections::VecDeque,
11    path::PathBuf,
12    time::{Duration, SystemTime, UNIX_EPOCH},
13};
14
15pub const DRY_RUN_DIR: &str = "dry-run";
16
17#[derive(Clone, Serialize, Deserialize)]
18pub struct NestedValue {
19    pub internal_type: String,
20    pub value: String,
21}
22
23/// Sensitive values from the transactions in a script sequence
24#[derive(Clone, Default, Serialize, Deserialize)]
25pub struct SensitiveTransactionMetadata {
26    pub rpc: String,
27}
28
29/// Sensitive info from the script sequence which is saved into the cache folder
30#[derive(Clone, Default, Serialize, Deserialize)]
31pub struct SensitiveScriptSequence {
32    pub transactions: VecDeque<SensitiveTransactionMetadata>,
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub recovery_generation: Option<TxHash>,
35}
36
37/// Helper that saves the transactions sequence and its state on which transactions have been
38/// broadcasted
39#[derive(Clone, Serialize, Deserialize)]
40#[serde(bound(
41    serialize = "N::TransactionRequest: Serialize, N::TxEnvelope: Serialize",
42    deserialize = "N::TransactionRequest: for<'de2> Deserialize<'de2>, N::TxEnvelope: for<'de2> Deserialize<'de2>"
43))]
44pub struct ScriptSequence<N: Network> {
45    pub transactions: VecDeque<TransactionWithMetadata<N>>,
46    pub receipts: Vec<N::ReceiptResponse>,
47    pub libraries: Vec<String>,
48    pub pending: Vec<TxHash>,
49    #[serde(skip)]
50    /// Contains paths to the sequence files
51    /// None if sequence should not be saved to disk (e.g. part of a multi-chain sequence)
52    pub paths: Option<(PathBuf, PathBuf)>,
53    pub returns: HashMap<String, NestedValue>,
54    pub timestamp: u128,
55    pub chain: u64,
56    pub commit: Option<String>,
57    #[serde(default, skip_serializing_if = "Option::is_none")]
58    pub recovery_generation: Option<TxHash>,
59}
60
61impl<N: Network> Default for ScriptSequence<N> {
62    fn default() -> Self {
63        Self {
64            transactions: Default::default(),
65            receipts: Default::default(),
66            libraries: Default::default(),
67            pending: Default::default(),
68            paths: Default::default(),
69            returns: Default::default(),
70            timestamp: Default::default(),
71            chain: Default::default(),
72            commit: Default::default(),
73            recovery_generation: Default::default(),
74        }
75    }
76}
77
78impl<N: Network> From<&ScriptSequence<N>> for SensitiveScriptSequence {
79    fn from(sequence: &ScriptSequence<N>) -> Self {
80        Self {
81            transactions: sequence
82                .transactions
83                .iter()
84                .map(|tx| SensitiveTransactionMetadata { rpc: tx.rpc.clone() })
85                .collect(),
86            recovery_generation: sequence.recovery_generation,
87        }
88    }
89}
90
91impl<N: Network> ScriptSequence<N> {
92    /// Loads The sequence for the corresponding json file
93    pub fn load(
94        config: &Config,
95        sig: &str,
96        target: &ArtifactId,
97        chain_id: u64,
98        dry_run: bool,
99    ) -> Result<Self>
100    where
101        N::TxEnvelope: for<'d> Deserialize<'d>,
102    {
103        let (path, sensitive_path) = Self::get_paths(config, sig, target, chain_id, dry_run)?;
104
105        let mut script_sequence: Self = fs::read_json_file(&path)
106            .wrap_err(format!("Deployment not found for chain `{chain_id}`."))?;
107
108        let sensitive_script_sequence: SensitiveScriptSequence = fs::read_json_file(
109            &sensitive_path,
110        )
111        .wrap_err(format!("Deployment's sensitive details not found for chain `{chain_id}`."))?;
112
113        script_sequence.fill_sensitive(&sensitive_script_sequence).wrap_err(format!(
114            "Deployment's sensitive details are out of sync with the broadcast file for chain `{chain_id}`; restore matching broadcast and sensitive-cache files before resuming."
115        ))?;
116
117        script_sequence.paths = Some((path, sensitive_path));
118
119        Ok(script_sequence)
120    }
121
122    /// Saves the transactions as file if it's a standalone deployment.
123    /// `save_ts` should be set to true for checkpoint updates, which might happen many times and
124    /// could result in us saving many identical files.
125    pub fn save(&mut self, silent: bool, save_ts: bool) -> Result<()>
126    where
127        N::TxEnvelope: Serialize,
128    {
129        self.sort_receipts();
130
131        if self.transactions.is_empty() {
132            return Ok(());
133        }
134
135        self.timestamp = now().as_millis();
136        let ts_name = format!("run-{}.json", self.timestamp);
137
138        let sensitive_script_sequence = SensitiveScriptSequence::from(&*self);
139
140        let Some((path, sensitive_path)) = self.paths.as_ref() else { return Ok(()) };
141
142        // broadcast folder writes
143        //../run-latest.json
144        fs::write_pretty_json_file(path, &self)?;
145        if save_ts {
146            //../run-[timestamp].json
147            fs::copy(path, path.with_file_name(&ts_name))?;
148        }
149
150        // cache folder writes
151        //../run-latest.json
152        fs::write_sensitive_json_file(sensitive_path, &sensitive_script_sequence)?;
153        if save_ts {
154            //../run-[timestamp].json
155            fs::copy(sensitive_path, sensitive_path.with_file_name(&ts_name))?;
156        }
157
158        if !silent {
159            if shell::is_json() {
160                sh_println!(
161                    "{}",
162                    serde_json::json!({
163                        "status": "success",
164                        "transactions": path.display().to_string(),
165                        "sensitive": sensitive_path.display().to_string(),
166                    })
167                )?;
168            } else {
169                sh_println!("\nTransactions saved to: {}\n", path.display())?;
170                sh_println!("Sensitive values saved to: {}\n", sensitive_path.display())?;
171            }
172        }
173
174        Ok(())
175    }
176
177    pub fn add_receipt(&mut self, receipt: N::ReceiptResponse) {
178        self.receipts.push(receipt);
179    }
180
181    /// Sorts all receipts with ascending transaction index
182    pub fn sort_receipts(&mut self) {
183        self.receipts.sort_by_key(|r| (r.block_number(), r.transaction_index()));
184    }
185
186    pub fn add_pending(&mut self, index: usize, tx_hash: TxHash) {
187        if !self.pending.contains(&tx_hash) {
188            self.transactions[index].hash = Some(tx_hash);
189            self.pending.push(tx_hash);
190        }
191    }
192
193    pub fn remove_pending(&mut self, tx_hash: TxHash) {
194        self.pending.retain(|element| element != &tx_hash);
195    }
196
197    /// Gets paths in the formats
198    /// `./broadcast/[contract_filename]/[chain_id]/[sig]-latest.json` and
199    /// `./cache/[contract_filename]/[chain_id]/[sig]-latest.json`.
200    pub fn get_paths(
201        config: &Config,
202        sig: &str,
203        target: &ArtifactId,
204        chain_id: u64,
205        dry_run: bool,
206    ) -> Result<(PathBuf, PathBuf)> {
207        let mut broadcast = config.broadcast.clone();
208        let mut cache = config.cache_path.clone();
209        let mut common = PathBuf::new();
210
211        let target_fname = target.source.file_name().wrap_err("No filename.")?;
212        common.push(target_fname);
213        common.push(chain_id.to_string());
214        if dry_run {
215            common.push(DRY_RUN_DIR);
216        }
217
218        broadcast.push(common.clone());
219        cache.push(common);
220
221        fs::create_dir_all(&broadcast)?;
222        fs::create_dir_all(&cache)?;
223
224        // TODO: ideally we want the name of the function here if sig is calldata
225        let filename = sig_to_file_name(sig);
226        let filename_with_ext = format!("{filename}-latest.json");
227
228        broadcast.push(&filename_with_ext);
229        cache.push(&filename_with_ext);
230
231        Ok((broadcast, cache))
232    }
233
234    /// Returns the first RPC URL of this sequence.
235    pub fn rpc_url(&self) -> &str {
236        self.transactions.front().expect("empty sequence").rpc.as_str()
237    }
238
239    /// Returns the list of the transactions without the metadata.
240    pub fn transactions(&self) -> impl Iterator<Item = &TransactionMaybeSigned<N>> {
241        self.transactions.iter().map(|tx| tx.tx())
242    }
243
244    /// Copies RPC URLs from a matching sensitive-cache sequence.
245    pub fn fill_sensitive(&mut self, sensitive: &SensitiveScriptSequence) -> Result<()> {
246        if self.recovery_generation != sensitive.recovery_generation {
247            eyre::bail!(
248                "the broadcast file and its sensitive-cache counterpart belong to different recovery generations"
249            );
250        }
251        let transactions_len = self.transactions.len();
252        let sensitive_len = sensitive.transactions.len();
253        if transactions_len != sensitive_len {
254            eyre::bail!(
255                "sensitive-cache entry count ({sensitive_len}) does not match transaction count \
256                 ({transactions_len}); the broadcast file and its sensitive-cache counterpart are \
257                 out of sync"
258            );
259        }
260        for (i, tx) in self.transactions.iter_mut().enumerate() {
261            tx.rpc.clone_from(&sensitive.transactions[i].rpc);
262        }
263        Ok(())
264    }
265}
266
267/// Converts the `sig` argument into the corresponding file path.
268///
269/// This accepts either the signature of the function or the raw calldata.
270pub fn sig_to_file_name(sig: &str) -> String {
271    if let Some((name, _)) = sig.split_once('(') {
272        // strip until call argument parenthesis
273        return name.to_string();
274    }
275    // assume calldata if `sig` is hex
276    if let Ok(calldata) = hex::decode(sig) {
277        // in which case we return the function selector if available
278        if let Some(selector) = calldata.get(..SELECTOR_LEN) {
279            return hex::encode(selector);
280        }
281        // fallback to original string if calldata is too short to contain selector
282        return sig.to_string();
283    }
284
285    sig.to_string()
286}
287
288pub fn now() -> Duration {
289    SystemTime::now().duration_since(UNIX_EPOCH).expect("time went backwards")
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use alloy_network::Ethereum;
296
297    fn sequence_with_two_transactions() -> ScriptSequence<Ethereum> {
298        let mut sequence = ScriptSequence::default();
299        for rpc in ["first", "second"] {
300            let mut tx = TransactionWithMetadata::from_tx_request(
301                TransactionMaybeSigned::Unsigned(Default::default()),
302            );
303            tx.rpc = rpc.to_string();
304            sequence.transactions.push_back(tx);
305        }
306        sequence
307    }
308
309    #[test]
310    fn fill_sensitive_rejects_mismatched_counts_without_mutation() {
311        for count in [1, 3] {
312            let mut sequence = sequence_with_two_transactions();
313            let sensitive = SensitiveScriptSequence {
314                transactions: (0..count)
315                    .map(|_| SensitiveTransactionMetadata { rpc: "replacement".to_string() })
316                    .collect(),
317                recovery_generation: None,
318            };
319            assert_eq!(
320                sequence.fill_sensitive(&sensitive).unwrap_err().to_string(),
321                format!(
322                    "sensitive-cache entry count ({count}) does not match transaction count (2); \
323                     the broadcast file and its sensitive-cache counterpart are out of sync"
324                )
325            );
326            assert_eq!(
327                sequence.transactions.iter().map(|tx| tx.rpc.as_str()).collect::<Vec<_>>(),
328                ["first", "second"]
329            );
330        }
331    }
332
333    #[test]
334    fn fill_sensitive_restores_matching_cache() {
335        let mut sequence = sequence_with_two_transactions();
336        let sensitive = SensitiveScriptSequence {
337            transactions: ["restored-first", "restored-second"]
338                .into_iter()
339                .map(|rpc| SensitiveTransactionMetadata { rpc: rpc.to_string() })
340                .collect(),
341            recovery_generation: None,
342        };
343        sequence.fill_sensitive(&sensitive).unwrap();
344        assert_eq!(
345            sequence.transactions.iter().map(|tx| tx.rpc.as_str()).collect::<Vec<_>>(),
346            ["restored-first", "restored-second"]
347        );
348    }
349
350    #[test]
351    fn fill_sensitive_rejects_a_different_recovery_generation() {
352        let mut sequence = sequence_with_two_transactions();
353        sequence.recovery_generation = Some(TxHash::ZERO);
354        let sensitive = SensitiveScriptSequence {
355            transactions: Default::default(),
356            recovery_generation: Some(TxHash::with_last_byte(1)),
357        };
358
359        assert_eq!(
360            sequence.fill_sensitive(&sensitive).unwrap_err().to_string(),
361            "the broadcast file and its sensitive-cache counterpart belong to different recovery generations"
362        );
363    }
364
365    #[test]
366    fn can_convert_sig() {
367        assert_eq!(sig_to_file_name("run()").as_str(), "run");
368        assert_eq!(
369            sig_to_file_name(
370                "522bb704000000000000000000000000f39fd6e51aad88f6f4ce6ab8827279cfFFb92266"
371            )
372            .as_str(),
373            "522bb704"
374        );
375        // valid calldata with 0x prefix
376        assert_eq!(
377            sig_to_file_name(
378                "0x522bb704000000000000000000000000f39fd6e51aad88f6f4ce6ab8827279cfFFb92266"
379            )
380            .as_str(),
381            "522bb704"
382        );
383        // short calldata: should not panic and should return input as-is
384        assert_eq!(sig_to_file_name("0x1234").as_str(), "0x1234");
385        assert_eq!(sig_to_file_name("123").as_str(), "123");
386        // invalid hex: should return input as-is
387        assert_eq!(sig_to_file_name("0xnotahex").as_str(), "0xnotahex");
388        // non-hex non-signature: should return input as-is
389        assert_eq!(sig_to_file_name("not_a_sig_or_hex").as_str(), "not_a_sig_or_hex");
390    }
391}