trident_syn/codegen/
trident_flow_executor.rs

1use proc_macro2::TokenStream;
2use quote::{quote, ToTokens};
3
4use crate::types::trident_flow_executor::TridentFlowExecutorImpl;
5
6impl ToTokens for TridentFlowExecutorImpl {
7    fn to_tokens(&self, tokens: &mut TokenStream) {
8        let type_name = &self.type_name;
9        let impl_items = &self.impl_block;
10        let methods = &self.flow_methods;
11        let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();
12
13        // Generate init call if present
14        let init_call = if let Some(init_method) = &self.init_method {
15            quote! {
16                self.#init_method();
17            }
18        } else {
19            quote! {}
20        };
21
22        // Use default_random_flows if there are no flow methods
23        let execute_impl = if methods.is_empty() {
24            quote! {
25                // No flow methods or all are ignored, use default implementation
26                self.default_random_transactions(fuzzer_data, &mut accounts)?;
27            }
28        } else {
29            let random_tail = if self.args.random_tail {
30                quote! {
31                    self.default_random_transactions(fuzzer_data, &mut accounts)?;
32                }
33            } else {
34                quote! {}
35            };
36
37            quote! {
38                // Execute all defined flow methods
39                #(self.#methods(fuzzer_data, &mut accounts)?;)*
40
41                // Optional random tail transactions
42                #random_tail
43            }
44        };
45
46        let expanded = quote! {
47            impl #impl_generics #type_name #ty_generics #where_clause {
48                #(#impl_items)*
49            }
50
51            impl #impl_generics #type_name #ty_generics #where_clause {
52                pub fn execute_flows(
53                    &mut self,
54                    fuzzer_data: &mut FuzzerData,
55                ) -> std::result::Result<(), FuzzingError> {
56                    let mut accounts = FuzzAccounts::default();
57
58                    #init_call
59                    #execute_impl
60                    Ok(())
61                }
62
63                fn default_random_transactions(
64                    &mut self,
65                    fuzzer_data: &mut FuzzerData,
66                    accounts: &mut FuzzAccounts,
67                ) -> std::result::Result<(), FuzzingError> {
68                    let mut transactions = <Vec<FuzzTransactions>>::arbitrary(fuzzer_data)?;
69                    for transaction in transactions.iter_mut() {
70                        transaction.transaction_selector(&mut self.client, accounts)?
71                    }
72                    Ok(())
73                }
74
75                fn fuzz(&mut self) {
76                    if cfg!(honggfuzz) {
77                        loop {
78                            fuzz_honggfuzz(|fuzzer_data| {
79                                let mut buf = Unstructured::new(fuzzer_data);
80                                let _ = self.execute_flows(&mut buf);
81                                self.client.clear_accounts();
82                            });
83                        }
84                    } else if cfg!(afl) {
85                        fuzz_afl(true, |fuzzer_data| {
86                            let mut buf = Unstructured::new(fuzzer_data);
87                            let _ = self.execute_flows(&mut buf);
88                            self.client.clear_accounts();
89                        });
90                    } else if cfg!(honggfuzz_debug) {
91                        let mut crash_file = String::new();
92                        std::io::stdin()
93                            .read_line(&mut crash_file)
94                            .expect("Failed to read crash file path from stdin");
95                        let crash_file = crash_file.trim();
96
97                        let fuzzer_data = std::fs::read(crash_file).expect("Failed to read crash file");
98                        let mut buf = Unstructured::new(&fuzzer_data);
99                        let _ = self.execute_flows(&mut buf);
100                        self.client.clear_accounts();
101                    } else {
102                        panic!("Select Honggfuzz or AFL for fuzzing!!!")
103                    }
104                }
105            }
106        };
107
108        tokens.extend(expanded);
109    }
110}
trident_syn/codegen/trident_flow_executor.rs

trident_syn/codegen/
trident_flow_executor.rs
