forge/cmd/
selectors.rs

1use alloy_primitives::hex;
2use clap::Parser;
3use comfy_table::{Table, modifiers::UTF8_ROUND_CORNERS, presets::ASCII_MARKDOWN};
4use eyre::Result;
5use foundry_cli::{
6    opts::{BuildOpts, CompilerOpts, ProjectPathOpts},
7    utils::{FoundryPathExt, cache_local_signatures},
8};
9use foundry_common::{
10    compile::{PathOrContractInfo, ProjectCompiler, compile_target},
11    selectors::{SelectorImportData, import_selectors},
12    shell,
13};
14use foundry_compilers::{artifacts::output_selection::ContractOutputSelection, info::ContractInfo};
15use std::{collections::BTreeMap, fs::canonicalize};
16
17/// CLI arguments for `forge selectors`.
18#[derive(Clone, Debug, Parser)]
19pub enum SelectorsSubcommands {
20    /// Check for selector collisions between contracts
21    #[command(visible_alias = "co")]
22    Collision {
23        /// The first of the two contracts for which to look selector collisions for, in the form
24        /// `(<path>:)?<contractname>`.
25        first_contract: ContractInfo,
26
27        /// The second of the two contracts for which to look selector collisions for, in the form
28        /// `(<path>:)?<contractname>`.
29        second_contract: ContractInfo,
30
31        #[command(flatten)]
32        build: Box<BuildOpts>,
33    },
34
35    /// Upload selectors to registry
36    #[command(visible_alias = "up")]
37    Upload {
38        /// The name of the contract to upload selectors for.
39        /// Can also be in form of `path:contract name`.
40        #[arg(required_unless_present = "all")]
41        contract: Option<PathOrContractInfo>,
42
43        /// Upload selectors for all contracts in the project.
44        #[arg(long, required_unless_present = "contract")]
45        all: bool,
46
47        #[command(flatten)]
48        project_paths: ProjectPathOpts,
49    },
50
51    /// List selectors from current workspace
52    #[command(visible_alias = "ls")]
53    List {
54        /// The name of the contract to list selectors for.
55        #[arg(help = "The name of the contract to list selectors for.")]
56        contract: Option<String>,
57
58        #[command(flatten)]
59        project_paths: ProjectPathOpts,
60
61        #[arg(long, help = "Do not group the selectors by contract in separate tables.")]
62        no_group: bool,
63    },
64
65    /// Find if a selector is present in the project
66    #[command(visible_alias = "f")]
67    Find {
68        /// The selector to search for
69        #[arg(help = "The selector to search for (with or without 0x prefix)")]
70        selector: String,
71
72        #[command(flatten)]
73        project_paths: ProjectPathOpts,
74    },
75
76    /// Cache project selectors (enables trace with local contracts functions and events).
77    #[command(visible_alias = "c")]
78    Cache {
79        #[command(flatten)]
80        project_paths: ProjectPathOpts,
81    },
82}
83
84impl SelectorsSubcommands {
85    pub async fn run(self) -> Result<()> {
86        match self {
87            Self::Cache { project_paths } => {
88                sh_println!("Caching selectors for contracts in the project...")?;
89                let build_args = BuildOpts {
90                    project_paths,
91                    compiler: CompilerOpts {
92                        extra_output: vec![ContractOutputSelection::Abi],
93                        ..Default::default()
94                    },
95                    ..Default::default()
96                };
97
98                // compile the project to get the artifacts/abis
99                let project = build_args.project()?;
100                let outcome = ProjectCompiler::new().quiet(true).compile(&project)?;
101                cache_local_signatures(&outcome)?;
102            }
103            Self::Upload { contract, all, project_paths } => {
104                let build_args = BuildOpts {
105                    project_paths: project_paths.clone(),
106                    compiler: CompilerOpts {
107                        extra_output: vec![ContractOutputSelection::Abi],
108                        ..Default::default()
109                    },
110                    ..Default::default()
111                };
112
113                let project = build_args.project()?;
114                let output = if let Some(contract_info) = &contract {
115                    let Some(contract_name) = contract_info.name() else {
116                        eyre::bail!("No contract name provided.")
117                    };
118
119                    let target_path = contract_info
120                        .path()
121                        .map(Ok)
122                        .unwrap_or_else(|| project.find_contract_path(contract_name))?;
123                    compile_target(&target_path, &project, false)?
124                } else {
125                    ProjectCompiler::new().compile(&project)?
126                };
127                let artifacts = if all {
128                    output
129                        .into_artifacts_with_files()
130                        .filter(|(file, _, _)| {
131                            let is_sources_path = file.starts_with(&project.paths.sources);
132                            let is_test = file.is_sol_test();
133
134                            is_sources_path && !is_test
135                        })
136                        .map(|(_, contract, artifact)| (contract, artifact))
137                        .collect()
138                } else {
139                    let contract_info = contract.unwrap();
140                    let contract = contract_info.name().unwrap().to_string();
141
142                    let found_artifact = if let Some(path) = contract_info.path() {
143                        output.find(project.root().join(path).as_path(), &contract)
144                    } else {
145                        output.find_first(&contract)
146                    };
147
148                    let artifact = found_artifact
149                        .ok_or_else(|| {
150                            eyre::eyre!(
151                                "Could not find artifact `{contract}` in the compiled artifacts"
152                            )
153                        })?
154                        .clone();
155                    vec![(contract, artifact)]
156                };
157
158                let mut artifacts = artifacts.into_iter().peekable();
159                while let Some((contract, artifact)) = artifacts.next() {
160                    let abi = artifact.abi.ok_or_else(|| eyre::eyre!("Unable to fetch abi"))?;
161                    if abi.functions.is_empty() && abi.events.is_empty() && abi.errors.is_empty() {
162                        continue;
163                    }
164
165                    sh_println!("Uploading selectors for {contract}...")?;
166
167                    // upload abi to selector database
168                    import_selectors(SelectorImportData::Abi(vec![abi])).await?.describe();
169
170                    if artifacts.peek().is_some() {
171                        sh_println!()?
172                    }
173                }
174            }
175            Self::Collision { mut first_contract, mut second_contract, build } => {
176                // Compile the project with the two contracts included
177                let project = build.project()?;
178                let mut compiler = ProjectCompiler::new().quiet(true);
179
180                if let Some(contract_path) = &mut first_contract.path {
181                    let target_path = canonicalize(&*contract_path)?;
182                    *contract_path = target_path.to_string_lossy().to_string();
183                    compiler = compiler.files([target_path]);
184                }
185                if let Some(contract_path) = &mut second_contract.path {
186                    let target_path = canonicalize(&*contract_path)?;
187                    *contract_path = target_path.to_string_lossy().to_string();
188                    compiler = compiler.files([target_path]);
189                }
190
191                let output = compiler.compile(&project)?;
192
193                // Check method selectors for collisions
194                let methods = |contract: &ContractInfo| -> eyre::Result<_> {
195                    let artifact = output
196                        .find_contract(contract)
197                        .ok_or_else(|| eyre::eyre!("Could not find artifact for {contract}"))?;
198                    artifact.method_identifiers.as_ref().ok_or_else(|| {
199                        eyre::eyre!("Could not find method identifiers for {contract}")
200                    })
201                };
202                let first_method_map = methods(&first_contract)?;
203                let second_method_map = methods(&second_contract)?;
204
205                let colliding_methods: Vec<(&String, &String, &String)> = first_method_map
206                    .iter()
207                    .filter_map(|(k1, v1)| {
208                        second_method_map
209                            .iter()
210                            .find_map(|(k2, v2)| if **v2 == *v1 { Some((k2, v2)) } else { None })
211                            .map(|(k2, v2)| (v2, k1, k2))
212                    })
213                    .collect();
214
215                if colliding_methods.is_empty() {
216                    sh_println!("No colliding method selectors between the two contracts.")?;
217                } else {
218                    let mut table = Table::new();
219                    if shell::is_markdown() {
220                        table.load_preset(ASCII_MARKDOWN);
221                    } else {
222                        table.apply_modifier(UTF8_ROUND_CORNERS);
223                    }
224                    table.set_header([
225                        String::from("Selector"),
226                        first_contract.name,
227                        second_contract.name,
228                    ]);
229                    for method in &colliding_methods {
230                        table.add_row([method.0, method.1, method.2]);
231                    }
232                    sh_println!("{} collisions found:", colliding_methods.len())?;
233                    sh_println!("\n{table}\n")?;
234                }
235            }
236            Self::List { contract, project_paths, no_group } => {
237                sh_println!("Listing selectors for contracts in the project...")?;
238                let build_args = BuildOpts {
239                    project_paths,
240                    compiler: CompilerOpts {
241                        extra_output: vec![ContractOutputSelection::Abi],
242                        ..Default::default()
243                    },
244                    ..Default::default()
245                };
246
247                // compile the project to get the artifacts/abis
248                let project = build_args.project()?;
249                let outcome = ProjectCompiler::new().quiet(true).compile(&project)?;
250                let artifacts = if let Some(contract) = contract {
251                    let found_artifact = outcome.find_first(&contract);
252                    let artifact = found_artifact
253                        .ok_or_else(|| {
254                            let candidates = outcome
255                                .artifacts()
256                                .map(|(name, _,)| name)
257                                .collect::<Vec<_>>();
258                            let suggestion = if let Some(suggestion) = foundry_cli::utils::did_you_mean(&contract, candidates).pop() {
259                                format!("\nDid you mean `{suggestion}`?")
260                            } else {
261                                String::new()
262                            };
263                            eyre::eyre!(
264                                "Could not find artifact `{contract}` in the compiled artifacts{suggestion}",
265                            )
266                        })?
267                        .clone();
268                    vec![(contract, artifact)]
269                } else {
270                    outcome
271                        .into_artifacts_with_files()
272                        .filter(|(file, _, _)| {
273                            let is_sources_path = file.starts_with(&project.paths.sources);
274                            let is_test = file.is_sol_test();
275
276                            is_sources_path && !is_test
277                        })
278                        .map(|(_, contract, artifact)| (contract, artifact))
279                        .collect()
280                };
281
282                let mut artifacts = artifacts.into_iter().peekable();
283
284                #[derive(PartialEq, PartialOrd, Eq, Ord)]
285                enum SelectorType {
286                    Function,
287                    Event,
288                    Error,
289                }
290                impl std::fmt::Display for SelectorType {
291                    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
292                        match self {
293                            Self::Function => write!(f, "Function"),
294                            Self::Event => write!(f, "Event"),
295                            Self::Error => write!(f, "Error"),
296                        }
297                    }
298                }
299
300                let mut selectors =
301                    BTreeMap::<String, BTreeMap<SelectorType, Vec<(String, String)>>>::new();
302
303                for (contract, artifact) in artifacts.by_ref() {
304                    let abi = artifact.abi.ok_or_else(|| eyre::eyre!("Unable to fetch abi"))?;
305
306                    let contract_selectors = selectors.entry(contract.clone()).or_default();
307
308                    for func in abi.functions() {
309                        let sig = func.signature();
310                        let selector = func.selector();
311                        contract_selectors
312                            .entry(SelectorType::Function)
313                            .or_default()
314                            .push((hex::encode_prefixed(selector), sig));
315                    }
316
317                    for event in abi.events() {
318                        let sig = event.signature();
319                        let selector = event.selector();
320                        contract_selectors
321                            .entry(SelectorType::Event)
322                            .or_default()
323                            .push((hex::encode_prefixed(selector), sig));
324                    }
325
326                    for error in abi.errors() {
327                        let sig = error.signature();
328                        let selector = error.selector();
329                        contract_selectors
330                            .entry(SelectorType::Error)
331                            .or_default()
332                            .push((hex::encode_prefixed(selector), sig));
333                    }
334                }
335
336                if no_group {
337                    let mut table = Table::new();
338                    if shell::is_markdown() {
339                        table.load_preset(ASCII_MARKDOWN);
340                    } else {
341                        table.apply_modifier(UTF8_ROUND_CORNERS);
342                    }
343                    table.set_header(["Type", "Signature", "Selector", "Contract"]);
344
345                    for (contract, contract_selectors) in selectors {
346                        for (selector_type, selectors) in contract_selectors {
347                            for (selector, sig) in selectors {
348                                table.add_row([
349                                    selector_type.to_string(),
350                                    sig,
351                                    selector,
352                                    contract.to_string(),
353                                ]);
354                            }
355                        }
356                    }
357
358                    sh_println!("\n{table}")?;
359                } else {
360                    for (idx, (contract, contract_selectors)) in selectors.into_iter().enumerate() {
361                        sh_println!("{}{contract}", if idx == 0 { "" } else { "\n" })?;
362                        let mut table = Table::new();
363                        if shell::is_markdown() {
364                            table.load_preset(ASCII_MARKDOWN);
365                        } else {
366                            table.apply_modifier(UTF8_ROUND_CORNERS);
367                        }
368                        table.set_header(["Type", "Signature", "Selector"]);
369
370                        for (selector_type, selectors) in contract_selectors {
371                            for (selector, sig) in selectors {
372                                table.add_row([selector_type.to_string(), sig, selector]);
373                            }
374                        }
375                        sh_println!("\n{table}")?;
376                    }
377                }
378            }
379
380            Self::Find { selector, project_paths } => {
381                sh_println!("Searching for selector {selector:?} in the project...")?;
382
383                let build_args = BuildOpts {
384                    project_paths,
385                    compiler: CompilerOpts {
386                        extra_output: vec![ContractOutputSelection::Abi],
387                        ..Default::default()
388                    },
389                    ..Default::default()
390                };
391
392                let project = build_args.project()?;
393                let outcome = ProjectCompiler::new().quiet(true).compile(&project)?;
394                let artifacts = outcome
395                    .into_artifacts_with_files()
396                    .filter(|(file, _, _)| {
397                        let is_sources_path = file.starts_with(&project.paths.sources);
398                        let is_test = file.is_sol_test();
399                        is_sources_path && !is_test
400                    })
401                    .collect::<Vec<_>>();
402
403                let mut table = Table::new();
404                if shell::is_markdown() {
405                    table.load_preset(ASCII_MARKDOWN);
406                } else {
407                    table.apply_modifier(UTF8_ROUND_CORNERS);
408                }
409
410                table.set_header(["Type", "Signature", "Selector", "Contract"]);
411
412                for (_file, contract, artifact) in artifacts {
413                    let abi = artifact.abi.ok_or_else(|| eyre::eyre!("Unable to fetch abi"))?;
414
415                    let selector_bytes =
416                        hex::decode(selector.strip_prefix("0x").unwrap_or(&selector))?;
417
418                    for func in abi.functions() {
419                        if func.selector().as_slice().starts_with(selector_bytes.as_slice()) {
420                            table.add_row([
421                                "Function",
422                                &func.signature(),
423                                &hex::encode_prefixed(func.selector()),
424                                contract.as_str(),
425                            ]);
426                        }
427                    }
428
429                    for event in abi.events() {
430                        if event.selector().as_slice().starts_with(selector_bytes.as_slice()) {
431                            table.add_row([
432                                "Event",
433                                &event.signature(),
434                                &hex::encode_prefixed(event.selector()),
435                                contract.as_str(),
436                            ]);
437                        }
438                    }
439
440                    for error in abi.errors() {
441                        if error.selector().as_slice().starts_with(selector_bytes.as_slice()) {
442                            table.add_row([
443                                "Error",
444                                &error.signature(),
445                                &hex::encode_prefixed(error.selector()),
446                                contract.as_str(),
447                            ]);
448                        }
449                    }
450                }
451
452                if table.row_count() > 0 {
453                    sh_println!("\nFound {} instance(s)...", table.row_count())?;
454                    sh_println!("\n{table}\n")?;
455                } else {
456                    return Err(eyre::eyre!("\nSelector not found in the project."));
457                }
458            }
459        }
460        Ok(())
461    }
462}