forge/cmd/
eip712.rs

1use clap::{Parser, ValueHint};
2use eyre::{Ok, OptionExt, Result};
3use foundry_cli::{opts::BuildOpts, utils::LoadConfig};
4use foundry_common::compile::ProjectCompiler;
5use foundry_compilers::artifacts::{
6    output_selection::OutputSelection,
7    visitor::{Visitor, Walk},
8    ContractDefinition, EnumDefinition, SourceUnit, StructDefinition, TypeDescriptions, TypeName,
9};
10use std::{collections::BTreeMap, fmt::Write, path::PathBuf};
11
12foundry_config::impl_figment_convert!(Eip712Args, build);
13
14/// CLI arguments for `forge eip712`.
15#[derive(Clone, Debug, Parser)]
16pub struct Eip712Args {
17    /// The path to the file from which to read struct definitions.
18    #[arg(value_hint = ValueHint::FilePath, value_name = "PATH")]
19    pub target_path: PathBuf,
20
21    #[command(flatten)]
22    build: BuildOpts,
23}
24
25impl Eip712Args {
26    pub fn run(self) -> Result<()> {
27        let config = self.load_config()?;
28        let mut project = config.ephemeral_project()?;
29        let target_path = dunce::canonicalize(self.target_path)?;
30        project.update_output_selection(|selection| {
31            *selection = OutputSelection::ast_output_selection();
32        });
33
34        let output = ProjectCompiler::new().files([target_path.clone()]).compile(&project)?;
35
36        // Collect ASTs by getting them from sources and converting into strongly typed
37        // `SourceUnit`s.
38        let asts = output
39            .into_output()
40            .sources
41            .into_iter()
42            .filter_map(|(path, mut sources)| Some((path, sources.swap_remove(0).source_file.ast?)))
43            .map(|(path, ast)| {
44                Ok((path, serde_json::from_str::<SourceUnit>(&serde_json::to_string(&ast)?)?))
45            })
46            .collect::<Result<BTreeMap<_, _>>>()?;
47
48        let resolver = Resolver::new(&asts);
49
50        let target_ast = asts
51            .get(&target_path)
52            .ok_or_else(|| eyre::eyre!("Could not find AST for target file {target_path:?}"))?;
53
54        let structs_in_target = {
55            let mut collector = StructCollector::default();
56            target_ast.walk(&mut collector);
57            collector.0
58        };
59
60        for id in structs_in_target.keys() {
61            if let Some(resolved) = resolver.resolve_struct_eip712(*id)? {
62                sh_println!("{resolved}\n")?;
63            }
64        }
65
66        Ok(())
67    }
68}
69
70/// AST [Visitor] used for collecting struct definitions.
71#[derive(Debug, Clone, Default)]
72pub struct StructCollector(pub BTreeMap<usize, StructDefinition>);
73
74impl Visitor for StructCollector {
75    fn visit_struct_definition(&mut self, def: &StructDefinition) {
76        self.0.insert(def.id, def.clone());
77    }
78}
79
80/// Collects mapping from AST id of type definition to representation of this type for EIP-712
81/// encoding.
82///
83/// For now, maps contract definitions to `address` and enums to `uint8`.
84#[derive(Debug, Clone, Default)]
85struct SimpleCustomTypesCollector(BTreeMap<usize, String>);
86
87impl Visitor for SimpleCustomTypesCollector {
88    fn visit_contract_definition(&mut self, def: &ContractDefinition) {
89        self.0.insert(def.id, "address".to_string());
90    }
91
92    fn visit_enum_definition(&mut self, def: &EnumDefinition) {
93        self.0.insert(def.id, "uint8".to_string());
94    }
95}
96
97pub struct Resolver {
98    simple_types: BTreeMap<usize, String>,
99    structs: BTreeMap<usize, StructDefinition>,
100}
101
102impl Resolver {
103    pub fn new(asts: &BTreeMap<PathBuf, SourceUnit>) -> Self {
104        let simple_types = {
105            let mut collector = SimpleCustomTypesCollector::default();
106            asts.values().for_each(|ast| ast.walk(&mut collector));
107
108            collector.0
109        };
110
111        let structs = {
112            let mut collector = StructCollector::default();
113            asts.values().for_each(|ast| ast.walk(&mut collector));
114            collector.0
115        };
116
117        Self { simple_types, structs }
118    }
119
120    /// Converts a given struct definition into EIP-712 `encodeType` representation.
121    ///
122    /// Returns `None` if struct contains any fields that are not supported by EIP-712 (e.g.
123    /// mappings or function pointers).
124    pub fn resolve_struct_eip712(&self, id: usize) -> Result<Option<String>> {
125        let mut subtypes = BTreeMap::new();
126        subtypes.insert(self.structs[&id].name.clone(), id);
127        self.resolve_eip712_inner(id, &mut subtypes, true, None)
128    }
129
130    fn resolve_eip712_inner(
131        &self,
132        id: usize,
133        subtypes: &mut BTreeMap<String, usize>,
134        append_subtypes: bool,
135        rename: Option<&str>,
136    ) -> Result<Option<String>> {
137        let def = &self.structs[&id];
138        let mut result = format!("{}(", rename.unwrap_or(&def.name));
139
140        for (idx, member) in def.members.iter().enumerate() {
141            let Some(ty) = self.resolve_type(
142                member.type_name.as_ref().ok_or_eyre("missing type name")?,
143                subtypes,
144            )?
145            else {
146                return Ok(None)
147            };
148
149            write!(result, "{ty} {name}", name = member.name)?;
150
151            if idx < def.members.len() - 1 {
152                result.push(',');
153            }
154        }
155
156        result.push(')');
157
158        if !append_subtypes {
159            return Ok(Some(result))
160        }
161
162        for (subtype_name, subtype_id) in
163            subtypes.iter().map(|(name, id)| (name.clone(), *id)).collect::<Vec<_>>()
164        {
165            if subtype_id == id {
166                continue
167            }
168            let Some(encoded_subtype) =
169                self.resolve_eip712_inner(subtype_id, subtypes, false, Some(&subtype_name))?
170            else {
171                return Ok(None)
172            };
173            result.push_str(&encoded_subtype);
174        }
175
176        Ok(Some(result))
177    }
178
179    /// Converts given [TypeName] into a type which can be converted to
180    /// [`alloy_dyn_abi::DynSolType`].
181    ///
182    /// Returns `None` if the type is not supported for EIP712 encoding.
183    pub fn resolve_type(
184        &self,
185        type_name: &TypeName,
186        subtypes: &mut BTreeMap<String, usize>,
187    ) -> Result<Option<String>> {
188        match type_name {
189            TypeName::FunctionTypeName(_) | TypeName::Mapping(_) => Ok(None),
190            TypeName::ElementaryTypeName(ty) => Ok(Some(ty.name.clone())),
191            TypeName::ArrayTypeName(ty) => {
192                let Some(inner) = self.resolve_type(&ty.base_type, subtypes)? else {
193                    return Ok(None)
194                };
195                let len = parse_array_length(&ty.type_descriptions)?;
196
197                Ok(Some(format!("{inner}[{}]", len.unwrap_or(""))))
198            }
199            TypeName::UserDefinedTypeName(ty) => {
200                if let Some(name) = self.simple_types.get(&(ty.referenced_declaration as usize)) {
201                    Ok(Some(name.clone()))
202                } else if let Some(def) = self.structs.get(&(ty.referenced_declaration as usize)) {
203                    let name =
204                        // If we've already seen struct with this ID, just use assigned name.
205                        if let Some((name, _)) = subtypes.iter().find(|(_, id)| **id == def.id) {
206                            name.clone()
207                        } else {
208                            // Otherwise, assign new name.
209                            let mut i = 0;
210                            let mut name = def.name.clone();
211                            while subtypes.contains_key(&name) {
212                                i += 1;
213                                name = format!("{}_{i}", def.name);
214                            }
215
216                            subtypes.insert(name.clone(), def.id);
217
218                            // iterate over members to check if they are resolvable and to populate subtypes
219                            for member in &def.members {
220                                if self.resolve_type(
221                                    member.type_name.as_ref().ok_or_eyre("missing type name")?,
222                                    subtypes,
223                                )?
224                                .is_none()
225                                {
226                                    return Ok(None)
227                                }
228                            }
229                            name
230                        };
231
232                    return Ok(Some(name))
233                } else {
234                    return Ok(None)
235                }
236            }
237        }
238    }
239}
240
241fn parse_array_length(type_description: &TypeDescriptions) -> Result<Option<&str>> {
242    let type_string =
243        type_description.type_string.as_ref().ok_or_eyre("missing typeString for array type")?;
244    let Some(inside_brackets) =
245        type_string.rsplit_once("[").and_then(|(_, right)| right.split("]").next())
246    else {
247        eyre::bail!("failed to parse array type string: {type_string}")
248    };
249
250    if inside_brackets.is_empty() {
251        Ok(None)
252    } else {
253        Ok(Some(inside_brackets))
254    }
255}