forge/cmd/
exploit.rs

1use clap::{Parser, ValueHint};
2use eyre::{Context, Result, eyre};
3use foundry_cli::{opts::BuildOpts, utils::LoadConfig};
4use foundry_common::{provider::mpp::transport::LazyAccountsProvider, shell};
5use foundry_compilers::utils::source_files_iter;
6use foundry_config::impl_figment_convert;
7use foundry_evm::constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, MONAD_CHEATCODE_ADDRESS};
8use mpp::client::Fetch;
9use serde::Serialize;
10use serde_json::{Value, json};
11use std::{
12    fmt::Write as _,
13    io::Write as _,
14    path::{Path, PathBuf},
15};
16use url::Url;
17
18use crate::{
19    cmd::properties::{
20        SELECT_ALL_TESTS, Sandbox, TestFile, restrict_config, shorten, stderr, suite_results,
21    },
22    result::{TestResult, TestStatus},
23    traces::{CallKind, TraceKind},
24    workspace::relative_to_root,
25};
26
27/// Tempo Gateway that serves OpenAI-compatible chat completions paid over MPP.
28const DEFAULT_ENDPOINT: &str = "https://gateway.tempo-dev.workers.dev";
29
30/// Model used when none is given.
31const DEFAULT_MODEL: &str = "@cf/zai-org/glm-5.3";
32
33/// Reply that ends an attempt without a finding.
34const NO_EXPLOIT: &str = "NO_EXPLOIT";
35
36/// Longest source file sent to the model.
37const MAX_SOURCE_BYTES: usize = 64_000;
38
39/// CLI arguments for `forge exploit`.
40///
41/// Asks a model served by the Tempo Gateway to find a vulnerability in a contract and prove it
42/// with a Foundry test. Forge compiles and runs every test the model writes and sends compiler
43/// errors and failing results back until a test passes. Requests are paid through the Machine
44/// Payments Protocol from the Tempo account that `cast tempo login` authorized.
45///
46/// The output of every test goes back to the model, so Forge runs the tests as `forge properties`
47/// runs generated tests: in a temporary copy of the project, without `ffi`, network access, RPC
48/// and Etherscan settings, or Forge's environment variables. A test that passes is written to
49/// `test/exploit/` under a name that no file there has yet.
50#[derive(Clone, Debug, Parser)]
51pub struct ExploitArgs {
52    /// The contract source to attack.
53    #[arg(long, value_hint = ValueHint::FilePath, value_name = "PATH")]
54    target: PathBuf,
55
56    /// Extra files for the model, such as a specification or related contracts.
57    #[arg(long, value_hint = ValueHint::FilePath, value_name = "PATH")]
58    context: Vec<PathBuf>,
59
60    /// Gateway serving `/v1/chat/completions` over MPP.
61    #[arg(long, env = "FORGE_EXPLOIT_ENDPOINT", default_value = DEFAULT_ENDPOINT)]
62    endpoint: Url,
63
64    /// Model to ask.
65    #[arg(long, env = "FORGE_EXPLOIT_MODEL", default_value = DEFAULT_MODEL)]
66    model: String,
67
68    /// Independent attempts; each later attempt is told which bugs were already proven.
69    #[arg(long, default_value_t = 3)]
70    attempts: usize,
71
72    /// Model replies per attempt, counting each retry after a compiler error or failing test.
73    #[arg(long, default_value_t = 4)]
74    turns: usize,
75
76    /// Most tokens per model reply.
77    #[arg(long, default_value_t = 8_192)]
78    max_tokens: u32,
79
80    #[command(flatten)]
81    build: BuildOpts,
82}
83
84impl_figment_convert!(ExploitArgs, build);
85
86/// A test the model wrote that passes, proving its claimed bug.
87#[derive(Debug, Serialize)]
88struct Finding {
89    attempt: usize,
90    turns: usize,
91    file: PathBuf,
92    tests: Vec<String>,
93    bugs: Vec<String>,
94}
95
96#[derive(Debug, Serialize)]
97struct Report<'a> {
98    target: &'a Path,
99    endpoint: &'a str,
100    model: &'a str,
101    findings: &'a [Finding],
102}
103
104enum Run {
105    Proven(Vec<String>),
106    Rejected(String),
107}
108
109impl ExploitArgs {
110    pub async fn run(self) -> Result<()> {
111        let mut config = self.load_config()?;
112        restrict_config(&mut config, false, false);
113        let root = config.root.clone();
114        let target = root.join(&self.target);
115        let target_rel = relative_to_root(&root, &target);
116        let target_source = read_source(&target)?;
117        let mut context = String::new();
118        for path in &self.context {
119            let path = root.join(path);
120            write!(
121                context,
122                "\n\n--- {} ---\n{}",
123                relative_to_root(&root, &path).display(),
124                read_source(&path)?
125            )?;
126        }
127        let remappings =
128            config.get_all_remappings().map(|r| r.to_string()).collect::<Vec<_>>().join("\n");
129        let mut sources = source_files_iter(&config.src, &["sol"])
130            .map(|path| relative_to_root(&root, &path).display().to_string())
131            .collect::<Vec<_>>();
132        sources.sort();
133        let sources = sources.join("\n");
134        let mut sandbox = Sandbox::new(&config, vec![])?;
135        sandbox.end_on_interrupt(Vec::new())?;
136        if let Some(errors) = sandbox.build()? {
137            return Err(eyre!(
138                "the project must compile before Forge looks for exploits: {errors}"
139            ));
140        }
141
142        let out_dir = config.test.join("exploit");
143        let chat = Chat::new(&self.endpoint, &self.model, self.max_tokens)?;
144        let mut findings = Vec::<Finding>::new();
145        let mut next = 1;
146
147        for attempt in 1..=self.attempts {
148            // Proofs of earlier runs keep their files: each attempt names its test after the first
149            // number that no file has.
150            while out_dir.join(format!("Exploit{next}Test.t.sol")).exists() {
151                next += 1;
152            }
153            let contract = format!("Exploit{next}Test");
154            let file = out_dir.join(format!("{contract}.t.sol"));
155            let file_rel = relative_to_root(&root, &file);
156            let mut prompt = format!(
157                "Target: {target}\n\n{target_source}{context}\n\nProject sources:\n{sources}\n\n\
158                 Remappings:\n{remappings}\n\nWrite the test file at {file}. Name the test \
159                 contract {contract}.",
160                target = target_rel.display(),
161                file = file_rel.display(),
162            );
163            if !findings.is_empty() {
164                prompt.push_str("\n\nThese bugs are already proven; find a different one:\n");
165                for bug in findings.iter().flat_map(|f| &f.bugs) {
166                    writeln!(prompt, "{bug}")?;
167                }
168            }
169            let mut messages = vec![
170                json!({"role": "system", "content": SYSTEM}),
171                json!({"role": "user", "content": prompt}),
172            ];
173
174            sh_status!("Attempt {attempt}/{}: asking {}", self.attempts, self.model)?;
175            for turn in 1..=self.turns {
176                let reply = chat.complete(&messages).await?;
177                messages.push(json!({"role": "assistant", "content": reply}));
178
179                let Some(code) = solidity_block(&reply) else {
180                    if reply.contains(NO_EXPLOIT) {
181                        sh_status!("Attempt {attempt}: the model found no exploit")?;
182                    } else {
183                        sh_warn!("Attempt {attempt}: the reply had no Solidity code block")?;
184                    }
185                    break;
186                };
187                let feedback = match run_test(&sandbox, &file_rel, code)? {
188                    Run::Proven(tests) => {
189                        std::fs::create_dir_all(&out_dir)?;
190                        std::fs::OpenOptions::new()
191                            .write(true)
192                            .create_new(true)
193                            .open(&file)
194                            .and_then(|mut target| target.write_all(code.as_bytes()))
195                            .wrap_err_with(|| format!("failed to write {}", file_rel.display()))?;
196                        let bugs = bug_tags(code);
197                        sh_status!("Attempt {attempt}: proven by {}", file_rel.display())?;
198                        findings.push(Finding {
199                            attempt,
200                            turns: turn,
201                            file: file_rel.clone(),
202                            tests,
203                            bugs,
204                        });
205                        break;
206                    }
207                    Run::Rejected(feedback) => feedback,
208                };
209                sh_status!("Attempt {attempt}, turn {turn}: not proven yet")?;
210                messages.push(json!({"role": "user", "content": feedback}));
211            }
212        }
213
214        let report_path = config.cache_path.join("exploit").join("report.json");
215        std::fs::create_dir_all(report_path.parent().unwrap())?;
216        let report = Report {
217            target: &target_rel,
218            endpoint: self.endpoint.as_str(),
219            model: &self.model,
220            findings: &findings,
221        };
222        std::fs::write(&report_path, serde_json::to_string_pretty(&report)?)?;
223
224        if shell::is_json() {
225            sh_println!("{}", serde_json::to_string(&report)?)?;
226        } else {
227            for finding in &findings {
228                for bug in &finding.bugs {
229                    sh_println!("{}\t{bug}", finding.file.display())?;
230                }
231            }
232            sh_status!(
233                "{} of {} attempts proved an exploit; report in {}",
234                findings.len(),
235                self.attempts,
236                relative_to_root(&root, &report_path).display()
237            )?;
238        }
239        Ok(())
240    }
241}
242
243const SYSTEM: &str = "You are a smart contract security researcher auditing a Foundry project. \
244Find one concrete, exploitable vulnerability in the target contract and prove it with a Foundry \
245test.
246
247Reply with exactly one solidity fenced code block holding a complete test file:
248- an SPDX line, a pragma compatible with the target, and imports of forge-std/Test.sol and of the \
249contracts you need, using the project's remappings or paths relative to the test file;
250- one contract with the name you are given, inheriting Test, that deploys everything it needs in \
251setUp();
252- one or more functions named test_exploit_* that pass only when the exploit works: assert what \
253the attacker gained or which guarantee broke, never something that holds anyway;
254- each such function successfully calls a function of a target contract that can change its \
255state; view calls and reverted calls alone do not count;
256- directly above each such function, a comment line \
257// @bug Contract.function path:line - one sentence saying what is wrong;
258- no ffi, no forks, no network access and no environment variables.
259
260When the test does not compile or does not pass you get the output; reply with the whole fixed \
261file. If you find no exploitable vulnerability, reply NO_EXPLOIT and nothing else.";
262
263struct Chat {
264    client: reqwest::Client,
265    provider: LazyAccountsProvider,
266    url: Url,
267    model: String,
268    max_tokens: u32,
269}
270
271impl Chat {
272    fn new(endpoint: &Url, model: &str, max_tokens: u32) -> Result<Self> {
273        let mut base = endpoint.clone();
274        if !base.path().ends_with('/') {
275            base.set_path(&format!("{}/", base.path()));
276        }
277        Ok(Self {
278            client: reqwest::Client::new(),
279            provider: LazyAccountsProvider::new(endpoint.to_string()),
280            url: base.join("v1/chat/completions")?,
281            model: model.to_string(),
282            max_tokens,
283        })
284    }
285
286    async fn complete(&self, messages: &[Value]) -> Result<String> {
287        let body = json!({
288            "model": self.model,
289            "messages": messages,
290            "max_tokens": self.max_tokens,
291        });
292        let response = self
293            .client
294            .post(self.url.clone())
295            .json(&body)
296            .send_with_payment(&self.provider)
297            .await
298            .map_err(|e| eyre!("request to {} failed: {e}", self.url))?;
299        let status = response.status();
300        let text = response.text().await?;
301        if !status.is_success() {
302            return Err(eyre!("{} answered {status}: {text}", self.url));
303        }
304        let value: Value =
305            serde_json::from_str(&text).wrap_err_with(|| format!("invalid response: {text}"))?;
306        let message = &value["choices"][0]["message"];
307        message["content"]
308            .as_str()
309            .filter(|content| !content.trim().is_empty())
310            .or_else(|| message["reasoning_content"].as_str())
311            .map(str::to_string)
312            .ok_or_else(|| eyre!("response has no message content: {text}"))
313    }
314}
315
316/// Runs the tests of a model-written file in a fresh copy of the project.
317fn run_test(sandbox: &Sandbox<'_>, file: &Path, code: &str) -> Result<Run> {
318    let test = TestFile { path: file.to_path_buf(), content: code.to_string() };
319    let workspace = sandbox.workspace(&[test])?;
320    let mut command = sandbox.command(workspace.path());
321    command
322        // The traces show which contracts each test calls.
323        .args(["test", "--json", "-vvv", "--no-network", "--restrict-inline-config"])
324        .arg("--match-path")
325        .arg(file)
326        .args(SELECT_ALL_TESTS)
327        // A fuzzed exploit must not pass on inputs that revert.
328        .env("FOUNDRY_FUZZ_FAIL_ON_REVERT", "true");
329    let Some(output) = sandbox.output(&mut command)? else {
330        return Ok(Run::Rejected(
331            "The tests reached the timeout; they must not loop without end.".to_string(),
332        ));
333    };
334    if !output.status.success() && output.stdout.is_empty() {
335        return Ok(Run::Rejected(format!(
336            "The test file did not compile or could not run:\n{}",
337            stderr(&output)
338        )));
339    }
340    let results = suite_results(&output)?;
341    let proven = results
342        .iter()
343        .filter(|(_, name, result)| {
344            result.status == TestStatus::Success
345                && name.starts_with("test_exploit")
346                && changes_state(result)
347        })
348        .map(|(_, name, _)| name.clone())
349        .collect::<Vec<_>>();
350    if !proven.is_empty() {
351        return Ok(Run::Proven(proven));
352    }
353
354    let mut feedback = String::from("No test_exploit_* test passed, so nothing is proven:\n");
355    for (_, name, result) in &results {
356        let _ = write!(feedback, "\n{name}: {:?}", result.status);
357        if let Some(reason) = &result.reason {
358            let _ = write!(feedback, " ({reason})");
359        } else if result.status == TestStatus::Success && !changes_state(result) {
360            feedback.push_str(
361                " (makes no successful call that can change state into another contract, so it \
362                 cannot show the exploit)",
363            );
364        }
365        for log in &result.decoded_logs {
366            let _ = write!(feedback, "\n  {log}");
367        }
368    }
369    if results.is_empty() {
370        feedback.push_str("\nNo tests ran; name the functions test_exploit_*.");
371    }
372    Ok(Run::Rejected(shorten(&feedback)))
373}
374
375/// Returns whether a test makes a successful plain call, which can change the callee's state,
376/// into a contract other than the test, the cheatcodes of any network, and the console. A test
377/// that does not, such as an empty one, one that only reads, or one whose calls all revert, passes
378/// whatever the target does. The tests run in another process, so the network's own list of
379/// cheatcode addresses is not at hand.
380fn changes_state(result: &TestResult) -> bool {
381    result.traces.iter().filter(|(kind, _)| *kind == TraceKind::Execution).any(|(_, arena)| {
382        let nodes = arena.arena.nodes();
383        let test = &nodes[0].trace;
384        test.success
385            && nodes[1..].iter().map(|node| &node.trace).any(|trace| {
386                trace.success
387                    && trace.kind == CallKind::Call
388                    && ![
389                        test.address,
390                        CHEATCODE_ADDRESS,
391                        MONAD_CHEATCODE_ADDRESS,
392                        HARDHAT_CONSOLE_ADDRESS,
393                    ]
394                    .contains(&trace.address)
395            })
396    })
397}
398
399fn solidity_block(reply: &str) -> Option<&str> {
400    let start = reply.find("```solidity").or_else(|| reply.find("```sol"))?;
401    let rest = &reply[start..];
402    let body = &rest[rest.find('\n')? + 1..];
403    let end = body.find("```").unwrap_or(body.len());
404    let code = body[..end].trim();
405    (!code.is_empty()).then_some(code)
406}
407
408fn bug_tags(code: &str) -> Vec<String> {
409    code.lines()
410        .filter_map(|line| line.trim().strip_prefix("//"))
411        .map(str::trim)
412        .filter(|line| line.starts_with("@bug"))
413        .map(str::to_string)
414        .collect()
415}
416
417fn read_source(path: &Path) -> Result<String> {
418    let mut source = std::fs::read_to_string(path)
419        .wrap_err_with(|| format!("failed to read {}", path.display()))?;
420    if source.len() > MAX_SOURCE_BYTES {
421        source.truncate(source.floor_char_boundary(MAX_SOURCE_BYTES));
422        source.push_str("\n// [truncated]");
423    }
424    Ok(source)
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430
431    #[test]
432    fn extracts_solidity_block_and_bug_tags() {
433        let reply = "Here:\n```solidity\n// SPDX-License-Identifier: UNLICENSED\ncontract A {\n    // @bug Vault.withdraw src/Vault.sol:12 - no reentrancy guard\n    function test_exploit_reenter() public {}\n}\n```\ndone";
434        let code = solidity_block(reply).unwrap();
435        assert!(code.starts_with("// SPDX"));
436        assert!(code.ends_with('}'));
437        assert_eq!(bug_tags(code), ["@bug Vault.withdraw src/Vault.sol:12 - no reentrancy guard"]);
438        assert_eq!(solidity_block("NO_EXPLOIT"), None);
439    }
440}
forge/cmd/exploit.rs

forge/cmd/
exploit.rs
