Skip to main content

forge/brutalizer/
mod.rs

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
48/// Brutalize all .sol source files in a temp project directory.
49///
50/// Walks the src directory under `temp_dir`, parses each .sol file, applies all
51/// brutalizations (value XOR, memory, FMP), and writes the result back in-place.
52///
53/// Returns the number of files brutalized.
54pub 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
101/// Applies the splitmix64 finalizer to produce a well-distributed 64-bit hash.
102const 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
111/// Derives a deterministic seed from a span's byte offsets.
112fn 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}