1use std::path::{Path, PathBuf};
2
3use foundry_config::Config;
4use solar::{
5 ast::{
6 Span,
7 interface::{Session, source_map::FileName},
8 },
9 parse::Parser,
10};
11
12mod assembly;
13mod transform;
14mod value;
15mod visitor;
16
17use transform::apply_transforms;
18use visitor::collect_transforms;
19
20pub fn brutalize_source(path: &Path, source: &str) -> eyre::Result<String> {
21 let sess = Session::builder().with_silent_emitter(None).build();
22
23 let result = sess.enter(|| -> solar::interface::Result<_> {
24 let arena = solar::ast::Arena::new();
25 let ast = {
26 let mut parser = Parser::from_lazy_source_code(
27 &sess,
28 &arena,
29 FileName::from(path.to_path_buf()),
30 || Ok(source.to_string()),
31 )?;
32 parser.parse_file().map_err(|e| e.emit())?
33 };
34
35 Ok(collect_transforms(source, &ast))
36 });
37
38 let transforms = match result {
39 Ok(t) => t,
40 Err(_) => {
41 eyre::bail!("failed to parse {}", path.display());
42 }
43 };
44
45 Ok(apply_transforms(source, transforms))
46}
47
48pub fn brutalize_project(config: &Config, temp_dir: &Path) -> eyre::Result<usize> {
55 let src_rel = config.src.strip_prefix(&config.root).unwrap_or(&config.src);
56 let src_dir = temp_dir.join(src_rel);
57
58 if !src_dir.exists() {
59 return Ok(0);
60 }
61
62 let src_rel = src_rel.to_path_buf();
63 let skipped_dirs = crate::workspace::handled_project_roots(config)?
64 .into_iter()
65 .filter(|rel| !rel.as_os_str().is_empty() && *rel != src_rel)
66 .map(|rel| temp_dir.join(rel))
67 .collect::<Vec<_>>();
68
69 brutalize_sol_files_in_dir(&src_dir, &skipped_dirs)
70}
71
72fn brutalize_sol_files_in_dir(dir: &Path, skipped_dirs: &[PathBuf]) -> eyre::Result<usize> {
73 let mut count = 0;
74 let entries = std::fs::read_dir(dir)?;
75 for entry in entries {
76 let entry = entry?;
77 let path = entry.path();
78 if skipped_dirs.contains(&path) {
79 continue;
80 }
81 if path.is_dir() {
82 count += brutalize_sol_files_in_dir(&path, skipped_dirs)?;
83 } else if path.extension().is_some_and(|ext| ext == "sol") && !is_test_or_script(&path) {
84 let source = std::fs::read_to_string(&path)?;
85 let brutalized = brutalize_source(&path, &source)?;
86 if brutalized != source {
87 std::fs::write(&path, brutalized)?;
88 count += 1;
89 }
90 }
91 }
92 Ok(count)
93}
94
95fn is_test_or_script(path: &Path) -> bool {
96 path.file_name()
97 .and_then(|name| name.to_str())
98 .is_some_and(|name| name.ends_with(".t.sol") || name.ends_with(".s.sol"))
99}
100
101const fn splitmix64(mut x: u64) -> u64 {
103 x ^= x >> 30;
104 x = x.wrapping_mul(0xbf58476d1ce4e5b9);
105 x ^= x >> 27;
106 x = x.wrapping_mul(0x94d049bb133111eb);
107 x ^= x >> 31;
108 x
109}
110
111fn span_seed(span: Span) -> u64 {
113 let lo = span.lo().0 as u64;
114 let hi = span.hi().0 as u64;
115 splitmix64(lo.wrapping_mul(0x9e3779b97f4a7c15) ^ hi.wrapping_mul(0xff51afd7ed558ccd))
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121
122 fn brutalize(source: &str) -> String {
123 brutalize_source(Path::new("test.sol"), source).unwrap()
124 }
125
126 #[test]
127 fn deterministic_output() {
128 let source = r#"
129pragma solidity ^0.8.0;
130contract T {
131 function f(uint160 x) external pure returns (address) {
132 return address(x);
133 }
134}
135"#;
136 let r1 = brutalize(source);
137 let r2 = brutalize(source);
138 assert_eq!(r1, r2);
139 }
140
141 #[test]
142 fn no_change_without_casts_or_assembly() {
143 let source = r#"
144pragma solidity ^0.8.0;
145contract T {
146 uint256 public x;
147 function set(uint256 v) external {
148 x = v;
149 }
150}
151"#;
152 let result = brutalize(source);
153 assert_eq!(result, source);
154 }
155
156 #[test]
157 fn special_functions_not_brutalized() {
158 for body in [
159 "constructor() { assembly { sstore(0, 1) } }",
160 "fallback() external { assembly { sstore(0, 1) } }",
161 "receive() external payable { assembly { sstore(0, 1) } }",
162 ] {
163 let source = format!("pragma solidity ^0.8.0;\ncontract T {{ {body} }}\n");
164 let result = brutalize(&source);
165 assert!(!result.contains("mstore(0x00,"), "should not inject for: {body}");
166 }
167
168 let free_fn = r#"
169pragma solidity ^0.8.0;
170function freeFunc() pure returns (uint256 r) {
171 assembly { r := 42 }
172}
173"#;
174 let result = brutalize(free_fn);
175 assert!(!result.contains("mstore(0x00,"));
176 }
177
178 #[test]
179 fn assembly_in_nested_control_flow() {
180 for body in [
181 "if (true) { assembly { r := 42 } }",
182 "for (uint256 i; i < 1; i++) { assembly { r := 42 } }",
183 "unchecked { assembly { r := 42 } }",
184 ] {
185 let source = format!(
186 "pragma solidity ^0.8.0;\ncontract T {{\n\
187 function f() external pure returns (uint256 r) {{ {body} }}\n}}\n"
188 );
189 let result = brutalize(&source);
190 assert!(result.contains("mstore(0x00,"), "should inject for: {body}");
191 assert!(
192 result.contains("mstore(0x40, add(mload(0x40),"),
193 "should inject FMP for: {body}"
194 );
195 }
196 }
197
198 #[test]
199 fn cast_and_assembly_in_same_function() {
200 let source = r#"
201pragma solidity ^0.8.0;
202contract T {
203 function f(uint256 x) external pure returns (uint8 r) {
204 r = uint8(x);
205 assembly { r := add(r, 1) }
206 }
207}
208"#;
209 let result = brutalize(source);
210 assert!(result.contains("uint8(uint256(uint8(x)) | (uint256(0x"));
211 assert!(result.contains("mstore(0x00,"));
212 assert!(result.contains("mstore(0x40, add(mload(0x40),"));
213 }
214
215 #[test]
216 fn visibility_gates_injection_not_casts() {
217 let source = r#"
218pragma solidity ^0.8.0;
219contract T {
220 function f(uint256 x) internal pure returns (uint8) {
221 return uint8(x);
222 }
223 function g(uint256 x) public pure returns (uint8) {
224 return uint8(x);
225 }
226}
227"#;
228 let result = brutalize(source);
229 let count = result.matches("uint8(uint256(uint8(x)) | (uint256(0x").count();
230 assert_eq!(count, 2, "casts in internal/public should be brutalized");
231 assert!(!result.contains("mstore(0x00,"), "no memory injection for non-external");
232 }
233}