fuchsia_fuzzctl/
corpus.rs

1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::input::save_input;
6use anyhow::{Context as _, Error, Result};
7use futures::TryStreamExt;
8use std::cell::RefCell;
9use std::path::Path;
10use {fidl_fuchsia_fuzzer as fuzz, zx_status as zx};
11
12/// Returns which type of corpus is represented by the `fuchsia.fuzzer.Corpus` enum.
13///
14/// A seed corpus is immutable. A fuzzer can add or modify inputs in its live corpus.
15pub fn get_type(seed: bool) -> fuzz::Corpus {
16    if seed {
17        fuzz::Corpus::Seed
18    } else {
19        fuzz::Corpus::Live
20    }
21}
22
23/// Get the corresponding name for a `fuchsia.fuzzer.Corpus` enum.
24pub fn get_name(corpus_type: fuzz::Corpus) -> &'static str {
25    match corpus_type {
26        fuzz::Corpus::Seed => "seed",
27        fuzz::Corpus::Live => "live",
28        other => unreachable!("unsupported type: {:?}", other),
29    }
30}
31
32/// Basic corpus information returned by `read`.
33#[derive(Debug, Default, PartialEq)]
34pub struct Stats {
35    pub num_inputs: u64,
36    pub total_size: u64,
37}
38
39/// Receives and saves inputs from a corpus.
40///
41/// Takes a `stream` and serves `fuchsia.fuzzer.CorpusReader`. A fuzzer can publish a sequence of
42/// test inputs using this protocol, typically in response to a `fuchsia.fuzzer.Controller/Fetch`
43/// request or similar. The inputs are saved under `out_dir`, or in the current working directory
44/// if `out_dir` is `None.
45pub async fn read<P: AsRef<Path>>(
46    stream: fuzz::CorpusReaderRequestStream,
47    out_dir: P,
48) -> Result<Stats> {
49    // Without these `RefCell`s, the compiler will complain about references in the async block
50    // below that escape the closure.
51    let num_inputs: RefCell<u64> = RefCell::new(0);
52    let total_size: RefCell<u64> = RefCell::new(0);
53    stream
54        .try_for_each(|request| async {
55            match request {
56                fuzz::CorpusReaderRequest::Next { test_input, responder } => {
57                    {
58                        let mut num_inputs = num_inputs.borrow_mut();
59                        let mut total_size = total_size.borrow_mut();
60                        *num_inputs += 1;
61                        *total_size += test_input.size;
62                    }
63                    let result = match save_input(test_input, out_dir.as_ref()).await {
64                        Ok(_) => zx::Status::OK,
65                        Err(_) => zx::Status::IO,
66                    };
67                    responder.send(result.into_raw())
68                }
69            }
70        })
71        .await
72        .map_err(Error::msg)
73        .context("failed to handle fuchsia.fuzzer.CorpusReader request")?;
74    let num_inputs = num_inputs.borrow();
75    let total_size = total_size.borrow();
76    Ok(Stats { num_inputs: *num_inputs, total_size: *total_size })
77}
78
79#[cfg(test)]
80mod tests {
81    use super::{get_name, get_type, read, Stats};
82    use crate::input::InputPair;
83    use crate::util::digest_path;
84    use anyhow::{Error, Result};
85    use fuchsia_fuzzctl_test::{verify_saved, Test};
86    use futures::join;
87    use {fidl_fuchsia_fuzzer as fuzz, zx_status as zx};
88
89    // Writes a test input using the given `corpus_reader`.
90    async fn send_one_input(corpus_reader: &fuzz::CorpusReaderProxy, data: Vec<u8>) -> Result<()> {
91        let input_pair = InputPair::try_from_data(data)?;
92        let (fidl_input, input) = input_pair.as_tuple();
93        let (response, _) = futures::try_join!(
94            async move { corpus_reader.next(fidl_input).await.map_err(Error::msg) },
95            input.send(),
96        )?;
97        zx::Status::ok(response).map_err(Error::msg)
98    }
99
100    #[test]
101    fn test_get_type() -> Result<()> {
102        assert_eq!(get_type(true), fuzz::Corpus::Seed);
103        assert_eq!(get_type(false), fuzz::Corpus::Live);
104        Ok(())
105    }
106
107    #[test]
108    fn test_get_name() -> Result<()> {
109        assert_eq!(get_name(fuzz::Corpus::Seed), "seed");
110        assert_eq!(get_name(fuzz::Corpus::Live), "live");
111        Ok(())
112    }
113
114    #[fuchsia::test]
115    async fn test_read() -> Result<()> {
116        let test = Test::try_new()?;
117        let corpus_dir = test.create_dir("corpus")?;
118        let corpus = vec![b"hello".to_vec(), b"world".to_vec(), b"".to_vec()];
119        let cloned = corpus.clone();
120
121        let (proxy, stream) =
122            fidl::endpoints::create_proxy_and_stream::<fuzz::CorpusReaderMarker>();
123        let read_fut = read(stream, &corpus_dir);
124        let send_fut = || async move {
125            for input in corpus.iter() {
126                send_one_input(&proxy, input.to_vec()).await?;
127            }
128            Ok::<(), Error>(())
129        };
130        let send_fut = send_fut();
131        let results = join!(read_fut, send_fut);
132        assert_eq!(results.0.ok(), Some(Stats { num_inputs: 3, total_size: 10 }));
133        assert!(results.1.is_ok());
134        for input in cloned.iter() {
135            let saved = digest_path(&corpus_dir, None, input);
136            verify_saved(&saved, input)?;
137        }
138        Ok(())
139    }
140}