phink_lib/fuzzer/
environment.rs

1use crate::{
2    cli::{
3        config::{
4            PFiles::{
5                AllowlistPath,
6                CorpusPath,
7                DictPath,
8            },
9            PhinkFiles,
10        },
11        env::PhinkEnv::AllowList,
12        ziggy::ZiggyConfig,
13    },
14    contract::selectors::{
15        database::SelectorDatabase,
16        selector::Selector,
17    },
18    EmptyResult,
19    ResultOf,
20};
21use anyhow::Context;
22use std::{
23    fs,
24    fs::{
25        File,
26        OpenOptions,
27    },
28    io,
29    io::Write,
30    path::PathBuf,
31};
32
33pub struct AllowListBuilder;
34
35impl AllowListBuilder {
36    pub const FUNCTIONS: [&str; 5] = [
37        "*_ZN9phink_lib6fuzzer6parser*",
38        "*redirect_coverage*",
39        "*decode*",
40        "*harness*",
41        "*black_box*",
42    ];
43    /// Builds the LLVM allowlist if it doesn't already exist.
44    pub fn build(fuzz_output: PathBuf) -> io::Result<()> {
45        let allowlist_path = PhinkFiles::new(fuzz_output).path(AllowlistPath);
46
47        if allowlist_path.exists() {
48            println!("❗ {AllowList} already exists... skipping");
49            return Ok(());
50        }
51
52        fs::create_dir_all(allowlist_path.parent().unwrap())?;
53        let mut allowlist_file = File::create(allowlist_path)?;
54
55        for func in Self::FUNCTIONS {
56            writeln!(allowlist_file, "fun: {}", func)?;
57        }
58
59        println!("✅ {AllowList} created successfully");
60        Ok(())
61    }
62}
63
64pub struct Dict {
65    file_path: PathBuf,
66}
67
68impl Dict {
69    pub fn write_dict_entry(&self, selector: &Selector) -> EmptyResult {
70        let mut file = OpenOptions::new()
71            .append(true)
72            .open(&self.file_path)
73            .with_context(|| format!("Failed to open file for appending: {:?}", self.file_path))?;
74
75        writeln!(file, "\"{}\"", selector)
76            .with_context(|| format!("Couldn't write selector '{}' into the dict", selector))?;
77
78        Ok(())
79    }
80
81    pub fn new(phink_file: PhinkFiles, max_message: usize) -> io::Result<Dict> {
82        let path_buf = phink_file.path(DictPath);
83        if let Some(parent) = path_buf.parent() {
84            fs::create_dir_all(parent)?;
85        }
86        let mut dict_file = OpenOptions::new()
87            .write(true)
88            .create(true)
89            .truncate(true)
90            .open(path_buf.clone())?;
91
92        writeln!(dict_file, "# Dictionary file for selectors")?;
93        writeln!(
94            dict_file,
95            "# Lines starting with '#' and empty lines are ignored."
96        )?;
97
98        // We only add delimiters if we want to fuzz more than one message
99        if max_message > 1 {
100            writeln!(dict_file, "delimiter=\"********\"")?;
101        }
102
103        Ok(Self {
104            file_path: path_buf,
105        })
106    }
107}
108
109#[derive(Clone, Debug)]
110pub struct CorpusManager {
111    corpus_dir: PathBuf,
112}
113
114impl CorpusManager {
115    pub fn new(phink_file: &PhinkFiles) -> ResultOf<CorpusManager> {
116        let corpus_dir = phink_file.path(CorpusPath);
117        if !corpus_dir.exists() {
118            fs::create_dir_all(&corpus_dir)?;
119        }
120        Ok(Self { corpus_dir })
121    }
122
123    /// Write a seed to a given `name`.bin. It pre-append `00000000 01` which is basically
124    /// null-value and Alice as origin
125    pub fn write_seed(&self, name: &str, seed: &[u8]) -> io::Result<()> {
126        let mut data = vec![0x00, 0x00, 0x00, 0x00, 0x01];
127        let file_path = self.corpus_dir.join(format!("{name}.bin"));
128        data.extend_from_slice(seed);
129        fs::write(file_path, data)
130    }
131
132    pub fn write_corpus_file(&self, index: usize, selector: &Selector) -> io::Result<()> {
133        let mut data = vec![0x00, 0x00, 0x00, 0x00, 0x01]; // 00010000 01 fa80c2f6 00
134        let file_path = self.corpus_dir.join(format!("selector_{index}.bin"));
135        data.extend_from_slice(selector.0.as_ref());
136        data.extend(vec![0x0, 0x0]); // Padding for SCALE
137        fs::write(file_path, data)
138    }
139}
140
141pub struct EnvironmentBuilder {
142    database: SelectorDatabase,
143}
144
145impl EnvironmentBuilder {
146    pub fn new(database: SelectorDatabase) -> EnvironmentBuilder {
147        Self { database }
148    }
149
150    /// This function builds both the correct seeds and the dict file for AFL++
151    pub fn build_env(self, conf: ZiggyConfig) -> EmptyResult {
152        let phink_file = PhinkFiles::new(conf.clone().fuzz_output());
153
154        let dict = Dict::new(
155            phink_file.clone(),
156            conf.config().max_messages_per_exec.unwrap_or_default(),
157        )?;
158        let corpus_manager =
159            CorpusManager::new(&phink_file).context("Couldn't create a new corpus manager")?;
160
161        for (i, selector) in self
162            .database
163            .get_unique_messages()
164            .context("Couldn't load messages")?
165            .iter()
166            .enumerate()
167        {
168            corpus_manager
169                .write_corpus_file(i, selector)
170                .context("Couldn't write corpus file")?;
171
172            dict.write_dict_entry(selector)
173                .context("Couldn't write the dictionnary entries")?;
174        }
175
176        Ok(())
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183    use crate::EmptyResult;
184    use std::{
185        fs,
186        io,
187    };
188    use tempfile::tempdir;
189
190    fn create_test_selector() -> Selector {
191        Selector([0x01, 0x02, 0x03, 0x04])
192    }
193
194    fn create_temp_phink_file(file_name: &str) -> PathBuf {
195        let dir = tempdir().unwrap();
196        dir.path().join(file_name)
197    }
198
199    #[test]
200    fn test_dict_new_creates_file() -> io::Result<()> {
201        let path = create_temp_phink_file("test_dict");
202        let phink_file = PhinkFiles::new(path.clone());
203
204        let dict = Dict::new(phink_file, 4)?;
205        assert!(dict.file_path.exists());
206
207        let contents = fs::read_to_string(dict.file_path)?;
208        assert!(contents.contains("# Dictionary file for selectors"));
209
210        Ok(())
211    }
212
213    #[test]
214    fn test_write_dict_entry() -> EmptyResult {
215        let path = create_temp_phink_file("test_dict");
216        let phink_file = PhinkFiles::new(path.clone());
217        let dict = Dict::new(phink_file, 3)?;
218        let selector = create_test_selector(); //        Selector([0x01, 0x02, 0x03, 0x04])
219        dict.write_dict_entry(&selector)?;
220        let contents = fs::read_to_string(dict.file_path)?;
221        assert!(contents.contains("01020304"));
222
223        Ok(())
224    }
225
226    #[test]
227    fn test_corpus_manager_new_creates_dir() -> EmptyResult {
228        let path = create_temp_phink_file("test_corpus");
229        let phink_file = PhinkFiles::new(path.clone());
230
231        let corpus_manager = CorpusManager::new(&phink_file)?;
232        assert!(corpus_manager.corpus_dir.exists());
233
234        Ok(())
235    }
236
237    #[test]
238    fn test_write_corpus_file() -> io::Result<()> {
239        let path = create_temp_phink_file("test_corpus");
240        let phink_file = PhinkFiles::new(path.clone());
241        let corpus_manager = CorpusManager::new(&phink_file).unwrap();
242
243        let selector = create_test_selector();
244        corpus_manager.write_corpus_file(0, &selector)?;
245
246        let file_path = corpus_manager.corpus_dir.join("selector_0.bin");
247        assert!(file_path.exists());
248
249        let data = fs::read(file_path)?;
250        assert_eq!(data[5..9], selector.0);
251
252        Ok(())
253    }
254}
phink_lib/fuzzer/environment.rs

phink_lib/fuzzer/
environment.rs
