seeding/
chunked.rs

1//! Utilities for grouping protocol exchanges
2//!
3//! Takes request-response pairs from a (collected/experimental) ExchangeDataset and breaks them into "chunks" based on protocol-derived featured
4//! Function code and expected response length (if known)
5//! Each chunk corresponds to a function code/length keyed set of values that can be used to seed fuzzing sessions.
6//!
7
8use navigator::protocol::Protocols;
9use std::collections::HashMap;
10use std::fs::File;
11use std::io::{BufWriter, Write};
12use std::path::Path;
13
14use crate::exchange_dataset::ExchangeDataset;
15
16pub type Chunked = HashMap<(usize, usize), Vec<Vec<u8>>>;
17
18/// Group request-response pairs into protocol-aware chunks
19///
20/// If the protocol defines a required response length that is used, else actual response length used as fallback
21/// Returns an error if function code or length cannot be derived
22///
23pub fn chunk(exchange_dataset: ExchangeDataset, protocol: &Protocols) -> Result<Chunked, String> {
24    let func_code = |request: &[u8]| protocol.implementation().function_code(request);
25    let n_fuzz_bytes = |request: &[u8]| {
26        protocol
27            .implementation()
28            .len_of_response_content_required(request)
29    };
30
31    exchange_dataset.into_iter().try_fold(
32        HashMap::new(),
33        |mut acc: Chunked, (request, response)| {
34            // If n_fuzz_bytes is not defined by the protocol implementation, chunk by the response
35            // length.
36            let len_feature = match n_fuzz_bytes(&request)? {
37                Some(known_n_fuzz_bytes) => known_n_fuzz_bytes,
38                None => response.len(),
39            };
40            let key = (func_code(&request)?, len_feature);
41            acc.entry(key)
42                .and_modify(|grouped_responses| grouped_responses.push(response.clone()))
43                .or_insert(vec![response]);
44            Ok(acc)
45        },
46    )
47}
48
49/// Write the chunked dataset to file for use with the fuzzer
50pub fn write_chunked(chunked: Chunked, output_path: &Path) -> Result<(), String> {
51    let file = File::create(output_path).map_err(|e| format!("Unable to create file: {:?}", e))?;
52    let mut writer = BufWriter::new(file);
53
54    for (key, payloads) in chunked {
55        // Serialize the key-value pair as MessagePack
56        let chunk = rmp_serde::to_vec_named(&(key, payloads))
57            .map_err(|e| format!("Unable to serialize chunk: {:?}", e))?;
58
59        let len = chunk.len() as u64;
60        // Write the length as 4 bytes (big-endian)
61        writer
62            .write_all(&len.to_be_bytes())
63            .map_err(|e| format!("Unable to write length: {:?}", e))?;
64        // Write the chunk
65        writer
66            .write_all(&chunk)
67            .map_err(|e| format!("Unable to write chunk: {:?}", e))?;
68    }
69    Ok(())
70}
seeding/chunked.rs

seeding/
chunked.rs
