forge/cmd/
selectors.rs

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