├── test_bins ├── test_corpus │ └── 1.txt ├── simple_prog ├── timeout_prog ├── test_config.json ├── simple_prog.c └── timeout_prog.c ├── .gitignore ├── install.sh ├── proto ├── hello_world.proto ├── execution.proto ├── monitor.proto ├── mutation.proto ├── queue_management.proto ├── feedback.proto └── msg_type.proto ├── src ├── queue.rs ├── feedback.rs ├── executor.rs ├── main.rs ├── frontend.rs ├── mutator.rs ├── executor │ ├── pipes.rs │ ├── rpc.rs │ ├── shmem.rs │ └── frontend.rs ├── mutator │ ├── rpc.rs │ └── frontend.rs ├── monitor │ ├── rpc.rs │ ├── output_writer.rs │ └── stats.rs ├── frontend │ ├── toy_frontend.rs │ └── simple_frontend.rs ├── queue │ ├── rpc.rs │ ├── manager.rs │ └── frontend.rs ├── minimizer.rs ├── util.rs ├── feedback │ ├── rpc.rs │ ├── newbit_filter.rs │ └── bitmap_collector.rs ├── monitor.rs ├── datatype.rs └── rpc.rs ├── .github └── workflows │ ├── main.yml │ ├── fmt.yml │ └── cov.yml ├── LICENSE ├── README.md └── Cargo.toml /test_bins/test_corpus/1.txt: -------------------------------------------------------------------------------- 1 | 1234 2 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | Cargo.lock 3 | .cur_input* 4 | -------------------------------------------------------------------------------- /test_bins/simple_prog: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OMH4ck/mufuzz/HEAD/test_bins/simple_prog -------------------------------------------------------------------------------- /test_bins/timeout_prog: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/OMH4ck/mufuzz/HEAD/test_bins/timeout_prog -------------------------------------------------------------------------------- /test_bins/test_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "monitor_addr": "127.0.0.1:11111", 3 | "max_exec": 10, 4 | "core": 4 5 | } 6 | -------------------------------------------------------------------------------- /test_bins/simple_prog.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | int main() { 5 | if(getchar() == 'c' && getchar() == 'a' && getchar() == 'b' ){ 6 | abort(); 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- 1 | curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y 2 | 3 | wget https://github.com/protocolbuffers/protobuf/releases/download/v21.5/protoc-21.5-linux-x86_64.zip && \ 4 | sudo unzip protoc-21.5-linux-x86_64.zip -d /usr/local 5 | -------------------------------------------------------------------------------- /test_bins/timeout_prog.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | int main() { 6 | usleep(100 * 1000); 7 | printf("Sleep done!"); 8 | if(getchar() == 'c' && getchar() == 'a' && getchar() == 'b' ){ 9 | abort(); 10 | } 11 | } 12 | -------------------------------------------------------------------------------- /proto/hello_world.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package helloworld; 4 | 5 | service Greeter { 6 | rpc SayHello(HelloRequest) returns (HelloReply); 7 | } 8 | 9 | message HelloRequest { 10 | string name = 1; 11 | } 12 | 13 | message HelloReply { 14 | string message = 1; 15 | } 16 | -------------------------------------------------------------------------------- /proto/execution.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package microfuzz.execution; 4 | 5 | import "msg_type.proto"; 6 | 7 | service Executor { 8 | rpc ExecuteTestCases(msg_type.TestCases) returns (msg_type.Empty); 9 | rpc GetFeedbacks(msg_type.Number) returns (msg_type.Feedbacks); 10 | } 11 | -------------------------------------------------------------------------------- /proto/monitor.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package microfuzz.monitor; 4 | 5 | import "msg_type.proto"; 6 | 7 | service Monitor { 8 | rpc RegisterService(msg_type.ServiceInfo) returns (msg_type.RegistrationStatus); 9 | rpc GetServices(msg_type.Empty) returns (msg_type.ServiceList); 10 | } 11 | -------------------------------------------------------------------------------- /proto/mutation.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package microfuzz.mutation; 4 | 5 | import "msg_type.proto"; 6 | 7 | service Mutator { 8 | // Post interesting testcases for mutation. 9 | rpc MutateTestCases(msg_type.TestCases) returns (msg_type.Number); 10 | rpc GetMutatedTestCases(msg_type.Number) returns (msg_type.TestCases); 11 | rpc PostMutationFeedbacks(msg_type.Feedbacks) returns (msg_type.Empty); 12 | } 13 | -------------------------------------------------------------------------------- /proto/queue_management.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package microfuzz.queue_management; 4 | 5 | import "msg_type.proto"; 6 | 7 | service QueueManager { 8 | // Get new test cases that we want to test/evaluate in executor. 9 | rpc GetInterestingTestCases(msg_type.Number) returns (msg_type.TestCases); 10 | // Receive the summary of execution feedback. 11 | rpc PostInterestingTestCases(msg_type.TestCases) returns (msg_type.Empty); 12 | 13 | rpc PostTestCaseFeedbacks(msg_type.Feedbacks) returns (msg_type.Empty); 14 | } 15 | -------------------------------------------------------------------------------- /proto/feedback.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package microfuzz.feedback; 4 | 5 | import "msg_type.proto"; 6 | 7 | service FeedbackCollector { 8 | rpc CheckFeedbacks(msg_type.Feedbacks) returns (msg_type.Empty); 9 | rpc GetInterestingTestCases(msg_type.Number) returns (msg_type.TestCases); 10 | rpc GetMutationFeedbacks(msg_type.Number) returns (msg_type.Feedbacks); 11 | rpc GetTestCaseFeedbacks(msg_type.Number) returns (msg_type.Feedbacks); 12 | 13 | rpc RetrieveMonitorData(msg_type.Empty) returns (msg_type.MonitorData); 14 | } 15 | -------------------------------------------------------------------------------- /src/queue.rs: -------------------------------------------------------------------------------- 1 | pub mod frontend; 2 | pub mod manager; 3 | pub mod rpc; 4 | 5 | use crate::datatype::TestCase; 6 | 7 | pub trait QueueManager { 8 | // Check the feedback that tells which test cases are good for exploring the program. 9 | // The queue mananger should utilize this to fine tune their schduling algorithm. 10 | fn process_test_case_feedbacks(&mut self, feedbacks: &[crate::datatype::Feedback]); 11 | 12 | // Save the interesting test cases into the queue. 13 | fn receive_interesting_test_cases(&mut self, test_cases: Vec); 14 | 15 | // Select n test cases for mutation. 16 | fn select_interesting_inputs(&mut self, num: usize) -> Vec; 17 | } 18 | -------------------------------------------------------------------------------- /.github/workflows/main.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: 4 | push: 5 | workflow_dispatch: 6 | 7 | # A workflow run is made up of one or more jobs that can run sequentially or in parallel 8 | jobs: 9 | build: 10 | runs-on: ubuntu-22.04 11 | #runs-on: ${{ matrix.os }} 12 | 13 | # Steps represent a sequence of tasks that will be executed as part of the job 14 | steps: 15 | # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it 16 | - uses: actions/checkout@v3 17 | 18 | - name: Install Prerequisites 19 | run: sh ./install.sh 20 | 21 | - name: add path 22 | run: echo "$HOME/.cargo/bin" >> $GITHUB_PATH 23 | 24 | - name: Run test 25 | run: cargo test 26 | -------------------------------------------------------------------------------- /src/feedback.rs: -------------------------------------------------------------------------------- 1 | pub mod bitmap_collector; 2 | pub mod frontend; 3 | mod newbit_filter; 4 | pub mod rpc; 5 | pub use newbit_filter::NewBitFilter; 6 | 7 | use crate::datatype::Feedback; 8 | 9 | pub trait FeedbackCollector { 10 | fn process_feedback(&mut self, feedbacks: Vec); 11 | 12 | // Get test cases that should be saved in the queue manager. 13 | fn get_interesting_test_cases(&mut self, num: Option) -> Vec; 14 | // Get test cases that cause crashes. 15 | fn get_crash_test_cases(&mut self, num: Option) -> Vec; 16 | // Get feedback about the mutation quality. 17 | fn get_mutation_feedbacks(&mut self, num: Option) -> Vec; 18 | // Get feedback about whether mutating a test case produce good variants. 19 | fn get_test_case_feedbacks(&mut self, num: Option) -> Vec; 20 | } 21 | -------------------------------------------------------------------------------- /.github/workflows/fmt.yml: -------------------------------------------------------------------------------- 1 | name: fmt 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | paths: 8 | - "src/**" 9 | - "tests/**" 10 | - "scripts/**" 11 | - "proto/**" 12 | pull_request: 13 | types: [ opened, synchronize ] 14 | workflow_dispatch: 15 | 16 | # A workflow run is made up of one or more jobs that can run sequentially or in parallel 17 | jobs: 18 | fmt: 19 | runs-on: ubuntu-22.04 20 | 21 | # Steps represent a sequence of tasks that will be executed as part of the job 22 | steps: 23 | # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it 24 | - uses: actions/checkout@v2 25 | 26 | - name: Install Prerequisites 27 | run: sh ./install.sh 28 | 29 | - name: add path 30 | run: echo "$HOME/.cargo/bin" >> $GITHUB_PATH 31 | 32 | - name: Run fmt check 33 | run: cargo fmt --all -- --check 34 | 35 | - name: Run clippy 36 | run: cargo clippy --all-targets --all-features -- -D warnings -------------------------------------------------------------------------------- /src/executor.rs: -------------------------------------------------------------------------------- 1 | use crate::Error; 2 | //use nix::unistd::Pid; 3 | 4 | pub mod forkserver; 5 | pub mod pipes; 6 | pub mod rpc; 7 | pub mod shmem; 8 | 9 | pub use forkserver::BitmapTracer; 10 | pub use forkserver::ForkServerExecutor; 11 | pub mod frontend; 12 | 13 | pub use crate::datatype::ExecutionStatus; 14 | 15 | // Constructor the Feedback after execution. 16 | pub trait Tracer<'a, T> { 17 | fn get_feedback(&'a self) -> T; 18 | } 19 | 20 | // Ideally we only need to specify `I` as input for Executor. 21 | pub trait Executor<'a, T, F, I> { 22 | // This shouldn't be in trait. 23 | fn new( 24 | tracer: T, 25 | args: Vec, 26 | working_dir: Option, 27 | timeout: u64, 28 | bind_to_cpu: bool, 29 | ) -> Result 30 | where 31 | Self: Sized; 32 | // TODO(yongheng): This should be a inner api. 33 | fn run_target(&'a mut self, input: &I) -> Result; 34 | 35 | // Execute the input and produce the feedback for feedback collector. 36 | fn execute(&'a mut self, input: I) -> crate::datatype::Feedback; 37 | } 38 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 OMH4ck 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.github/workflows/cov.yml: -------------------------------------------------------------------------------- 1 | on: 2 | push: 3 | branches: 4 | - main 5 | paths: 6 | - "src/**" 7 | - "tests/**" 8 | - "scripts/**" 9 | - "proto/**" 10 | workflow_dispatch: 11 | 12 | 13 | name: Coverage 14 | 15 | jobs: 16 | check: 17 | name: Rust project 18 | runs-on: ubuntu-latest 19 | steps: 20 | - name: Checkout repository 21 | uses: actions/checkout@v2 22 | 23 | - uses: actions/cache@v2 24 | with: 25 | path: | 26 | ~/.cargo/bin/ 27 | ~/.cargo/registry/index/ 28 | ~/.cargo/registry/cache/ 29 | ~/.cargo/git/db/ 30 | target/ 31 | key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }} 32 | 33 | - name: Install stable toolchain 34 | uses: actions-rs/toolchain@v1 35 | with: 36 | toolchain: stable 37 | override: true 38 | 39 | - name: Install Prerequisites 40 | run: sh ./install.sh 41 | 42 | - name: Run cargo-tarpaulin 43 | uses: actions-rs/tarpaulin@v0.1 44 | with: 45 | version: '0.19.1' 46 | args: '-- --test-threads 1' 47 | 48 | - name: Upload to codecov.io 49 | uses: codecov/codecov-action@v1.0.2 50 | with: 51 | token: ${{secrets.CODECOV_TOKEN}} 52 | 53 | - name: Archive code coverage results 54 | uses: actions/upload-artifact@v1 55 | with: 56 | name: code-coverage-report 57 | path: cobertura.xml 58 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # mufuzz, a parallel fuzzing framework 2 | 3 | [![build](https://github.com/OMH4ck/mufuzz/actions/workflows/main.yml/badge.svg)](https://github.com/OMH4ck/mufuzz/actions/workflows/main.yml) 4 | [![codecov](https://codecov.io/gh/OMH4ck/mufuzz/branch/main/graph/badge.svg?token=S29YN47226)](https://codecov.io/gh/OMH4ck/mufuzz) 5 | [![fmt](https://github.com/OMH4ck/mufuzz/actions/workflows/fmt.yml/badge.svg)](https://github.com/OMH4ck/mufuzz/actions/workflows/fmt.yml) 6 | 7 | The orginal paper can be [here](https://changochen.github.io/publication/mufuzz_usenix_2023.pdf). 8 | 9 | ## Build 10 | 1. Install `cargo` and `protoc` 11 | ``` 12 | curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y 13 | wget https://github.com/protocolbuffers/protobuf/releases/download/v21.5/protoc-21.5-linux-x86_64.zip && \ 14 | sudo unzip protoc-21.5-linux-x86_64.zip -d /usr/local 15 | ``` 16 | 2. `cargo build --release` 17 | 18 | 19 | ## Run 20 | Basic run (for example, run on 20 cores): 21 | ``` 22 | taskset -c 0-19 cargo run --release -- -c "your/command @@" -i input_corpus -o output --core 20 23 | ``` 24 | 25 | Check the usage: 26 | ``` 27 | cargo run --release -- -help 28 | ``` 29 | 30 | ## Note 31 | 1. The code is still under cleanup. Some features are still missing (such as saving the corpus/crashes to disk). 32 | 2. Assign exclusive cores to `mufuzz` for better performance. 33 | 3. Part of the forkserver code is borrowed from [LibAFL](https://github.com/AFLplusplus/LibAFL), a great project for building fuzzers. 34 | 4. After you exit the fuzzer, you might need to run `ipcrm -a` to remove the share memory. 35 | -------------------------------------------------------------------------------- /proto/msg_type.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package microfuzz.msg_type; 4 | 5 | message Number { 6 | uint32 num = 1; 7 | } 8 | 9 | message TestCases { 10 | repeated TestCase test_cases = 1; 11 | } 12 | 13 | message TestCase { 14 | bytes content = 2; 15 | uint32 pid = 3; 16 | } 17 | 18 | message Feedbacks { 19 | repeated Feedback feedbacks = 1; 20 | } 21 | 22 | message NewBit { 23 | uint32 index = 1; 24 | uint32 value = 2; 25 | } 26 | 27 | message NewCoverage { 28 | repeated NewBit new_bits = 1; 29 | } 30 | 31 | message Counter { 32 | uint32 counter = 1; 33 | } 34 | 35 | message FakeFeedback { 36 | string fake = 1; 37 | } 38 | 39 | enum ExecutionStatus { 40 | OK = 0; 41 | TIMEOUT = 1; 42 | CRASH = 2; 43 | INTERESTING = 3; 44 | } 45 | 46 | message MutationInfo { 47 | uint32 pid = 1; 48 | uint32 mutator_id = 2; 49 | } 50 | 51 | message FeedbackData { 52 | oneof feedback_data { 53 | NewCoverage new_coverage = 4; 54 | Counter counter = 5; 55 | FakeFeedback fake_feedback= 6; 56 | } 57 | } 58 | 59 | message Feedback { 60 | TestCase test_case = 1; 61 | ExecutionStatus status = 2; 62 | MutationInfo mutation_info = 3; 63 | FeedbackData data = 4; 64 | } 65 | 66 | message Empty {} 67 | 68 | message RegistrationStatus { 69 | bool success = 1; 70 | } 71 | 72 | enum ServiceType { 73 | MUTATOR = 0; 74 | EXECUTOR = 1; 75 | QUEUEMANAGER = 2; 76 | FEEDBACKCOLLECTOR = 3; 77 | } 78 | 79 | message ServiceInfo { 80 | ServiceType service_type = 1; 81 | string socket_addr = 2; 82 | } 83 | 84 | message ServiceList { 85 | repeated ServiceInfo services = 1; 86 | } 87 | 88 | message MonitorData { 89 | repeated string data = 1; 90 | } -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "mufuzz" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | tonic = { version = "0.8.0", features = ["transport"] } 10 | prost = "0.11.0" 11 | lazy_static = "1.4.0" 12 | clap = { version = "3.2.17", features = ["derive"] } 13 | 14 | tokio = { version = "1.20.1", features = ["macros", "rt-multi-thread"] } 15 | #tokio = { version = "1.17.0", features = ["macros", "rt-multi-thread", "full", "tracing"] } 16 | tokio-stream = { version = "0.1.9", features = ["net"] } 17 | libc = "0.2.132" 18 | nix = { version = "0.25.0", optional = true } 19 | core_affinity = { version = "0.5", git = "https://github.com/s1341/core_affinity_rs", rev = "6648a7a", optional = true } 20 | serde = { version = "1.0.136", default-features = false, features = [ 21 | "alloc", 22 | "derive", 23 | ] } # serialization lib 24 | #quote = "1.0" 25 | #syn = "1.0" 26 | async-trait = "0.1.57" 27 | futures = "0.3.23" 28 | rand = "0.8.5" 29 | tower = "0.4.13" 30 | http = "0.2.8" 31 | serde_json = "1.0.83" 32 | itertools = "0.10.3" 33 | rand_distr = "0.4.3" 34 | async-channel = "1.7.1" 35 | rayon = "1.5.3" 36 | dashmap = "5.3.4" 37 | #crossbeam = "0.8.1" 38 | psutil = "3.2.2" 39 | bitflags = {version ="2.2.1", features = ["serde"] } 40 | 41 | #[lib] 42 | #proc-macro = true 43 | 44 | [build-dependencies] 45 | tonic-build = "0.8.0" 46 | 47 | [dev-dependencies] 48 | serial_test = "0.9.0" 49 | criterion = "0.3.6" 50 | 51 | [features] 52 | default = ["std", "fork"] 53 | std = ["core_affinity", "nix"] # print, env, launcher ... support 54 | fork = [] 55 | 56 | [profile.release] 57 | debug = true 58 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | use clap::Parser; 2 | use mufuzz::monitor::Monitor; 3 | use mufuzz::FuzzerConfig; 4 | use std::error::Error; 5 | use tokio::runtime; 6 | 7 | // In the future we will let users provide a config file instead. 8 | fn main() -> Result<(), Box> { 9 | let fuzzer = FuzzerConfig::parse(); 10 | 11 | if !fuzzer.config.is_empty() { 12 | let fuzzer_stat = mufuzz::monitor::stats::get_fuzzer_info(); 13 | fuzzer_stat 14 | .lock() 15 | .unwrap() 16 | .sync_with_json_config(fuzzer.config.clone()); 17 | } 18 | 19 | match fuzzer.mode { 20 | 2 => { 21 | let rt = runtime::Builder::new_multi_thread() 22 | .worker_threads(fuzzer.core as usize * 6) 23 | .enable_all() 24 | .build()?; 25 | rt.block_on(mufuzz::run_fuzzer_rpc_async_mode(fuzzer)); 26 | } 27 | 3 => { 28 | let rt = runtime::Builder::new_multi_thread() 29 | .worker_threads(fuzzer.core as usize * 5) 30 | .enable_all() 31 | .build()?; 32 | rt.block_on(mufuzz::run_fuzzer_local_async_mode(fuzzer)); 33 | } 34 | 4 => { 35 | mufuzz::monitor::get_monitor() 36 | .write() 37 | .unwrap() 38 | .get_fuzzer_info_mut() 39 | .sync_with_fuzzer_config(&fuzzer); 40 | for i in 0..fuzzer.core { 41 | let config = fuzzer.clone(); 42 | std::thread::spawn(move || mufuzz::run_single_fuzzer_mode(config, i == 0)); 43 | } 44 | 45 | std::thread::spawn(|| { 46 | let monitor = mufuzz::monitor::get_monitor(); 47 | loop { 48 | std::thread::sleep(std::time::Duration::from_millis(2000)); 49 | monitor.read().unwrap().show_statistics(); 50 | } 51 | }) 52 | .join() 53 | .unwrap(); 54 | } 55 | 5 => { 56 | let rt = runtime::Builder::new_multi_thread() 57 | .worker_threads(fuzzer.core as usize * 6) 58 | .enable_all() 59 | .build()?; 60 | rt.block_on(mufuzz::run_fuzzer_local_async_lock_free_mode(fuzzer)); 61 | } 62 | 6 => { 63 | let rt = runtime::Builder::new_multi_thread() 64 | .worker_threads(fuzzer.core as usize * 6) 65 | .enable_all() 66 | .build()?; 67 | rt.block_on(mufuzz::run_fuzzer_local_work_stealing_mode(fuzzer)); 68 | } 69 | _ => { 70 | unreachable!(); 71 | } 72 | } 73 | Ok(()) 74 | } 75 | -------------------------------------------------------------------------------- /src/frontend.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::TestCaseID; 2 | use crate::datatype::{Feedback, TestCase}; 3 | use serde_json::Value; 4 | 5 | pub mod load_balance_frontend; 6 | pub mod simple_frontend; 7 | mod toy_frontend; 8 | pub mod work_stealing_frontend; 9 | 10 | // The basic unit controlled by the Frontend. 11 | pub trait Worker: Send { 12 | type Input: Clone + Send + 'static; 13 | type Output: Send + 'static; 14 | 15 | // The main function for Worker. Accept some inputs and produce some results. 16 | fn handle_one_input(&mut self, input: Vec) -> Vec; 17 | 18 | // Other than producing some results, a worker might also produce some extra side data that 19 | // we want to keep. For example, the executor might tell how many test cases it has tried. 20 | fn retrieve_monitor_data(&mut self) -> Vec { 21 | unreachable!(); 22 | } 23 | } 24 | 25 | macro_rules! io_type_marcro { 26 | ($($type_name:tt($data_type:ty)), *) => { 27 | #[derive(Clone, Debug)] 28 | pub enum FuzzerIO { 29 | $($type_name($data_type)), * 30 | } 31 | 32 | #[derive(Debug, Clone, Copy, std::hash::Hash, PartialEq, Eq, std::cmp::Ord, std::cmp::PartialOrd)] 33 | pub enum FuzzerIOType { 34 | $($type_name), * 35 | } 36 | 37 | impl FuzzerIO { 38 | pub fn get_type(&self) -> FuzzerIOType { 39 | match self{ 40 | $(FuzzerIO::$type_name(_) => FuzzerIOType::$type_name),* 41 | } 42 | } 43 | 44 | pub fn len(&self) -> usize { 45 | match self{ 46 | $(FuzzerIO::$type_name(data) => data.len()),* 47 | } 48 | 49 | } 50 | 51 | pub fn is_empty(&self) -> bool { 52 | match self{ 53 | $(FuzzerIO::$type_name(data) => data.is_empty()),* 54 | } 55 | 56 | } 57 | 58 | pub fn get_default(io_type: FuzzerIOType) -> FuzzerIO { 59 | match io_type { 60 | $(FuzzerIOType::$type_name => FuzzerIO::$type_name(Vec::default())),* 61 | } 62 | } 63 | } 64 | 65 | }; 66 | } 67 | 68 | // For each element we will generate a enum variant in FuzzerIO and FuzzerIOType 69 | // respectively. For example, we will have 70 | // FuzzerIO::TestCase(_).get_type() = FuzzerIOType::TestCase 71 | io_type_marcro! { 72 | TestCase(Vec), 73 | Feedback(Vec), 74 | MutationFeedback(Vec), 75 | TestCaseFeedback(Vec), 76 | MonitorData(Vec), 77 | TestCaseScoreChange(Vec<(TestCaseID, i64)>), 78 | MutatorScoreChange(Vec<(u32, i64)>) 79 | } 80 | -------------------------------------------------------------------------------- /src/mutator.rs: -------------------------------------------------------------------------------- 1 | pub mod afl_mutator; 2 | pub mod frontend; 3 | pub mod rpc; 4 | 5 | use crate::datatype::{Feedback, TestCase}; 6 | 7 | type MutationRng = rand::rngs::StdRng; 8 | /* 9 | pub struct MutatorFunc { 10 | pub func: fn(&TestCase, &mut MutationRng) -> Option, 11 | pub score: i64, 12 | pub id: u32, 13 | } 14 | 15 | impl MutatorFunc { 16 | pub fn new( 17 | id: u32, 18 | func: fn(&TestCase, &mut MutationRng) -> Option, 19 | score: i64, 20 | ) -> Self { 21 | Self { func, score, id } 22 | } 23 | 24 | pub fn mutate_once( 25 | &self, 26 | test_case: &TestCase, 27 | rng: &mut MutationRng, 28 | mutator_id: u32, 29 | ) -> Option { 30 | if let Some(mut test_case) = (self.func)(test_case, rng) { 31 | let mid = (mutator_id << 16) + self.id; 32 | test_case.set_mutator_id(mid as u32); 33 | Some(test_case) 34 | } else { 35 | None 36 | } 37 | } 38 | } 39 | */ 40 | 41 | pub trait Mutator { 42 | //fn get_mutator_funcs(&self) -> &Vec; 43 | 44 | // Each mutator should have its own unique id. 45 | // TODO: We should automatically assign this id. 46 | fn get_id(&self) -> u32; 47 | 48 | fn get_rng(&mut self) -> &mut MutationRng; 49 | //fn get_dist(&self) -> &WeightedAliasIndex; 50 | //fn pick_mutator_function(&mut self) -> &MutatorFunc; 51 | fn get_result_pool(&mut self) -> &mut Vec; 52 | fn mutate(&mut self, test_case: &TestCase, round: usize); 53 | /* { 54 | if test_case.get_buffer().is_empty() { 55 | return; 56 | } 57 | 58 | let mut rng = self.get_rng().clone(); 59 | let id = self.get_id(); 60 | for _i in 0..round { 61 | let mutator_func = self.pick_mutator_function(); 62 | if let Some(test_case) = mutator_func.mutate_once(test_case, &mut rng, id) { 63 | self.get_result_pool().push(test_case); 64 | } 65 | } 66 | } 67 | */ 68 | 69 | fn get_mutated_test_cases(&mut self, num: Option) -> Vec { 70 | let n_new_test_case = self.get_result_pool().len(); 71 | match num { 72 | Some(sz) => { 73 | let mut nsize = sz as usize; 74 | nsize = nsize.min(n_new_test_case); 75 | self.get_result_pool().split_off(n_new_test_case - nsize) 76 | } 77 | None => self.get_result_pool().split_off(0), 78 | } 79 | } 80 | 81 | /// Check the feedbacks about the mutation and fine-tune the scheduling of mutator functions. 82 | fn process_mutation_feedback(&mut self, feedbacks: &[Feedback]); 83 | } 84 | -------------------------------------------------------------------------------- /src/executor/pipes.rs: -------------------------------------------------------------------------------- 1 | //! Unix `pipe` wrapper for `LibAFL` 2 | use crate::Error; 3 | //#[cfg(feature = "std")] 4 | use nix::unistd::{close, pipe, read, write}; 5 | //#[cfg(feature = "std")] 6 | use std::{ 7 | io::{self, ErrorKind, Read, Write}, 8 | os::unix::io::RawFd, 9 | }; 10 | 11 | //#[cfg(not(feature = "std"))] 12 | //type RawFd = i32; 13 | 14 | //#[cfg(feature = "std")] 15 | #[derive(Debug, Clone)] 16 | pub struct Pipe { 17 | read_end: Option, 18 | write_end: Option, 19 | } 20 | 21 | //#[cfg(feature = "std")] 22 | impl Pipe { 23 | pub fn new() -> Result { 24 | let (read_end, write_end) = pipe()?; 25 | Ok(Self { 26 | read_end: Some(read_end), 27 | write_end: Some(write_end), 28 | }) 29 | } 30 | 31 | pub fn close_read_end(&mut self) { 32 | if let Some(read_end) = self.read_end { 33 | let _ = close(read_end); 34 | self.read_end = None; 35 | } 36 | } 37 | 38 | pub fn close_write_end(&mut self) { 39 | if let Some(write_end) = self.write_end { 40 | let _ = close(write_end); 41 | self.write_end = None; 42 | } 43 | } 44 | 45 | #[must_use] 46 | pub fn read_end(&self) -> Option { 47 | self.read_end 48 | } 49 | 50 | #[must_use] 51 | pub fn write_end(&self) -> Option { 52 | self.write_end 53 | } 54 | } 55 | 56 | #[cfg(feature = "std")] 57 | impl Read for Pipe { 58 | /// Reads a few bytes 59 | fn read(&mut self, buf: &mut [u8]) -> Result { 60 | match self.read_end { 61 | Some(read_end) => match read(read_end, buf) { 62 | Ok(res) => Ok(res), 63 | Err(e) => Err(io::Error::from_raw_os_error(e as i32)), 64 | }, 65 | None => Err(io::Error::new( 66 | ErrorKind::BrokenPipe, 67 | "Read pipe end was already closed", 68 | )), 69 | } 70 | } 71 | } 72 | 73 | #[cfg(feature = "std")] 74 | impl Write for Pipe { 75 | /// Writes a few bytes 76 | fn write(&mut self, buf: &[u8]) -> Result { 77 | match self.write_end { 78 | Some(write_end) => match write(write_end, buf) { 79 | Ok(res) => Ok(res), 80 | Err(e) => Err(io::Error::from_raw_os_error(e as i32)), 81 | }, 82 | None => Err(io::Error::new( 83 | ErrorKind::BrokenPipe, 84 | "Write pipe end was already closed", 85 | )), 86 | } 87 | } 88 | 89 | fn flush(&mut self) -> Result<(), io::Error> { 90 | Ok(()) 91 | } 92 | } 93 | 94 | //#[cfg(feature = "std")] 95 | impl Drop for Pipe { 96 | fn drop(&mut self) { 97 | if let Some(read_end) = self.read_end { 98 | let _ = close(read_end); 99 | } 100 | if let Some(write_end) = self.write_end { 101 | let _ = close(write_end); 102 | } 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /src/mutator/rpc.rs: -------------------------------------------------------------------------------- 1 | use crate::frontend::simple_frontend::{AsyncFrontend, BasicFrontend}; 2 | use crate::mutator::frontend::BitFlipMutatorFrontend; 3 | use crate::rpc::msg_type::{Empty, Feedbacks, Number, TestCase, TestCases}; 4 | use crate::rpc::mutation::mutator_server::{Mutator, MutatorServer}; 5 | use std::net::SocketAddr; 6 | use tonic::{Request, Response, Status}; 7 | 8 | pub use crate::rpc::mutation::mutator_client::MutatorClient; 9 | pub struct MyMutator { 10 | frontend: BitFlipMutatorFrontend, 11 | } 12 | 13 | impl MyMutator { 14 | async fn new(worker_num: u32) -> Self { 15 | let mut frontend = BitFlipMutatorFrontend::default(); 16 | frontend.create_worker(worker_num); 17 | frontend.run().await; 18 | MyMutator { frontend } 19 | } 20 | } 21 | 22 | #[tonic::async_trait] 23 | impl Mutator for MyMutator { 24 | async fn mutate_test_cases( 25 | &self, 26 | request: Request, 27 | ) -> Result, Status> { 28 | let request = request.into_inner(); 29 | 30 | let num = request.test_cases.len(); 31 | 32 | self.frontend 33 | .handle_inputs(request.test_cases.into_iter().map(TestCase::into).collect()) 34 | .await; 35 | 36 | Ok(Response::new(Number { num: num as u32 })) 37 | } 38 | 39 | async fn get_mutated_test_cases( 40 | &self, 41 | request: Request, 42 | ) -> Result, Status> { 43 | let num = request.into_inner().num; 44 | let test_cases = self.frontend.get_results(Some(num)).await; 45 | let test_cases = TestCases { 46 | test_cases: test_cases.into_iter().map(TestCase::from).collect(), 47 | }; 48 | 49 | Ok(Response::new(test_cases)) 50 | } 51 | 52 | async fn post_mutation_feedbacks( 53 | &self, 54 | request: Request, 55 | ) -> Result, Status> { 56 | let feedbacks = request 57 | .into_inner() 58 | .feedbacks 59 | .into_iter() 60 | .map(|feedback| feedback.into()) 61 | .collect(); 62 | self.frontend.process_feedbacks(feedbacks).await; 63 | Ok(Response::new(Empty::default())) 64 | } 65 | } 66 | 67 | impl MyMutator { 68 | pub async fn create_service(worker_num: u32) -> MutatorServer { 69 | MutatorServer::new(MyMutator::new(worker_num).await) 70 | } 71 | 72 | pub async fn run_service(server_addr: Option, worker_num: u32) -> SocketAddr { 73 | crate::rpc::run_service_in_background(Self::create_service(worker_num).await, server_addr) 74 | .await 75 | } 76 | } 77 | 78 | #[cfg(test)] 79 | mod tests { 80 | use super::*; 81 | 82 | #[tokio::test] 83 | async fn mutation_rpc_test() { 84 | let addr = MyMutator::run_service(None, 2).await; 85 | 86 | let mut client = MutatorClient::connect(format!("http://{}", addr)) 87 | .await 88 | .unwrap(); 89 | 90 | let mut test_cases = TestCases::default(); 91 | test_cases.test_cases.push(TestCase { 92 | content: (0..255).collect::>(), 93 | pid: 0, 94 | }); 95 | test_cases.test_cases.push(TestCase { 96 | content: (0..255).rev().collect::>(), 97 | pid: 0, 98 | }); 99 | let request = tonic::Request::new(test_cases); 100 | 101 | let response = client 102 | .mutate_test_cases(request) 103 | .await 104 | .unwrap() 105 | .into_inner(); 106 | assert!(response.num == 2); 107 | 108 | let test_case_num = 500; 109 | let request = tonic::Request::new(Number { num: test_case_num }); 110 | 111 | let response = client 112 | .get_mutated_test_cases(request) 113 | .await 114 | .unwrap() 115 | .into_inner(); 116 | println!("{}", response.test_cases.len()); 117 | assert!(response.test_cases.len() > 90); 118 | } 119 | } 120 | -------------------------------------------------------------------------------- /src/monitor/rpc.rs: -------------------------------------------------------------------------------- 1 | use crate::rpc::monitor::monitor_server::{Monitor, MonitorServer}; 2 | use crate::rpc::msg_type::{Empty, RegistrationStatus, ServiceInfo, ServiceList}; 3 | use std::net::SocketAddr; 4 | use tonic::{Request, Response, Status}; 5 | 6 | #[derive(Debug, Default)] 7 | pub struct MySimpleMonitor {} 8 | 9 | #[tonic::async_trait] 10 | impl Monitor for MySimpleMonitor { 11 | async fn register_service( 12 | &self, 13 | request: Request, 14 | ) -> Result, Status> { 15 | println!("Got a request: {:?}", request); 16 | let service_info = request.into_inner(); 17 | let success = match service_info.into() { 18 | Some(service_info) => crate::monitor::SimpleMonitor::register_service(service_info), 19 | None => false, 20 | }; 21 | Ok(Response::new(RegistrationStatus { success })) 22 | } 23 | 24 | async fn get_services(&self, request: Request) -> Result, Status> { 25 | println!("Got a request: {:?}", request); 26 | 27 | let service_list = crate::monitor::SimpleMonitor::get_services(); 28 | Ok(Response::new(ServiceList { 29 | services: service_list.into_iter().map(ServiceInfo::from).collect(), 30 | })) 31 | } 32 | } 33 | 34 | impl MySimpleMonitor { 35 | pub fn create_service() -> MonitorServer { 36 | MonitorServer::new(MySimpleMonitor::default()) 37 | } 38 | 39 | pub async fn run_service(server_addr: Option) -> SocketAddr { 40 | crate::rpc::run_service_in_background(Self::create_service(), server_addr).await 41 | } 42 | } 43 | 44 | #[cfg(test)] 45 | mod tests { 46 | use super::*; 47 | use crate::rpc::monitor::monitor_client::MonitorClient; 48 | use crate::rpc::msg_type::{Empty, ServiceInfo, ServiceType}; 49 | 50 | #[tokio::test] 51 | async fn monitor_rpc_test() { 52 | let addr = MySimpleMonitor::run_service(None).await; 53 | 54 | let mut client = MonitorClient::connect(format!("http://{}", addr)) 55 | .await 56 | .unwrap(); 57 | 58 | let request = tonic::Request::new(ServiceInfo { 59 | service_type: ServiceType::Feedbackcollector as i32, 60 | socket_addr: "127.0.0.1:22233".to_string(), 61 | }); 62 | let response = client.register_service(request).await.unwrap().into_inner(); 63 | assert!(response.success); 64 | 65 | let request = tonic::Request::new(Empty::default()); 66 | let response = client.get_services(request).await.unwrap().into_inner(); 67 | println!("{:#?}", response); 68 | assert_eq!(response.services.len(), 1); 69 | 70 | let request = tonic::Request::new(ServiceInfo { 71 | service_type: ServiceType::Mutator as i32, 72 | socket_addr: "127.0.0.1:22234".to_string(), 73 | }); 74 | let response = client.register_service(request).await.unwrap().into_inner(); 75 | assert!(response.success); 76 | 77 | let request = tonic::Request::new(ServiceInfo { 78 | service_type: ServiceType::Queuemanager as i32, 79 | socket_addr: "127.0.0.1:22235".to_string(), 80 | }); 81 | let response = client.register_service(request).await.unwrap().into_inner(); 82 | assert!(response.success); 83 | 84 | let request = tonic::Request::new(ServiceInfo { 85 | service_type: ServiceType::Executor as i32, 86 | socket_addr: "127.0.0.1:22236".to_string(), 87 | }); 88 | let response = client.register_service(request).await.unwrap().into_inner(); 89 | assert!(response.success); 90 | 91 | // A service with duplicated socket address should not be registered. 92 | let request = tonic::Request::new(ServiceInfo { 93 | service_type: ServiceType::Executor as i32, 94 | socket_addr: "127.0.0.1:22236".to_string(), 95 | }); 96 | let response = client.register_service(request).await.unwrap().into_inner(); 97 | assert!(!response.success); 98 | 99 | let request = tonic::Request::new(Empty::default()); 100 | let response = client.get_services(request).await.unwrap().into_inner(); 101 | println!("{:#?}", response); 102 | assert_eq!(response.services.len(), 4); 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /src/frontend/toy_frontend.rs: -------------------------------------------------------------------------------- 1 | use crate::frontend::*; 2 | use crate::simple_frontend::*; 3 | use async_trait::async_trait; 4 | use std::sync::{Arc, Mutex}; 5 | 6 | #[derive(Default)] 7 | struct ToyWorker { 8 | pub handled_workload: usize, 9 | } 10 | 11 | struct ToyFrontend { 12 | worker_pool: Vec>>, 13 | sender: Option>>, 14 | output_receiver: Option>>, 15 | output_sender: Option>>, 16 | } 17 | 18 | impl Worker for ToyWorker { 19 | type Input = u32; 20 | type Output = u32; 21 | fn handle_one_input(&mut self, input: Vec) -> Vec { 22 | self.handled_workload += input.len(); 23 | input.into_iter().map(|x| x * x).collect() 24 | } 25 | } 26 | 27 | impl ToyFrontend { 28 | #[allow(dead_code)] 29 | fn new() -> Self { 30 | ToyFrontend { 31 | worker_pool: Vec::default(), 32 | sender: None, 33 | output_receiver: None, 34 | output_sender: None, 35 | } 36 | } 37 | } 38 | 39 | impl BasicFrontend for ToyFrontend { 40 | type Worker = ToyWorker; 41 | type Output = u32; 42 | 43 | fn create_worker(&mut self, num: u32) { 44 | for _i in 0..num { 45 | self.worker_pool 46 | .push(Arc::new(Mutex::new(ToyWorker::default()))); 47 | } 48 | } 49 | 50 | fn output_transfrom(&self, input: Vec) -> Vec { 51 | input 52 | } 53 | 54 | crate::frontend_default!(); 55 | } 56 | 57 | impl AsyncFrontendChannel for ToyFrontend { 58 | crate::async_frontend_default!(); 59 | } 60 | 61 | #[async_trait] 62 | impl AsyncFrontend for ToyFrontend {} 63 | 64 | #[cfg(test)] 65 | mod tests { 66 | use super::*; 67 | use tokio::time::{sleep, Duration}; 68 | 69 | #[tokio::test] 70 | async fn frontend_evenly_distribute_workload_to_workers() { 71 | let mut toy_frontend = ToyFrontend::new(); 72 | toy_frontend.create_worker(10); 73 | toy_frontend.run().await; 74 | let mut join_handle = Vec::new(); 75 | for _j in 0..10 { 76 | join_handle.push(async { 77 | let mut inputs = Vec::new(); 78 | for i in 0..10 { 79 | inputs.push(i); 80 | } 81 | toy_frontend.handle_inputs(inputs).await; 82 | }); 83 | } 84 | futures::future::join_all(join_handle).await; 85 | sleep(Duration::from_millis(10)).await; 86 | 87 | for worker in toy_frontend.get_worker_pool() { 88 | let worker = worker.lock().unwrap(); 89 | assert_eq!(worker.handled_workload, 10); 90 | } 91 | } 92 | 93 | #[tokio::test] 94 | async fn frontend_gather_output_from_all_workers() { 95 | let mut toy_frontend = ToyFrontend::new(); 96 | toy_frontend.create_worker(10); 97 | toy_frontend.run().await; 98 | let mut join_handle = Vec::new(); 99 | for _j in 0..10 { 100 | join_handle.push(async { 101 | let mut inputs = Vec::new(); 102 | for i in 0..10 { 103 | inputs.push(i); 104 | } 105 | toy_frontend.handle_inputs(inputs).await; 106 | }); 107 | } 108 | futures::future::join_all(join_handle).await; 109 | sleep(Duration::from_millis(10)).await; 110 | 111 | let result = toy_frontend.get_results(None).await; 112 | assert_eq!(result.len(), 100); 113 | } 114 | 115 | #[test] 116 | fn forkserver_executor_frontend_add_delete_worker_correctly() { 117 | let mut toy_frontend = ToyFrontend::new(); 118 | toy_frontend.create_worker(15); 119 | 120 | assert!(toy_frontend.worker_pool.len() == 15); 121 | toy_frontend.delete_worker(5); 122 | assert!(toy_frontend.worker_pool.len() == 10); 123 | } 124 | 125 | #[tokio::test] 126 | #[should_panic] 127 | async fn fronted_panic_at_handling_inputs_when_no_worker_is_created() { 128 | let toy_frontend = ToyFrontend::new(); 129 | toy_frontend.handle_inputs(vec![0, 1]).await; 130 | } 131 | 132 | #[tokio::test] 133 | #[should_panic] 134 | async fn fronted_panic_at_handling_inputs_when_it_is_not_run() { 135 | let mut toy_frontend = ToyFrontend::new(); 136 | toy_frontend.create_worker(10); 137 | toy_frontend.handle_inputs(vec![0, 1]).await; 138 | } 139 | } 140 | -------------------------------------------------------------------------------- /src/queue/rpc.rs: -------------------------------------------------------------------------------- 1 | use crate::frontend::simple_frontend::{AsyncFrontend, BasicFrontend}; 2 | use crate::queue::frontend::SimpleQueueManagerFrontend; 3 | use crate::rpc::msg_type::{Empty, Feedbacks, Number, TestCase, TestCases}; 4 | use crate::rpc::queue_management::queue_manager_server::{QueueManager, QueueManagerServer}; 5 | use std::net::SocketAddr; 6 | use tonic::{Request, Response, Status}; 7 | 8 | pub use crate::rpc::queue_management::queue_manager_client::QueueManagerClient; 9 | 10 | pub struct MyQueueManager { 11 | frontend: SimpleQueueManagerFrontend, 12 | } 13 | 14 | impl MyQueueManager { 15 | async fn new(worker_num: u32) -> Self { 16 | let mut frontend = SimpleQueueManagerFrontend::default(); 17 | frontend.create_worker(worker_num); 18 | frontend.run().await; 19 | Self { frontend } 20 | } 21 | } 22 | 23 | #[tonic::async_trait] 24 | impl QueueManager for MyQueueManager { 25 | async fn get_interesting_test_cases( 26 | &self, 27 | request: Request, 28 | ) -> Result, Status> { 29 | let request = request.into_inner(); 30 | let test_cases = self.frontend.get_results(Some(request.num)).await; 31 | Ok(Response::new(TestCases { 32 | test_cases: test_cases.into_iter().map(TestCase::from).collect(), 33 | })) 34 | } 35 | 36 | async fn post_interesting_test_cases( 37 | &self, 38 | request: Request, 39 | ) -> Result, Status> { 40 | let test_cases = request 41 | .into_inner() 42 | .test_cases 43 | .into_iter() 44 | .map(|test_case| test_case.into()) 45 | .collect(); 46 | self.frontend.handle_inputs(test_cases).await; 47 | Ok(Response::new(Empty::default())) 48 | } 49 | 50 | async fn post_test_case_feedbacks( 51 | &self, 52 | request: Request, 53 | ) -> Result, Status> { 54 | let feedbacks = request 55 | .into_inner() 56 | .feedbacks 57 | .into_iter() 58 | .map(|feedback| feedback.into()) 59 | .collect(); 60 | self.frontend.process_feedbacks(feedbacks).await; 61 | Ok(Response::new(Empty::default())) 62 | } 63 | } 64 | 65 | impl MyQueueManager { 66 | pub async fn create_service(worker_num: u32) -> QueueManagerServer { 67 | QueueManagerServer::new(MyQueueManager::new(worker_num).await) 68 | } 69 | 70 | pub async fn run_service(server_addr: Option, worker_num: u32) -> SocketAddr { 71 | crate::rpc::run_service_in_background(Self::create_service(worker_num).await, server_addr) 72 | .await 73 | } 74 | } 75 | 76 | #[cfg(test)] 77 | mod tests { 78 | use super::*; 79 | use crate::datatype; 80 | use crate::rpc::msg_type::{self, Number}; 81 | use crate::rpc::queue_management::queue_manager_client::QueueManagerClient; 82 | 83 | #[tokio::test] 84 | async fn queue_manager_rpc_test() { 85 | let addr = MyQueueManager::run_service(None, 2).await; 86 | 87 | let mut client = QueueManagerClient::connect(format!("http://{}", addr)) 88 | .await 89 | .unwrap(); 90 | 91 | let test_cases_str = [ 92 | "select 1;", 93 | "create table v(c); select c from v;", 94 | "select printf(a, b);", 95 | ]; 96 | 97 | let request = tonic::Request::new(TestCases { 98 | test_cases: test_cases_str 99 | .iter() 100 | .map(|s| { 101 | msg_type::TestCase::from(datatype::TestCase::new(s.as_bytes().to_vec(), 0)) 102 | }) 103 | .collect(), 104 | }); 105 | 106 | let response = client.post_interesting_test_cases(request).await; 107 | assert!(response.is_ok()); 108 | 109 | let request = tonic::Request::new(Number { num: 3 }); 110 | let response = client 111 | .get_interesting_test_cases(request) 112 | .await 113 | .unwrap() 114 | .into_inner(); 115 | assert_eq!(response.test_cases.len(), 3); 116 | 117 | let request = tonic::Request::new(Number { num: 4 }); 118 | let response = client 119 | .get_interesting_test_cases(request) 120 | .await 121 | .unwrap() 122 | .into_inner(); 123 | assert_eq!(response.test_cases.len(), 4); 124 | } 125 | } 126 | -------------------------------------------------------------------------------- /src/minimizer.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::TestCase; 2 | use crate::executor::shmem::ShMem; 3 | use crate::executor::BitmapTracer; 4 | use crate::executor::ExecutionStatus; 5 | use crate::executor::Executor; 6 | use crate::executor::ForkServerExecutor; 7 | use std::collections::hash_map::DefaultHasher; 8 | use std::hash::{Hash, Hasher}; 9 | 10 | pub struct TestCaseMinimizer { 11 | executor: ForkServerExecutor, 12 | } 13 | 14 | impl TestCaseMinimizer { 15 | pub fn new(args: Vec, map_size: usize, timeout: u64) -> Self { 16 | let bitmap_tracer = BitmapTracer::new(map_size); 17 | bitmap_tracer.bitmap.write_to_env("__AFL_SHM_ID").unwrap(); 18 | 19 | let executor = ForkServerExecutor::new(bitmap_tracer, args, None, timeout, false) 20 | .ok() 21 | .unwrap(); 22 | 23 | TestCaseMinimizer { executor } 24 | } 25 | 26 | fn get_tracer_bitmap_hash(&self) -> u64 { 27 | let mut s = DefaultHasher::new(); 28 | self.executor.tracer.bitmap.map().hash(&mut s); 29 | s.finish() 30 | } 31 | 32 | fn next_p2(&self, val: usize) -> usize { 33 | let mut ret = 1; 34 | while val > ret { 35 | ret <<= 1; 36 | } 37 | ret 38 | } 39 | 40 | fn keep_same_state_and_bitmap( 41 | &mut self, 42 | testcase: &TestCase, 43 | save_exec_result: &ExecutionStatus, 44 | save_hash: u64, 45 | ) -> bool { 46 | let exec_result = self.executor.run_target(testcase); 47 | assert!(exec_result.is_ok()); 48 | if exec_result.unwrap() != *save_exec_result { 49 | return false; 50 | } 51 | let exec_hash = self.get_tracer_bitmap_hash(); 52 | self.executor.tracer.clear(); 53 | exec_hash == save_hash 54 | } 55 | 56 | // Minimize the length of the test case while keeping the coverage same. 57 | pub fn minimize(&mut self, testcase: TestCase) -> TestCase { 58 | const TRIM_START_STEPS: usize = 2; 59 | const TRIM_MIN_BYTES: usize = 1; 60 | const TRIM_END_STEPS: usize = 1024; 61 | let save_exec_result = self.executor.run_target(&testcase); 62 | assert!(save_exec_result.is_ok()); 63 | let save_exec_result = save_exec_result.unwrap(); 64 | 65 | let save_hash = self.get_tracer_bitmap_hash(); 66 | self.executor.tracer.clear(); 67 | 68 | let mut result = testcase; 69 | let mut len_p2 = self.next_p2(result.len()); 70 | let mut remove_len = (len_p2 / TRIM_START_STEPS).max(TRIM_MIN_BYTES); 71 | while remove_len >= (len_p2 / TRIM_END_STEPS).max(TRIM_MIN_BYTES) { 72 | let mut remove_pos: usize = remove_len; 73 | while remove_pos < result.len() { 74 | let trim_avail = remove_len.min(result.len() - remove_pos); 75 | let mut tmp_testcase = result.clone(); 76 | tmp_testcase 77 | .get_buffer_mut() 78 | .drain(remove_pos..remove_pos + trim_avail); 79 | if self.keep_same_state_and_bitmap(&tmp_testcase, &save_exec_result, save_hash) { 80 | result = tmp_testcase; 81 | len_p2 = self.next_p2(result.len()); 82 | } else { 83 | remove_pos += remove_len; 84 | } 85 | } 86 | remove_len >>= 1; 87 | } 88 | result 89 | } 90 | } 91 | 92 | #[cfg(test)] 93 | mod tests { 94 | use super::*; 95 | use crate::util; 96 | use serial_test::serial; 97 | 98 | #[test] 99 | #[serial] 100 | fn test_case_minimizer_remove_redundant_bytes() { 101 | const MAP_SIZE: usize = 65536; 102 | let args = vec![util::get_test_bin_by_name("simple_prog").unwrap()]; 103 | let mut minimizer = TestCaseMinimizer::new(args, MAP_SIZE, 20); 104 | 105 | let mut buffer = vec![b'f'; 10]; 106 | buffer[0] = b'c'; 107 | buffer[1] = b'a'; 108 | let testcase = TestCase::new(buffer, 0); 109 | 110 | let minimized_testcase = minimizer.minimize(testcase); 111 | assert_eq!(minimized_testcase.len(), 10); 112 | //assert_eq!(*minimized_testcase.get_buffer(), vec![b'c', b'a']); 113 | 114 | let mut buffer = vec![b'f'; 1000]; 115 | buffer[0] = b'c'; 116 | buffer[1] = b'a'; 117 | buffer[2] = b'b'; 118 | let testcase = TestCase::new(buffer, 0); 119 | let minimized_testcase = minimizer.minimize(testcase); 120 | // Triggering EOF will make the edge different. SO we have a 'f' at the end. 121 | assert!(minimized_testcase.len() < 10); 122 | assert_eq!(minimized_testcase.get_buffer()[0..3], [b'c', b'a', b'b']); 123 | } 124 | } 125 | -------------------------------------------------------------------------------- /src/executor/rpc.rs: -------------------------------------------------------------------------------- 1 | use crate::executor::frontend::ForkServerExecutorFrontend; 2 | use crate::frontend::simple_frontend::{AsyncFrontend, BasicFrontend}; 3 | use crate::rpc::execution::executor_server::{Executor, ExecutorServer}; 4 | use crate::rpc::msg_type::{Empty, Feedback, Feedbacks, Number, TestCases}; 5 | use std::net::SocketAddr; 6 | use std::sync::Arc; 7 | use tokio::sync::RwLock; 8 | use tonic::{Request, Response, Status}; 9 | 10 | pub use crate::rpc::execution::executor_client::ExecutorClient; 11 | 12 | pub struct MyExecutor { 13 | frontend: Arc>, 14 | //frontend: Arc, 15 | //frontend: ForkServerExecutorFrontend, 16 | } 17 | 18 | impl MyExecutor { 19 | async fn new( 20 | worker_num: u32, 21 | bitmap_size: usize, 22 | args: Vec, 23 | timeout: u64, 24 | working_dir: Option, 25 | ) -> Self { 26 | let mut frontend = ForkServerExecutorFrontend::new(bitmap_size, args, timeout, working_dir); 27 | frontend.create_worker(worker_num); 28 | frontend.run().await; 29 | Self { 30 | //frontend, 31 | //frontend: Arc::new(frontend), 32 | frontend: Arc::new(RwLock::new(frontend)), 33 | } 34 | } 35 | } 36 | 37 | #[tonic::async_trait] 38 | impl Executor for MyExecutor { 39 | async fn get_feedbacks(&self, request: Request) -> Result, Status> { 40 | let request = request.into_inner(); 41 | let num = request.num; 42 | let frontend = self.frontend.clone(); 43 | let feedbacks = tokio::task::spawn(async move { 44 | let feedbacks = frontend.read().await.get_results(Some(num)).await; 45 | //let feedbacks = self.frontend.read().unwrap().get_results(Some(num)).await; 46 | feedbacks 47 | .into_iter() 48 | .map(Feedback::from) 49 | .collect::>() 50 | }) 51 | .await 52 | .unwrap(); 53 | Ok(Response::new(Feedbacks { feedbacks })) 54 | //Ok(Response::new(Feedbacks::default())) 55 | } 56 | 57 | async fn execute_test_cases( 58 | &self, 59 | request: Request, 60 | ) -> Result, Status> { 61 | let request = request.into_inner(); 62 | let inputs = request 63 | .test_cases 64 | .into_iter() 65 | .map(|test_case| test_case.into()) 66 | .collect::>(); 67 | let frontend = self.frontend.clone(); 68 | frontend.read().await.handle_inputs(inputs).await; 69 | Ok(Response::new(Empty::default())) 70 | } 71 | } 72 | 73 | impl MyExecutor { 74 | pub async fn create_service( 75 | worker_num: u32, 76 | bitmap_size: usize, 77 | args: Vec, 78 | timeout: u64, 79 | working_dir: Option, 80 | ) -> ExecutorServer { 81 | ExecutorServer::new( 82 | MyExecutor::new(worker_num, bitmap_size, args, timeout, working_dir).await, 83 | ) 84 | } 85 | 86 | pub async fn run_service( 87 | server_addr: Option, 88 | worker_num: u32, 89 | args: Vec, 90 | map_size: usize, 91 | timeout: u64, 92 | working_dir: Option, 93 | ) -> SocketAddr { 94 | crate::rpc::run_service_in_background( 95 | Self::create_service(worker_num, map_size, args, timeout, working_dir).await, 96 | server_addr, 97 | ) 98 | .await 99 | } 100 | } 101 | 102 | #[cfg(test)] 103 | mod tests { 104 | use super::*; 105 | use crate::rpc::execution::executor_client::ExecutorClient; 106 | use crate::rpc::msg_type::{Number, TestCase, TestCases}; 107 | use crate::util; 108 | 109 | #[tokio::test] 110 | #[serial_test::serial] 111 | async fn executor_rpc_test() { 112 | const MAP_SIZE: usize = 65536; 113 | let args = vec![util::get_test_bin_by_name("simple_prog").unwrap()]; 114 | let addr = MyExecutor::run_service(None, 1, args, MAP_SIZE, 20, None).await; 115 | 116 | let mut client = ExecutorClient::connect(format!("http://{}", addr)) 117 | .await 118 | .unwrap(); 119 | 120 | let test_cases = ["select 1;", "create table v2;"] 121 | .iter() 122 | .map(|str| TestCase { 123 | content: str.as_bytes().to_vec(), 124 | pid: 0, 125 | }) 126 | .collect(); 127 | 128 | let request = tonic::Request::new(TestCases { test_cases }); 129 | let response = client.execute_test_cases(request).await; 130 | assert!(response.is_ok()); 131 | 132 | let request = tonic::Request::new(Number { num: 2 }); 133 | let feedback = client.get_feedbacks(request).await.unwrap().into_inner(); 134 | assert_eq!(feedback.feedbacks.len(), 2); 135 | println!("{:#?}", feedback.feedbacks); 136 | } 137 | } 138 | -------------------------------------------------------------------------------- /src/util.rs: -------------------------------------------------------------------------------- 1 | use super::datatype::{self, TestCase}; 2 | use core_affinity::CoreId; 3 | use lazy_static::lazy_static; 4 | use std::fs; 5 | use std::fs::File; 6 | use std::io; 7 | use std::io::prelude::*; 8 | use std::path::PathBuf; 9 | use std::sync::{Arc, Mutex}; 10 | use std::time::Instant; 11 | 12 | #[cfg(target_os = "linux")] 13 | use libc::{cpu_set_t, sched_setaffinity, CPU_SET}; 14 | #[cfg(target_os = "linux")] 15 | use std::mem; 16 | 17 | #[cfg(not(target_os = "linux"))] 18 | pub fn bind_to_cpu(_pid: i32) -> bool { 19 | true 20 | } 21 | 22 | // TODO: Make this portable. 23 | #[cfg(target_os = "linux")] 24 | pub fn bind_to_cpu(pid: i32) -> bool { 25 | let mut free_cpus = FREE_CPUS.lock().unwrap(); 26 | if free_cpus.is_empty() { 27 | println!("No free cpu left"); 28 | return false; 29 | } 30 | let core_id = free_cpus.pop().unwrap().id; 31 | // Turn `core_id` into a `libc::cpu_set_t` with only 32 | // one core active. 33 | let mut set = unsafe { mem::zeroed::() }; 34 | 35 | unsafe { CPU_SET(core_id, &mut set) }; 36 | 37 | // Set the current thread's core affinity. 38 | unsafe { 39 | sched_setaffinity( 40 | pid, // Defaults to current thread 41 | mem::size_of::(), 42 | &set, 43 | ); 44 | } 45 | println!("Bind to cpu {}", core_id); 46 | true 47 | } 48 | 49 | lazy_static! { 50 | pub static ref FREE_CPUS: Arc>> = 51 | Arc::new(Mutex::new(core_affinity::get_core_ids().unwrap())); 52 | pub static ref AFFINITY_NUM: usize = FREE_CPUS.lock().unwrap().len(); 53 | } 54 | 55 | pub fn pretty_format_num(num: u64) -> String { 56 | num.to_string() 57 | /* 58 | if num < 10_000 { 59 | num.to_string() 60 | } else if num < 10_000_000 { 61 | format!("{}K", num / 1000) 62 | } else { 63 | format!("{}M", num / 1000_0000) 64 | } 65 | */ 66 | } 67 | 68 | pub fn get_test_bin_path() -> String { 69 | PathBuf::from(env!("CARGO_MANIFEST_DIR")) 70 | .to_str() 71 | .unwrap() 72 | .to_string() 73 | + "/test_bins/" 74 | } 75 | 76 | pub fn get_test_bin_by_name(bin: &str) -> Option { 77 | let mut path_str = PathBuf::from(env!("CARGO_MANIFEST_DIR")) 78 | .to_str() 79 | .unwrap() 80 | .to_string() 81 | + "/test_bins/" 82 | + bin; 83 | 84 | if !cfg!(target_os = "linux") { 85 | path_str += "_osx"; 86 | } 87 | 88 | if std::path::Path::new(&path_str).is_file() { 89 | Some(path_str) 90 | } else { 91 | None 92 | } 93 | } 94 | 95 | pub fn read_corpus(path_str: String) -> io::Result> { 96 | let mut result = Vec::new(); 97 | let paths = fs::read_dir(path_str).unwrap(); 98 | for path in paths { 99 | let path = path.unwrap().path(); 100 | let path = path.to_str().unwrap(); 101 | 102 | let mut f = File::open(path)?; 103 | let mut buffer = Vec::new(); 104 | 105 | f.read_to_end(&mut buffer)?; 106 | 107 | result.push(datatype::TestCase::new(buffer, 0)); 108 | } 109 | Ok(result) 110 | } 111 | 112 | pub fn get_tcpdump_corpus() -> Vec { 113 | let corpus_path = get_test_bin_path() + "/test_corpus"; 114 | read_corpus(corpus_path).unwrap() 115 | } 116 | 117 | pub enum Component { 118 | QueueManager, 119 | Executor, 120 | FeedbackCollector, 121 | } 122 | 123 | pub struct FuzzTimer { 124 | queue_manager_time: u128, 125 | executor_time: u128, 126 | feedback_collector_time: u128, 127 | now: Instant, 128 | } 129 | 130 | impl std::default::Default for FuzzTimer { 131 | fn default() -> Self { 132 | FuzzTimer { 133 | queue_manager_time: 0, 134 | executor_time: 0, 135 | feedback_collector_time: 0, 136 | now: Instant::now(), 137 | } 138 | } 139 | } 140 | 141 | impl FuzzTimer { 142 | pub fn start(&mut self) { 143 | self.now = Instant::now(); 144 | } 145 | 146 | pub fn add_time(&mut self, t: Component) { 147 | match t { 148 | Component::QueueManager => self.queue_manager_time += self.now.elapsed().as_millis(), 149 | Component::Executor => self.executor_time += self.now.elapsed().as_millis(), 150 | Component::FeedbackCollector => { 151 | self.feedback_collector_time += self.now.elapsed().as_millis() 152 | } 153 | } 154 | } 155 | 156 | pub fn show(&self) { 157 | let queue_manager_time = self.queue_manager_time as f64; 158 | let executor_time = self.executor_time as f64; 159 | let feedback_time = self.feedback_collector_time as f64; 160 | let total = (queue_manager_time + executor_time + feedback_time).max(1.0); 161 | println!( 162 | "Total {:.2}, executor: {:.2}%, feedback: {:.2}%, queue manager: {:.2}% ", 163 | total, 164 | executor_time / total * 100.0, 165 | feedback_time / total * 100.0, 166 | queue_manager_time / total * 100.0 167 | ); 168 | } 169 | } 170 | 171 | pub fn take_n_elements_from_vector(v: &mut Vec, num: Option) -> Vec { 172 | let max_result_len = v.len(); 173 | let n = match num { 174 | Some(num) => num.min(max_result_len), 175 | _ => max_result_len, 176 | }; 177 | 178 | v.split_off(max_result_len - n) 179 | } 180 | 181 | #[cfg(test)] 182 | mod tests { 183 | use super::*; 184 | 185 | #[test] 186 | fn read_all_corpus() { 187 | let corpus_path = get_test_bin_path() + "/test_corpus"; 188 | let corpus = read_corpus(corpus_path); 189 | assert!(corpus.is_ok()); 190 | let corpus = corpus.unwrap(); 191 | for test_case in corpus { 192 | println!("{:?}", test_case); 193 | } 194 | } 195 | } 196 | -------------------------------------------------------------------------------- /src/feedback/rpc.rs: -------------------------------------------------------------------------------- 1 | use crate::feedback::frontend::BitmapCollectorFrontend; 2 | use crate::frontend::simple_frontend::{AsyncFrontend, BasicFrontend}; 3 | use crate::rpc::feedback::feedback_collector_server::{FeedbackCollector, FeedbackCollectorServer}; 4 | use crate::rpc::msg_type::{self, Empty, Feedbacks, MonitorData, Number, TestCases}; 5 | use std::net::SocketAddr; 6 | use tonic::{Request, Response, Status}; 7 | 8 | pub use crate::rpc::feedback::feedback_collector_client::FeedbackCollectorClient; 9 | 10 | pub struct MyFeedbackCollector { 11 | frontend: BitmapCollectorFrontend, 12 | } 13 | 14 | impl MyFeedbackCollector { 15 | pub async fn new(worker_num: u32, bitmap_size: usize) -> Self { 16 | let mut frontend = BitmapCollectorFrontend::new(bitmap_size); 17 | frontend.create_worker(worker_num); 18 | frontend.run().await; 19 | Self { frontend } 20 | } 21 | } 22 | 23 | #[tonic::async_trait] 24 | impl FeedbackCollector for MyFeedbackCollector { 25 | async fn retrieve_monitor_data( 26 | &self, 27 | _request: Request, 28 | ) -> Result, Status> { 29 | let monitor_data = self.frontend.retrieve_monitor_data().await; 30 | Ok(Response::new(MonitorData::from(monitor_data))) 31 | } 32 | 33 | async fn check_feedbacks( 34 | &self, 35 | request: Request, 36 | ) -> Result, Status> { 37 | let feedbacks = request 38 | .into_inner() 39 | .feedbacks 40 | .into_iter() 41 | .map(|feedback| feedback.into()) 42 | .collect(); 43 | self.frontend.handle_inputs(feedbacks).await; 44 | Ok(Response::new(Empty::default())) 45 | } 46 | 47 | async fn get_test_case_feedbacks( 48 | &self, 49 | request: Request, 50 | ) -> Result, Status> { 51 | let number = request.into_inner().num; 52 | let results = self.frontend.get_test_case_feedbacks(Some(number)).await; 53 | Ok(Response::new(msg_type::Feedbacks { 54 | feedbacks: results.into_iter().map(msg_type::Feedback::from).collect(), 55 | })) 56 | } 57 | 58 | async fn get_mutation_feedbacks( 59 | &self, 60 | request: Request, 61 | ) -> Result, Status> { 62 | let number = request.into_inner().num; 63 | let results = self.frontend.get_mutation_feedbacks(Some(number)).await; 64 | Ok(Response::new(msg_type::Feedbacks { 65 | feedbacks: results.into_iter().map(msg_type::Feedback::from).collect(), 66 | })) 67 | } 68 | 69 | async fn get_interesting_test_cases( 70 | &self, 71 | request: Request, 72 | ) -> Result, Status> { 73 | let number = request.into_inner().num; 74 | let results = self.frontend.get_results(Some(number)).await; 75 | Ok(Response::new(msg_type::TestCases { 76 | test_cases: results.into_iter().map(msg_type::TestCase::from).collect(), 77 | })) 78 | } 79 | } 80 | 81 | impl MyFeedbackCollector { 82 | pub async fn create_service( 83 | worker_num: u32, 84 | bitmap_size: usize, 85 | ) -> FeedbackCollectorServer { 86 | FeedbackCollectorServer::new(MyFeedbackCollector::new(worker_num, bitmap_size).await) 87 | } 88 | 89 | pub async fn run_service( 90 | server_addr: Option, 91 | worker_num: u32, 92 | map_size: usize, 93 | ) -> SocketAddr { 94 | crate::rpc::run_service_in_background( 95 | Self::create_service(worker_num, map_size).await, 96 | server_addr, 97 | ) 98 | .await 99 | } 100 | } 101 | 102 | #[cfg(test)] 103 | mod tests { 104 | use super::*; 105 | use crate::datatype; 106 | use crate::datatype::ExecutionStatus; 107 | use crate::rpc::feedback::feedback_collector_client::FeedbackCollectorClient; 108 | use crate::rpc::msg_type::{Feedback, Feedbacks, Number, TestCase}; 109 | 110 | #[tokio::test] 111 | async fn feedback_rpc_test() { 112 | let addr = MyFeedbackCollector::run_service(None, 2, 65536).await; 113 | 114 | let mut client = FeedbackCollectorClient::connect(format!("http://{}", addr)) 115 | .await 116 | .unwrap(); 117 | 118 | let test_case = TestCase::default(); 119 | let mut feedbacks = Feedbacks::default(); 120 | let feedback = datatype::Feedback::create_coverage_feedback( 121 | ExecutionStatus::Interesting, 122 | Some(test_case.clone().into()), 123 | Some(vec![datatype::NewBit::new(0, 1)]), 124 | ) 125 | .set_mutation_info(datatype::MutationInfo::new(0, 0)); 126 | 127 | feedbacks.feedbacks.push(Feedback::from(feedback.clone())); 128 | feedbacks.feedbacks.push(Feedback::from(feedback)); 129 | let feedback = datatype::Feedback::create_coverage_feedback( 130 | ExecutionStatus::Interesting, 131 | Some(test_case.clone().into()), 132 | Some(vec![datatype::NewBit::new(1, 1)]), 133 | ) 134 | .set_mutation_info(datatype::MutationInfo::new(0, 0)); 135 | 136 | feedbacks.feedbacks.push(Feedback::from(feedback)); 137 | let request = tonic::Request::new(feedbacks); 138 | let response = client.check_feedbacks(request).await; 139 | assert!(response.is_ok()); 140 | println!("go check"); 141 | 142 | let request = tonic::Request::new(Number { num: 2 }); 143 | let response: TestCases = client 144 | .get_interesting_test_cases(request) 145 | .await 146 | .unwrap() 147 | .into_inner(); 148 | 149 | println!("{:?}", response); 150 | let interesing_test_case_counter = response.test_cases.len(); 151 | assert_eq!(interesing_test_case_counter, 2); 152 | } 153 | } 154 | -------------------------------------------------------------------------------- /src/monitor/output_writer.rs: -------------------------------------------------------------------------------- 1 | use super::stats::FuzzerInfo; 2 | use crate::datatype::TestCase; 3 | use std::fs; 4 | use std::fs::File; 5 | use std::io::prelude::*; 6 | use std::sync::{Arc, Mutex}; 7 | 8 | /// Save the fuzzing process/statistics on disk 9 | pub struct OutputWriter { 10 | root_dir: String, 11 | crash_dir: String, 12 | queue_dir: String, 13 | hang_dir: String, 14 | fuzzer_stat_path: String, 15 | plot_data_file: Arc>, // We need to write this file often, so keep a file handler. 16 | } 17 | 18 | impl OutputWriter { 19 | pub fn new(root_dir: String) -> std::io::Result { 20 | fs::create_dir_all(root_dir.clone())?; 21 | let crash_dir = format!("{}/crash", root_dir); 22 | let queue_dir = format!("{}/queue", root_dir); 23 | let hang_dir = format!("{}/hang", root_dir); 24 | let plot_data_path = format!("{}/plot_data", root_dir); 25 | let fuzzer_stat_path = format!("{}/fuzzer_stat", root_dir); 26 | fs::create_dir(crash_dir.clone())?; 27 | fs::create_dir(queue_dir.clone())?; 28 | fs::create_dir(hang_dir.clone())?; 29 | let plot_data_file = Arc::new(Mutex::new(File::create(plot_data_path)?)); 30 | Ok(Self { 31 | crash_dir, 32 | queue_dir, 33 | hang_dir, 34 | root_dir, 35 | plot_data_file, 36 | fuzzer_stat_path, 37 | }) 38 | } 39 | 40 | fn create_file_name_for_test_case(prefix: &str, test_case: &TestCase) -> String { 41 | format!( 42 | "{}/{}_{}", 43 | prefix, 44 | test_case.get_id(), 45 | test_case.get_mutator_id() 46 | ) 47 | } 48 | 49 | fn save_test_case(file_name: String, test_case: &TestCase) -> std::io::Result<()> { 50 | let mut file = File::create(file_name)?; 51 | file.write_all(&test_case.get_buffer()[..])?; 52 | Ok(()) 53 | } 54 | 55 | pub fn save_crash(&self, test_case: &TestCase) -> std::io::Result<()> { 56 | let file_name = Self::create_file_name_for_test_case(&self.crash_dir, test_case); 57 | Self::save_test_case(file_name, test_case) 58 | } 59 | 60 | pub fn save_hang(&self, test_case: &TestCase) -> std::io::Result<()> { 61 | let file_name = Self::create_file_name_for_test_case(&self.hang_dir, test_case); 62 | Self::save_test_case(file_name, test_case) 63 | } 64 | 65 | pub fn save_queue(&self, test_case: &TestCase) -> std::io::Result<()> { 66 | let file_name = Self::create_file_name_for_test_case(&self.queue_dir, test_case); 67 | Self::save_test_case(file_name, test_case) 68 | } 69 | 70 | pub fn write_plot_data(&self, fuzzer_stat: &FuzzerInfo) -> std::io::Result<()> { 71 | let data = format!( 72 | "{}, {}, {}, {}\n", 73 | fuzzer_stat.get_fuzzing_time(), 74 | fuzzer_stat.get_exec(), 75 | fuzzer_stat.get_timeout_exec(), 76 | fuzzer_stat.get_crash(), 77 | ); 78 | self.plot_data_file 79 | .lock() 80 | .unwrap() 81 | .write_all(data.as_bytes())?; 82 | Ok(()) 83 | } 84 | 85 | pub fn write_fuzzer_stat(&self, fuzzer_stat: &FuzzerInfo) -> std::io::Result<()> { 86 | let data = format!( 87 | "{}, {}, {}, {}\n", 88 | fuzzer_stat.get_fuzzing_time(), 89 | fuzzer_stat.get_exec(), 90 | fuzzer_stat.get_timeout_exec(), 91 | fuzzer_stat.get_crash(), 92 | ); 93 | let mut file = fs::File::create(self.fuzzer_stat_path.clone())?; 94 | file.write_all(data.as_bytes())?; 95 | Ok(()) 96 | } 97 | 98 | pub fn get_root_path(&self) -> String { 99 | self.root_dir.clone() 100 | } 101 | } 102 | 103 | #[cfg(test)] 104 | mod tests { 105 | use super::*; 106 | use std::path::Path; 107 | 108 | #[test] 109 | fn output_writer_create_dir_when_it_does_not_exist() { 110 | let output_dir = "./test_output1".to_string(); 111 | let output_writer = OutputWriter::new(output_dir.clone()); 112 | assert!(output_writer.is_ok()); 113 | assert!(Path::new(&output_dir).is_dir()); 114 | assert!(Path::new(&format!("{}/queue", output_dir)).is_dir()); 115 | assert!(Path::new(&format!("{}/crash", output_dir)).is_dir()); 116 | assert!(Path::new(&format!("{}/hang", output_dir)).is_dir()); 117 | assert!(Path::new(&format!("{}/plot_data", output_dir)).is_file()); 118 | 119 | drop(output_writer); 120 | assert!(fs::remove_dir_all(output_dir).is_ok()); 121 | } 122 | 123 | #[test] 124 | fn output_writer_panick_when_the_dir_exist() { 125 | let output_dir = "./test_output2".to_string(); 126 | let output_writer = OutputWriter::new(output_dir.clone()); 127 | assert!(output_writer.is_ok()); 128 | 129 | let output_writer = OutputWriter::new(output_dir.clone()); 130 | assert!(output_writer.is_err()); 131 | drop(output_writer); 132 | 133 | assert!(fs::remove_dir_all(output_dir).is_ok()); 134 | } 135 | 136 | #[test] 137 | fn output_writer_save_file_in_disk() { 138 | let output_dir = "./test_output3".to_string(); 139 | let output_writer = OutputWriter::new(output_dir.clone()).unwrap(); 140 | 141 | let mut test_case = TestCase::new(vec![0xff; 100], 0); 142 | test_case.set_mutator_id(1); 143 | 144 | output_writer.save_crash(&test_case).unwrap(); 145 | 146 | let file_name = OutputWriter::create_file_name_for_test_case( 147 | &format!("{}/crash", output_dir), 148 | &test_case, 149 | ); 150 | 151 | assert!(Path::new(&file_name).is_file()); 152 | 153 | let mut f = File::open(file_name).unwrap(); 154 | let mut buffer = Vec::new(); 155 | 156 | f.read_to_end(&mut buffer).unwrap(); 157 | 158 | assert_eq!(buffer.len(), 100); 159 | 160 | assert_eq!(buffer.iter().filter(|&x| *x == 0xff).count(), 100); 161 | 162 | assert!(fs::remove_dir_all(output_dir).is_ok()); 163 | } 164 | } 165 | -------------------------------------------------------------------------------- /src/monitor.rs: -------------------------------------------------------------------------------- 1 | use lazy_static::lazy_static; 2 | use stats::FuzzerInfo; 3 | use std::collections::HashSet; 4 | use std::net::SocketAddr; 5 | use std::net::{IpAddr, Ipv4Addr}; 6 | use std::sync::{Arc, Mutex, RwLock}; 7 | 8 | use crate::datatype::TestCase; 9 | use crate::executor::ExecutionStatus; 10 | 11 | use self::output_writer::OutputWriter; 12 | 13 | pub mod output_writer; 14 | pub mod rpc; 15 | pub mod stats; 16 | 17 | #[derive(Clone, Copy)] 18 | pub enum Service { 19 | Mutator(SocketAddr), 20 | FeedbackCollector(SocketAddr), 21 | Executor(SocketAddr), 22 | QueueManager(SocketAddr), 23 | } 24 | 25 | pub trait Monitor { 26 | fn get_services(&self) -> Vec; 27 | fn register_service(&self, service: Service) -> bool; 28 | // Print the fuzzer status. 29 | fn show_statistics(&self); 30 | fn reset_start_time(&mut self); 31 | 32 | // Receive statistics from other components. To make this general we use json type. 33 | fn receive_statistics(&self, v: serde_json::Value); 34 | } 35 | 36 | pub struct SimpleMonitor { 37 | services: Arc>>, 38 | #[allow(dead_code)] 39 | addr: SocketAddr, 40 | service_addr_hashes: Arc>>, 41 | fuzzer_info: FuzzerInfo, 42 | } 43 | 44 | lazy_static! { 45 | pub static ref SIMPLE_MONITOR: Arc> = 46 | Arc::new(RwLock::new(SimpleMonitor::default())); 47 | } 48 | 49 | static SIMPLE_MONITOR_PORT: u16 = 12345; 50 | static WORKER_NUM: u32 = 1; 51 | 52 | pub fn get_worker_num() -> u32 { 53 | WORKER_NUM 54 | } 55 | 56 | fn dump_json_testcases( 57 | output_writer: &OutputWriter, 58 | stats: &serde_json::Value, 59 | testcase_type: ExecutionStatus, 60 | ) { 61 | let test_case_vec = stats.get("testcases").unwrap().as_array().unwrap(); 62 | for test_case_val in test_case_vec { 63 | let test_case_str = test_case_val.as_str().unwrap(); 64 | let mut test_case: TestCase = serde_json::from_str(test_case_str).unwrap(); 65 | match testcase_type { 66 | ExecutionStatus::Crash => { 67 | test_case.gen_id(); 68 | match output_writer.save_crash(&test_case) { 69 | Ok(_v) => {} 70 | Err(_e) => println!("Cannot write crash file"), 71 | } 72 | } 73 | ExecutionStatus::Interesting => { 74 | output_writer.save_queue(&test_case).unwrap(); 75 | } 76 | ExecutionStatus::Timeout => { 77 | test_case.gen_id(); 78 | output_writer.save_hang(&test_case).unwrap(); 79 | } 80 | _ => {} 81 | } 82 | } 83 | } 84 | 85 | impl SimpleMonitor { 86 | fn default() -> Self { 87 | SimpleMonitor { 88 | services: Arc::new(Mutex::new(Vec::default())), 89 | service_addr_hashes: Arc::new(Mutex::new(HashSet::default())), 90 | addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), SIMPLE_MONITOR_PORT), 91 | fuzzer_info: FuzzerInfo::default(), 92 | } 93 | } 94 | 95 | #[allow(dead_code)] 96 | fn get_monitor_addr() -> SocketAddr { 97 | SIMPLE_MONITOR.read().unwrap().addr 98 | } 99 | 100 | fn get_services() -> Vec { 101 | SIMPLE_MONITOR.read().unwrap().get_services() 102 | } 103 | 104 | fn register_service(service: Service) -> bool { 105 | SIMPLE_MONITOR.write().unwrap().register_service(service) 106 | } 107 | 108 | pub fn get_fuzzer_info(&self) -> &FuzzerInfo { 109 | &self.fuzzer_info 110 | } 111 | 112 | pub fn get_fuzzer_info_mut(&mut self) -> &mut FuzzerInfo { 113 | &mut self.fuzzer_info 114 | } 115 | } 116 | 117 | pub fn get_monitor() -> Arc> { 118 | SIMPLE_MONITOR.clone() 119 | } 120 | 121 | impl Monitor for SimpleMonitor { 122 | fn reset_start_time(&mut self) { 123 | self.fuzzer_info.reset_start_time(); 124 | } 125 | fn get_services(&self) -> Vec { 126 | let services = self.services.lock().unwrap(); 127 | services.clone() 128 | } 129 | 130 | fn register_service(&self, service: Service) -> bool { 131 | let mut services = self.services.lock().unwrap(); 132 | let mut addr_hash = self.service_addr_hashes.lock().unwrap(); 133 | let addr = match service { 134 | Service::Mutator(addr) => addr, 135 | Service::FeedbackCollector(addr) => addr, 136 | Service::Executor(addr) => addr, 137 | Service::QueueManager(addr) => addr, 138 | }; 139 | if addr_hash.contains(&addr) { 140 | return false; 141 | } 142 | addr_hash.insert(addr); 143 | services.push(service); 144 | true 145 | } 146 | 147 | // TODO: Receive crash/hang/interesting test cases for saving in disk. 148 | fn receive_statistics(&self, stats: serde_json::Value) { 149 | let mut testcase_type = ExecutionStatus::Ok; 150 | if let Some(v) = stats.get("exec") { 151 | self.fuzzer_info.add_exec(v.as_u64().unwrap()); 152 | } else if let Some(v) = stats.get("crash") { 153 | self.fuzzer_info.add_crash(v.as_u64().unwrap()); 154 | testcase_type = ExecutionStatus::Crash; 155 | } else if let Some(v) = stats.get("timeout") { 156 | self.fuzzer_info.add_timeout_exec(v.as_u64().unwrap()); 157 | } else if let Some(v) = stats.get("interesting_test_case") { 158 | self.fuzzer_info.add_coverage(v.as_u64().unwrap()); 159 | } else { 160 | unreachable!(); 161 | } 162 | if self.fuzzer_info.get_output_writer().is_some() && stats.get("testcases").is_some() { 163 | dump_json_testcases( 164 | self.fuzzer_info.get_output_writer().unwrap(), 165 | &stats, 166 | testcase_type, 167 | ) 168 | } 169 | } 170 | 171 | fn show_statistics(&self) { 172 | self.fuzzer_info.simple_show(); 173 | self.fuzzer_info.stop_if_done(); 174 | } 175 | } 176 | 177 | #[cfg(test)] 178 | mod tests { 179 | use super::*; 180 | use std::str::FromStr; 181 | 182 | #[test] 183 | fn simple_monitor_return_its_addr() { 184 | let monitor_addr = SimpleMonitor::get_monitor_addr(); 185 | assert_eq!( 186 | SocketAddr::from_str("127.0.0.1:12345").unwrap(), 187 | monitor_addr 188 | ); 189 | } 190 | } 191 | -------------------------------------------------------------------------------- /src/frontend/simple_frontend.rs: -------------------------------------------------------------------------------- 1 | pub use async_channel::bounded as frontend_channel; 2 | use async_trait::async_trait; 3 | use std::sync::{Arc, Mutex}; 4 | pub type AsyncReceiver = async_channel::Receiver; 5 | pub type AsyncSender = async_channel::Sender; 6 | use crate::frontend::Worker; 7 | use serde_json::Value; 8 | 9 | // This macro generates the trait impl that are the same in all frontend. 10 | // TODO(yongheng): I don't think this is a good practice and it might make the code difficult to debug. 11 | #[macro_export] 12 | macro_rules! async_frontend_default { 13 | () => { 14 | fn get_async_output_receiver( 15 | &self, 16 | ) -> &AsyncReceiver::Output>> { 17 | self.output_receiver.as_ref().unwrap() 18 | } 19 | 20 | fn set_async_output_receiver( 21 | &mut self, 22 | receiver: AsyncReceiver::Output>>, 23 | ) { 24 | self.output_receiver = Some(receiver); 25 | } 26 | 27 | fn set_async_output_sender( 28 | &mut self, 29 | sender: AsyncSender::Output>>, 30 | ) { 31 | self.output_sender = Some(sender); 32 | } 33 | 34 | fn set_async_input_sender( 35 | &mut self, 36 | input: AsyncSender::Input>>, 37 | ) { 38 | self.sender = Some(input); 39 | } 40 | 41 | fn get_async_input_sender( 42 | &self, 43 | ) -> AsyncSender::Input>> { 44 | self.sender.as_ref().unwrap().clone() 45 | } 46 | }; 47 | } 48 | 49 | #[macro_export] 50 | macro_rules! frontend_default { 51 | () => { 52 | fn get_worker_pool(&self) -> &Vec>> { 53 | &self.worker_pool 54 | } 55 | 56 | fn get_worker_pool_mut(&mut self) -> &mut Vec>> { 57 | &mut self.worker_pool 58 | } 59 | }; 60 | } 61 | 62 | pub trait BasicFrontend { 63 | type Worker: crate::frontend::Worker + 'static; 64 | type Output; 65 | 66 | // We don't provide default implementation because some worker requires 67 | // extra arguments for creation. 68 | fn create_worker(&mut self, num: u32); 69 | 70 | fn delete_worker(&mut self, num: u32) { 71 | let worker = self.get_worker_pool_mut(); 72 | let final_len = worker.len().saturating_sub(num as usize); 73 | worker.truncate(final_len); 74 | } 75 | 76 | fn get_worker_by_idx(&self, idx: usize) -> Arc> { 77 | self.get_worker_pool()[idx].clone() 78 | } 79 | 80 | fn get_worker_pool_mut(&mut self) -> &mut Vec>>; 81 | fn get_worker_pool(&self) -> &Vec>>; 82 | 83 | // Transform the output of the worker to that of frontend. In most case it does nothing. 84 | fn output_transfrom( 85 | &self, 86 | input: Vec<::Output>, 87 | ) -> Vec; 88 | } 89 | 90 | pub trait AsyncFrontendChannel: BasicFrontend { 91 | fn set_async_input_sender( 92 | &mut self, 93 | input: AsyncSender::Input>>, 94 | ); 95 | 96 | fn get_async_input_sender( 97 | &self, 98 | ) -> AsyncSender::Input>>; 99 | 100 | fn set_async_output_receiver( 101 | &mut self, 102 | receiver: AsyncReceiver::Output>>, 103 | ); 104 | 105 | fn set_async_output_sender( 106 | &mut self, 107 | sender: AsyncSender::Output>>, 108 | ); 109 | 110 | fn get_async_output_receiver( 111 | &self, 112 | ) -> &AsyncReceiver::Output>>; 113 | } 114 | 115 | #[async_trait] 116 | pub trait AsyncFrontend: AsyncFrontendChannel { 117 | async fn run(&mut self) { 118 | let worker_pool = self.get_worker_pool().clone(); 119 | let (input_sx, input_rx) = frontend_channel(1000); 120 | let (output_sender, output_receiver) = frontend_channel(100); 121 | self.set_async_output_receiver(output_receiver); 122 | self.set_async_output_sender(output_sender.clone()); 123 | self.set_async_input_sender(input_sx); 124 | tokio::spawn(async move { 125 | let (worker_sx, worker_rx) = frontend_channel(worker_pool.len()); 126 | for worker in worker_pool.into_iter() { 127 | assert!(worker_sx.send(worker).await.is_ok()); 128 | } 129 | 130 | while let Ok(inputs) = input_rx.recv().await { 131 | if let Ok(worker) = worker_rx.recv().await { 132 | let save_worker_sx = worker_sx.clone(); 133 | let output_sender_tmp = output_sender.clone(); 134 | tokio::spawn(async move { 135 | let send_back = worker.clone(); 136 | let worker_output = { 137 | let mut worker = worker.lock().unwrap(); 138 | worker.handle_one_input(inputs) 139 | }; 140 | if !worker_output.is_empty() { 141 | assert!(output_sender_tmp.send(worker_output).await.is_ok()); 142 | } 143 | assert!(save_worker_sx.send(send_back).await.is_ok()); 144 | }); 145 | } 146 | } 147 | }); 148 | } 149 | 150 | // Handle the inputs and save the result into the pool. 151 | // The inputs are splitted evenly and distributed to every workers. 152 | async fn handle_inputs(&self, inputs: Vec<::Input>) { 153 | if inputs.is_empty() { 154 | return; 155 | } 156 | 157 | let mut inputs = inputs; 158 | let worker_len = self.get_worker_pool().len(); 159 | assert_ne!(worker_len, 0); 160 | let mut split_workload = inputs.len() / worker_len; 161 | if inputs.len() % worker_len != 0 { 162 | split_workload += 1; 163 | } 164 | while !inputs.is_empty() { 165 | let sub_input = inputs.split_off(inputs.len().saturating_sub(split_workload)); 166 | assert!(self.get_async_input_sender().send(sub_input).await.is_ok()); 167 | } 168 | } 169 | 170 | // Fetch some results from the pool with best effort. It doesn't promise to 171 | // fetch as many result as you want. To avoid returning nothing, we try to wait for 172 | // at least one result to be produced. 173 | async fn get_results(&self, num: Option) -> Vec { 174 | let num = num.unwrap_or(2000) as usize; 175 | let receiver = self.get_async_output_receiver(); 176 | //let num = num.unwrap_or(receiver.len() as u32) as usize; 177 | 178 | let mut result = Vec::with_capacity(num); 179 | while result.len() < num { 180 | // Result avaible? 181 | if let Ok(some_result) = receiver.try_recv() { 182 | result.extend(some_result.into_iter()); 183 | } 184 | // If not and we have nothing now, we wait for some. 185 | else if result.is_empty() { 186 | if let Ok(some_result) = receiver.recv().await { 187 | result.extend(some_result.into_iter()); 188 | } 189 | } else { 190 | break; 191 | } 192 | } 193 | self.output_transfrom(result) 194 | } 195 | 196 | async fn retrieve_monitor_data(&self) -> Vec { 197 | unreachable!(); 198 | } 199 | } 200 | -------------------------------------------------------------------------------- /src/datatype.rs: -------------------------------------------------------------------------------- 1 | use crate::mutator::afl_mutator::DeterministicMutationPass; 2 | use serde::{Deserialize, Serialize}; 3 | use std::sync::atomic::{AtomicU32, Ordering}; 4 | 5 | pub type TestCaseID = u32; 6 | 7 | #[derive(Default, Clone, Debug, Serialize, Deserialize)] 8 | pub struct TestCaseMetaInfo { 9 | #[allow(dead_code)] 10 | done_pass: DeterministicMutationPass, // All the finished determinsitic passes. 11 | todo_pass: DeterministicMutationPass, // The determinsitic passes that we want to do next. 12 | } 13 | 14 | impl TestCaseMetaInfo { 15 | pub fn get_todo_pass(&self) -> DeterministicMutationPass { 16 | self.todo_pass 17 | } 18 | 19 | pub fn set_todo_pass(&mut self, todo_pass: DeterministicMutationPass) { 20 | self.todo_pass = todo_pass; 21 | } 22 | 23 | pub fn get_done_pass(&self) -> DeterministicMutationPass { 24 | self.done_pass 25 | } 26 | 27 | pub fn set_done_pass(&mut self, done_pass: DeterministicMutationPass) { 28 | self.done_pass = done_pass; 29 | } 30 | } 31 | 32 | #[derive(Default, Clone, Debug, Serialize, Deserialize)] 33 | pub struct TestCase { 34 | id: TestCaseID, // The id of the test_case itself or the id of parent. Since we won't use both of them at the same time, we reuse the same slot. 35 | mutator_id: u32, 36 | buffer: Vec, 37 | meta: Option, 38 | } 39 | 40 | impl TestCase { 41 | pub fn new(buffer: Vec, id: u32) -> Self { 42 | TestCase { 43 | id, 44 | buffer, 45 | mutator_id: 0, 46 | meta: None, 47 | } 48 | } 49 | 50 | pub fn clone_without_meta(&self) -> Self { 51 | let mut result = self.clone(); 52 | result.take_meta(); 53 | result 54 | } 55 | 56 | pub fn set_meta(&mut self, meta: TestCaseMetaInfo) { 57 | self.meta = Some(meta); 58 | } 59 | 60 | pub fn borrow_meta(&self) -> Option<&TestCaseMetaInfo> { 61 | self.meta.as_ref() 62 | } 63 | 64 | pub fn take_meta(&mut self) -> Option { 65 | self.meta.take() 66 | } 67 | 68 | pub fn borrow_meta_mut(&mut self) -> Option<&mut TestCaseMetaInfo> { 69 | self.meta.as_mut() 70 | } 71 | 72 | pub fn has_meta(&self) -> bool { 73 | self.meta.is_some() 74 | } 75 | 76 | pub fn len(&self) -> usize { 77 | self.buffer.len() 78 | } 79 | 80 | pub fn is_empty(&self) -> bool { 81 | self.buffer.is_empty() 82 | } 83 | 84 | pub fn extract_mutation_info(&self) -> MutationInfo { 85 | MutationInfo::new(self.id, self.mutator_id) 86 | } 87 | 88 | pub fn set_mutator_id(&mut self, mid: u32) { 89 | self.mutator_id = mid; 90 | } 91 | 92 | pub fn get_mutator_id(&self) -> u32 { 93 | self.mutator_id 94 | } 95 | 96 | pub fn get_buffer(&self) -> &Vec { 97 | &self.buffer 98 | } 99 | 100 | pub fn get_buffer_mut(&mut self) -> &mut Vec { 101 | &mut self.buffer 102 | } 103 | 104 | pub fn take_buffer(self) -> Vec { 105 | self.buffer 106 | } 107 | 108 | pub fn get_id(&self) -> u32 { 109 | self.id 110 | } 111 | 112 | pub fn set_id(&mut self, id: TestCaseID) { 113 | self.id = id; 114 | } 115 | 116 | pub fn get_pid(&self) -> u32 { 117 | self.id 118 | } 119 | 120 | pub fn create_variant(&self, buffer: Vec) -> TestCase { 121 | Self::new(buffer, self.id) 122 | } 123 | 124 | pub fn gen_id(&mut self) { 125 | self.id = get_id(); 126 | } 127 | } 128 | 129 | static COUNTER: AtomicU32 = AtomicU32::new(0); 130 | pub fn get_id() -> u32 { 131 | COUNTER.fetch_add(1, Ordering::Relaxed) 132 | } 133 | 134 | #[cfg(test)] 135 | pub fn reset_id() { 136 | COUNTER.store(0, Ordering::Relaxed) 137 | } 138 | 139 | #[derive(Clone, Copy, Debug, PartialEq, Eq)] 140 | pub enum ExecutionStatus { 141 | Ok, // Nothing special happens 142 | Interesting, // Find new coverage or something that worth notice 143 | Crash, 144 | Timeout, 145 | } 146 | 147 | #[derive(Default, Clone, Hash, PartialEq, Eq, Copy, Debug)] 148 | pub struct MutationInfo { 149 | pid: u32, 150 | mutation_id: u32, 151 | } 152 | 153 | impl MutationInfo { 154 | pub fn new(pid: u32, mutation_id: u32) -> Self { 155 | Self { pid, mutation_id } 156 | } 157 | 158 | pub fn get_pid(&self) -> u32 { 159 | self.pid 160 | } 161 | 162 | pub fn get_mutation_id(&self) -> u32 { 163 | self.mutation_id 164 | } 165 | 166 | pub fn get_mutator_id2(&self) -> u32 { 167 | self.mutation_id >> 16 168 | } 169 | 170 | pub fn get_mutator_func_id(&self) -> u32 { 171 | self.mutation_id & 0xFFFF 172 | } 173 | } 174 | 175 | #[derive(Default, Clone, Hash, PartialEq, Eq, Copy, Debug)] 176 | pub struct NewBit { 177 | pub index: usize, 178 | pub val: u8, 179 | } 180 | 181 | impl NewBit { 182 | pub fn new(index: usize, val: u8) -> Self { 183 | Self { index, val } 184 | } 185 | } 186 | 187 | #[derive(Clone, Debug)] 188 | pub enum FeedbackData { 189 | Unknown, 190 | NewCoverage(Vec), 191 | Counter(u32), // Indicate how many times such feedback is repeated. 192 | } 193 | 194 | impl FeedbackData { 195 | pub fn new_coverage(new_bits: Vec) -> Self { 196 | Self::NewCoverage(new_bits) 197 | } 198 | 199 | pub fn into_new_coverage(self) -> Option> { 200 | match self { 201 | Self::NewCoverage(coverage) => Some(coverage), 202 | _ => None, 203 | } 204 | } 205 | } 206 | 207 | #[derive(Clone, Debug)] 208 | pub struct Feedback { 209 | status: ExecutionStatus, 210 | test_case: Option, // The test case that produce this feedback 211 | data: Option, // The real feedback data 212 | mutation_info: Option, 213 | } 214 | 215 | impl Feedback { 216 | pub fn new(status: ExecutionStatus) -> Self { 217 | Self { 218 | status, 219 | test_case: None, 220 | data: None, 221 | mutation_info: None, 222 | } 223 | } 224 | 225 | pub fn get_status(&self) -> ExecutionStatus { 226 | self.status 227 | } 228 | 229 | pub fn get_mutation_info(&self) -> Option<&MutationInfo> { 230 | self.mutation_info.as_ref() 231 | } 232 | 233 | pub fn set_mutation_info(mut self, mutation_info: MutationInfo) -> Self { 234 | self.mutation_info = Some(mutation_info); 235 | self 236 | } 237 | 238 | pub fn is_valid(&self) -> bool { 239 | matches!(self.status, ExecutionStatus::Ok) && self.test_case.is_none() 240 | || (matches!(self.status, ExecutionStatus::Interesting) 241 | && (self.test_case.is_some() || self.data.is_some())) 242 | || (matches!(self.status, ExecutionStatus::Crash) && self.test_case.is_some()) 243 | || (matches!(self.status, ExecutionStatus::Timeout) && self.test_case.is_none()) 244 | } 245 | 246 | /* 247 | pub fn get_test_case(&self) -> Option<&TestCase> { 248 | self.test_case.as_ref() 249 | } 250 | */ 251 | pub fn take_mutation_info(&mut self) -> Option { 252 | self.mutation_info.take() 253 | } 254 | 255 | pub fn take_test_case(&mut self) -> Option { 256 | self.test_case.take() 257 | } 258 | 259 | pub fn borrow_test_case(&self) -> Option<&TestCase> { 260 | self.test_case.as_ref() 261 | } 262 | 263 | pub fn contain_test_case(&self) -> bool { 264 | self.test_case.is_some() 265 | } 266 | 267 | pub fn take_data(&mut self) -> Option { 268 | self.data.take() 269 | } 270 | 271 | pub fn borrow_data(&self) -> Option<&FeedbackData> { 272 | self.data.as_ref() 273 | } 274 | 275 | pub fn set_data(mut self, data: FeedbackData) -> Self { 276 | self.data = Some(data); 277 | self 278 | } 279 | 280 | pub fn set_test_case(mut self, test_case: TestCase) -> Self { 281 | if self.mutation_info.is_none() { 282 | self.mutation_info = Some(test_case.extract_mutation_info()); 283 | } 284 | self.test_case = Some(test_case); 285 | self 286 | } 287 | } 288 | 289 | #[cfg(test)] 290 | mod tests { 291 | use super::*; 292 | #[test] 293 | fn test_construct_test_case() { 294 | let test_case = TestCase::new(Vec::default(), 0); 295 | let test_case2 = test_case.create_variant(Vec::default()); 296 | assert_eq!(test_case.get_id(), test_case2.get_pid()); 297 | } 298 | } 299 | -------------------------------------------------------------------------------- /src/executor/shmem.rs: -------------------------------------------------------------------------------- 1 | extern crate alloc; 2 | use crate::Error; 3 | use serde::{Deserialize, Serialize}; 4 | use std::env; 5 | 6 | use core::fmt::{self, Debug, Display}; 7 | use core::{ptr, slice}; 8 | 9 | use libc::{c_int, c_uchar, shmat, shmctl, shmget}; 10 | use std::ptr::null_mut; 11 | 12 | /// Description of a shared map. 13 | /// May be used to restore the map by id. 14 | #[derive(Copy, Clone, Debug, Serialize, Deserialize)] 15 | pub struct ShMemDescription { 16 | /// Size of this map 17 | pub size: usize, 18 | /// Id of this map 19 | pub id: ShMemId, 20 | } 21 | 22 | impl ShMemDescription { 23 | /// Create a description from a `id_str` and a `size`. 24 | #[must_use] 25 | pub fn from_string_and_size(id_str: &str, size: usize) -> Self { 26 | Self { 27 | size, 28 | id: ShMemId::from_string(id_str), 29 | } 30 | } 31 | } 32 | 33 | pub trait ShMem: Sized + Debug + Clone { 34 | /// Get the id of this shared memory mapping 35 | fn id(&self) -> ShMemId; 36 | 37 | /// Get the size of this mapping 38 | fn len(&self) -> usize; 39 | 40 | /// Check if the mapping is empty 41 | fn is_empty(&self) -> bool { 42 | self.len() == 0 43 | } 44 | 45 | /// Get the description of the shared memory mapping 46 | fn description(&self) -> ShMemDescription { 47 | ShMemDescription { 48 | size: self.len(), 49 | id: self.id(), 50 | } 51 | } 52 | 53 | /// The actual shared map, in memory 54 | fn map(&self) -> &[u8]; 55 | 56 | /// The actual shared map, mutable 57 | fn map_mut(&mut self) -> &mut [u8]; 58 | 59 | /// Write this map's config to env 60 | #[cfg(feature = "std")] 61 | fn write_to_env(&self, env_name: &str) -> Result<(), Error> { 62 | let map_size = self.len(); 63 | let map_size_env = format!("{}_SIZE", env_name); 64 | env::set_var(env_name, self.id().to_string()); 65 | env::set_var(map_size_env, format!("{}", map_size)); 66 | Ok(()) 67 | } 68 | } 69 | 70 | /// An id associated with a given shared memory mapping ([`ShMem`]), which can be used to 71 | /// establish shared-mappings between proccesses. 72 | #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, Default)] 73 | pub struct ShMemId { 74 | id: [u8; 20], 75 | } 76 | 77 | impl ShMemId { 78 | /// Create a new id from a fixed-size string 79 | #[must_use] 80 | pub fn from_slice(slice: &[u8; 20]) -> Self { 81 | Self { id: *slice } 82 | } 83 | 84 | /// Create a new id from an int 85 | #[must_use] 86 | pub fn from_int(val: i32) -> Self { 87 | Self::from_string(&val.to_string()) 88 | } 89 | 90 | /// Create a new id from a string 91 | #[must_use] 92 | pub fn from_string(val: &str) -> Self { 93 | let mut slice: [u8; 20] = [0; 20]; 94 | for (i, val) in val.as_bytes().iter().enumerate() { 95 | slice[i] = *val; 96 | } 97 | Self { id: slice } 98 | } 99 | 100 | /// Get the id as a fixed-length slice 101 | #[must_use] 102 | pub fn as_slice(&self) -> &[u8; 20] { 103 | &self.id 104 | } 105 | 106 | /// Returns the first null-byte in or the end of the buffer 107 | #[must_use] 108 | pub fn null_pos(&self) -> usize { 109 | self.id.iter().position(|&c| c == 0).unwrap() 110 | } 111 | 112 | /// Returns a `str` representation of this [`ShMemId`] 113 | #[must_use] 114 | pub fn as_str(&self) -> &str { 115 | alloc::str::from_utf8(&self.id[..self.null_pos()]).unwrap() 116 | } 117 | } 118 | 119 | impl From for i32 { 120 | fn from(id: ShMemId) -> i32 { 121 | id.as_str().parse().unwrap() 122 | } 123 | } 124 | 125 | impl Display for ShMemId { 126 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 127 | write!(f, "{}", self.as_str()) 128 | } 129 | } 130 | 131 | pub trait ShMemProvider: Clone + Default + Debug { 132 | /// The actual shared map handed out by this [`ShMemProvider`]. 133 | type Mem: ShMem; 134 | 135 | /// Create a new instance of the provider 136 | fn new() -> Result; 137 | 138 | /// Create a new shared memory mapping 139 | fn new_map(&mut self, map_size: usize) -> Result; 140 | 141 | /// Get a mapping given its id and size 142 | fn get_by_id_and_size(&mut self, id: ShMemId, size: usize) -> Result; 143 | 144 | /// Get a mapping given a description 145 | fn get_by_description(&mut self, description: ShMemDescription) -> Result { 146 | self.get_by_id_and_size(description.id, description.size) 147 | } 148 | 149 | /// Create a new sharedmap reference from an existing `id` and `len` 150 | fn clone_ref(&mut self, mapping: &Self::Mem) -> Result { 151 | self.get_by_id_and_size(mapping.id(), mapping.len()) 152 | } 153 | 154 | /// Reads an existing map config from env vars, then maps it 155 | #[cfg(feature = "std")] 156 | fn existing_from_env(&mut self, env_name: &str) -> Result { 157 | let map_shm_str = env::var(env_name)?; 158 | let map_size = str::parse::(&env::var(format!("{}_SIZE", env_name))?)?; 159 | self.get_by_description(ShMemDescription::from_string_and_size( 160 | &map_shm_str, 161 | map_size, 162 | )) 163 | } 164 | 165 | /// This method should be called before a fork or a thread creation event, allowing the [`ShMemProvider`] to 166 | /// get ready for a potential reset of thread specific info, and for potential reconnects. 167 | /// Make sure to call [`Self::post_fork()`] after threading! 168 | fn pre_fork(&mut self) -> Result<(), Error> { 169 | // do nothing 170 | Ok(()) 171 | } 172 | 173 | /// This method should be called after a fork or after cloning/a thread creation event, allowing the [`ShMemProvider`] to 174 | /// reset thread specific info, and potentially reconnect. 175 | /// Make sure to call [`Self::pre_fork()`] before threading! 176 | fn post_fork(&mut self, _is_child: bool) -> Result<(), Error> { 177 | // do nothing 178 | Ok(()) 179 | } 180 | 181 | /// Release the resources associated with the given [`ShMem`] 182 | fn release_map(&mut self, _map: &mut Self::Mem) { 183 | // do nothing 184 | } 185 | } 186 | 187 | #[derive(Clone, Debug)] 188 | pub struct CommonUnixShMem { 189 | id: ShMemId, 190 | map: *mut u8, 191 | map_size: usize, 192 | } 193 | 194 | unsafe impl Send for CommonUnixShMem {} 195 | 196 | impl CommonUnixShMem { 197 | /// Create a new shared memory mapping, using shmget/shmat 198 | pub fn new(map_size: usize) -> Result { 199 | unsafe { 200 | let os_id = shmget( 201 | libc::IPC_PRIVATE, 202 | map_size, 203 | libc::IPC_CREAT | libc::IPC_EXCL | libc::SHM_R | libc::SHM_W, 204 | ); 205 | 206 | if os_id < 0_i32 { 207 | return Err(Error::Unknown(format!("Failed to allocate a shared mapping of size {} - check OS limits (i.e shmall, shmmax)", map_size))); 208 | } 209 | 210 | let map = shmat(os_id, ptr::null(), 0) as *mut c_uchar; 211 | 212 | if map as c_int == -1 || map.is_null() { 213 | shmctl(os_id, libc::IPC_RMID, ptr::null_mut()); 214 | return Err(Error::Unknown( 215 | "Failed to map the shared mapping".to_string(), 216 | )); 217 | } 218 | 219 | Ok(Self { 220 | id: ShMemId::from_int(os_id), 221 | map, 222 | map_size, 223 | }) 224 | } 225 | } 226 | 227 | /// Get a [`UnixShMem`] of the existing shared memory mapping identified by id 228 | pub fn from_id_and_size(id: ShMemId, map_size: usize) -> Result { 229 | unsafe { 230 | let id_int: i32 = id.into(); 231 | let map = shmat(id_int, ptr::null(), 0) as *mut c_uchar; 232 | 233 | if map.is_null() || map == null_mut::().wrapping_sub(1) { 234 | return Err(Error::Unknown( 235 | "Failed to map the shared mapping".to_string(), 236 | )); 237 | } 238 | 239 | Ok(Self { id, map, map_size }) 240 | } 241 | } 242 | } 243 | 244 | #[cfg(unix)] 245 | impl ShMem for CommonUnixShMem { 246 | fn id(&self) -> ShMemId { 247 | self.id 248 | } 249 | 250 | fn len(&self) -> usize { 251 | self.map_size 252 | } 253 | 254 | fn map(&self) -> &[u8] { 255 | unsafe { slice::from_raw_parts(self.map, self.map_size) } 256 | } 257 | 258 | fn map_mut(&mut self) -> &mut [u8] { 259 | unsafe { slice::from_raw_parts_mut(self.map, self.map_size) } 260 | } 261 | } 262 | 263 | /// [`Drop`] implementation for [`UnixShMem`], which cleans up the mapping. 264 | #[cfg(unix)] 265 | impl Drop for CommonUnixShMem { 266 | fn drop(&mut self) { 267 | unsafe { 268 | let id_int: i32 = self.id.into(); 269 | shmctl(id_int, libc::IPC_RMID, ptr::null_mut()); 270 | } 271 | } 272 | } 273 | 274 | /// A [`ShMemProvider`] which uses `shmget`/`shmat`/`shmctl` to provide shared memory mappings. 275 | #[cfg(unix)] 276 | #[derive(Clone, Debug)] 277 | pub struct CommonUnixShMemProvider {} 278 | 279 | unsafe impl Send for CommonUnixShMemProvider {} 280 | 281 | #[cfg(unix)] 282 | impl Default for CommonUnixShMemProvider { 283 | fn default() -> Self { 284 | Self::new().unwrap() 285 | } 286 | } 287 | 288 | /// Implement [`ShMemProvider`] for [`UnixShMemProvider`]. 289 | #[cfg(unix)] 290 | impl ShMemProvider for CommonUnixShMemProvider { 291 | type Mem = CommonUnixShMem; 292 | 293 | fn new() -> Result { 294 | Ok(Self {}) 295 | } 296 | fn new_map(&mut self, map_size: usize) -> Result { 297 | CommonUnixShMem::new(map_size) 298 | } 299 | 300 | fn get_by_id_and_size(&mut self, id: ShMemId, size: usize) -> Result { 301 | CommonUnixShMem::from_id_and_size(id, size) 302 | } 303 | } 304 | -------------------------------------------------------------------------------- /src/rpc.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype; 2 | use crate::monitor::Service; 3 | use msg_type::ServiceType; 4 | use serde_json::Value; 5 | use std::net::SocketAddr; 6 | use std::str::FromStr; 7 | 8 | pub mod msg_type { 9 | #![allow(clippy::derive_partial_eq_without_eq)] 10 | tonic::include_proto!("microfuzz.msg_type"); 11 | } 12 | 13 | pub mod execution { 14 | #![allow(clippy::derive_partial_eq_without_eq)] 15 | tonic::include_proto!("microfuzz.execution"); 16 | } 17 | 18 | pub mod feedback { 19 | #![allow(clippy::derive_partial_eq_without_eq)] 20 | tonic::include_proto!("microfuzz.feedback"); 21 | } 22 | 23 | pub mod mutation { 24 | #![allow(clippy::derive_partial_eq_without_eq)] 25 | tonic::include_proto!("microfuzz.mutation"); 26 | } 27 | 28 | pub mod queue_management { 29 | #![allow(clippy::derive_partial_eq_without_eq)] 30 | tonic::include_proto!("microfuzz.queue_management"); 31 | } 32 | 33 | pub mod monitor { 34 | #![allow(clippy::derive_partial_eq_without_eq)] 35 | tonic::include_proto!("microfuzz.monitor"); 36 | } 37 | 38 | impl msg_type::MonitorData { 39 | pub fn from(data: Vec) -> Self { 40 | Self { 41 | data: data.iter().map(serde_json::Value::to_string).collect(), 42 | } 43 | } 44 | 45 | pub fn into(self) -> Vec { 46 | self.data 47 | .into_iter() 48 | .filter_map(|s| { 49 | if let Ok(s) = serde_json::from_str(s.as_str()) { 50 | Some(s) 51 | } else { 52 | None 53 | } 54 | }) 55 | .collect::>() 56 | } 57 | } 58 | 59 | impl msg_type::ServiceInfo { 60 | pub fn from(service: Service) -> Self { 61 | match service { 62 | Service::FeedbackCollector(addr) => Self { 63 | service_type: ServiceType::Feedbackcollector as i32, 64 | socket_addr: addr.to_string(), 65 | }, 66 | Service::QueueManager(addr) => Self { 67 | service_type: ServiceType::Queuemanager as i32, 68 | socket_addr: addr.to_string(), 69 | }, 70 | Service::Mutator(addr) => Self { 71 | service_type: ServiceType::Mutator as i32, 72 | socket_addr: addr.to_string(), 73 | }, 74 | Service::Executor(addr) => Self { 75 | service_type: ServiceType::Executor as i32, 76 | socket_addr: addr.to_string(), 77 | }, 78 | } 79 | } 80 | 81 | pub fn into(self) -> Option { 82 | let addr = SocketAddr::from_str(&self.socket_addr).unwrap(); 83 | let service = match ServiceType::from_i32(self.service_type) { 84 | Some(ServiceType::Executor) => Service::Executor(addr), 85 | Some(ServiceType::Mutator) => Service::Mutator(addr), 86 | Some(ServiceType::Feedbackcollector) => Service::FeedbackCollector(addr), 87 | Some(ServiceType::Queuemanager) => Service::QueueManager(addr), 88 | None => return None, 89 | }; 90 | Some(service) 91 | } 92 | } 93 | 94 | impl msg_type::TestCase { 95 | pub fn from(test_case: datatype::TestCase) -> Self { 96 | Self { 97 | pid: test_case.get_pid(), 98 | content: test_case.take_buffer(), 99 | } 100 | } 101 | 102 | pub fn into(self) -> datatype::TestCase { 103 | crate::datatype::TestCase::new(self.content, self.pid) 104 | } 105 | } 106 | 107 | impl msg_type::MutationInfo { 108 | pub fn from(mutation_info: datatype::MutationInfo) -> Self { 109 | Self { 110 | pid: mutation_info.get_pid(), 111 | mutator_id: mutation_info.get_mutation_id(), 112 | } 113 | } 114 | 115 | pub fn into(self) -> datatype::MutationInfo { 116 | crate::datatype::MutationInfo::new(self.pid, self.mutator_id) 117 | } 118 | } 119 | 120 | impl msg_type::ExecutionStatus { 121 | pub fn from(status: datatype::ExecutionStatus) -> Self { 122 | match status { 123 | datatype::ExecutionStatus::Ok => msg_type::ExecutionStatus::Ok, 124 | datatype::ExecutionStatus::Timeout => msg_type::ExecutionStatus::Timeout, 125 | datatype::ExecutionStatus::Crash => msg_type::ExecutionStatus::Crash, 126 | datatype::ExecutionStatus::Interesting => msg_type::ExecutionStatus::Interesting, 127 | } 128 | } 129 | 130 | pub fn into(self) -> datatype::ExecutionStatus { 131 | match self { 132 | msg_type::ExecutionStatus::Ok => datatype::ExecutionStatus::Ok, 133 | msg_type::ExecutionStatus::Timeout => datatype::ExecutionStatus::Timeout, 134 | msg_type::ExecutionStatus::Crash => datatype::ExecutionStatus::Crash, 135 | msg_type::ExecutionStatus::Interesting => datatype::ExecutionStatus::Interesting, 136 | } 137 | } 138 | } 139 | 140 | impl msg_type::NewBit { 141 | pub fn from(new_bit: datatype::NewBit) -> Self { 142 | Self { 143 | index: new_bit.index as u32, 144 | value: new_bit.val as u32, 145 | } 146 | } 147 | 148 | pub fn into(self) -> datatype::NewBit { 149 | datatype::NewBit::new(self.index as usize, self.value as u8) 150 | } 151 | } 152 | 153 | impl msg_type::feedback_data::FeedbackData { 154 | pub fn from(data: datatype::FeedbackData) -> Self { 155 | match data { 156 | datatype::FeedbackData::NewCoverage(coverage) => { 157 | msg_type::feedback_data::FeedbackData::NewCoverage(msg_type::NewCoverage { 158 | new_bits: coverage.into_iter().map(msg_type::NewBit::from).collect(), 159 | }) 160 | } 161 | datatype::FeedbackData::Counter(counter) => { 162 | msg_type::feedback_data::FeedbackData::Counter(msg_type::Counter { counter }) 163 | } 164 | datatype::FeedbackData::Unknown => msg_type::feedback_data::FeedbackData::FakeFeedback( 165 | msg_type::FakeFeedback::default(), 166 | ), 167 | } 168 | } 169 | 170 | pub fn into(self) -> datatype::FeedbackData { 171 | match self { 172 | msg_type::feedback_data::FeedbackData::NewCoverage(coverage) => { 173 | datatype::FeedbackData::NewCoverage( 174 | coverage 175 | .new_bits 176 | .into_iter() 177 | .map(|new_bit| new_bit.into()) 178 | .collect(), 179 | ) 180 | } 181 | msg_type::feedback_data::FeedbackData::Counter(counter) => { 182 | datatype::FeedbackData::Counter(counter.counter) 183 | } 184 | msg_type::feedback_data::FeedbackData::FakeFeedback(_) => { 185 | datatype::FeedbackData::Unknown 186 | } 187 | } 188 | } 189 | } 190 | 191 | impl msg_type::Feedback { 192 | pub fn from(mut feedback: crate::datatype::Feedback) -> Self { 193 | let test_case = feedback.take_test_case().map(msg_type::TestCase::from); 194 | let mutation_info = feedback 195 | .take_mutation_info() 196 | .map(msg_type::MutationInfo::from); 197 | let status = msg_type::ExecutionStatus::from(feedback.get_status()) as i32; 198 | 199 | let data = feedback 200 | .take_data() 201 | .map(msg_type::feedback_data::FeedbackData::from); 202 | Self { 203 | test_case, 204 | status, 205 | mutation_info, 206 | data: Some(msg_type::FeedbackData { 207 | feedback_data: data, 208 | }), 209 | } 210 | } 211 | 212 | pub fn into(self) -> datatype::Feedback { 213 | let status = msg_type::ExecutionStatus::from_i32(self.status) 214 | .unwrap() 215 | .into(); 216 | let mut result = datatype::Feedback::new(status); 217 | 218 | if let Some(mutation_info) = self.mutation_info { 219 | result = result.set_mutation_info(mutation_info.into()); 220 | } 221 | 222 | if let Some(test_case) = self.test_case.map(msg_type::TestCase::into) { 223 | result = result.set_test_case(test_case); 224 | } 225 | 226 | if let Some(data) = self.data { 227 | if let Some(feedback_data) = data.feedback_data { 228 | result = result.set_data(feedback_data.into()); 229 | } 230 | } 231 | 232 | result 233 | } 234 | } 235 | 236 | pub async fn run_service_in_background(svc: S, addr: Option) -> SocketAddr 237 | where 238 | S: tower::Service< 239 | http::Request, 240 | Response = http::Response, 241 | Error = std::convert::Infallible, 242 | > + tonic::transport::NamedService 243 | + Clone 244 | + Send 245 | + 'static, 246 | S::Future: Send + 'static, 247 | S::Error: Into> + Send, 248 | { 249 | let mut server_addr_str = String::from("127.0.0.1:0"); 250 | if let Some(server_addr) = addr { 251 | server_addr_str = server_addr; 252 | } 253 | let listener = tokio::net::TcpListener::bind(server_addr_str) 254 | .await 255 | .unwrap(); 256 | let addr = listener.local_addr().unwrap(); 257 | tokio::spawn(async move { 258 | tonic::transport::Server::builder() 259 | .add_service(svc) 260 | .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener)) 261 | .await 262 | .unwrap(); 263 | }); 264 | addr 265 | } 266 | 267 | #[cfg(test)] 268 | mod tests { 269 | use super::*; 270 | #[test] 271 | fn test_case_conversion_keep_all_info() { 272 | let test_case_inner = crate::datatype::TestCase::new(vec![1, 2, 3], 123); 273 | let test_case_rpc = msg_type::TestCase::from(test_case_inner.clone()); 274 | assert_eq!(test_case_inner.get_pid(), test_case_rpc.pid); 275 | assert_eq!(test_case_inner.get_buffer().clone(), test_case_rpc.content); 276 | 277 | let test_case_rpc_inner = test_case_rpc.into(); 278 | assert_eq!(test_case_inner.get_pid(), test_case_rpc_inner.get_pid()); 279 | assert_eq!( 280 | test_case_inner.get_buffer(), 281 | test_case_rpc_inner.get_buffer() 282 | ); 283 | } 284 | 285 | #[test] 286 | fn feedback_conversion_keep_all_info() { 287 | // TODO 288 | } 289 | } 290 | -------------------------------------------------------------------------------- /src/monitor/stats.rs: -------------------------------------------------------------------------------- 1 | use crate::util::pretty_format_num; 2 | use lazy_static::lazy_static; 3 | use psutil::cpu::CpuPercentCollector; 4 | use serde_json::Value; 5 | use std::fs; 6 | use std::net::SocketAddr; 7 | use std::sync::atomic::{AtomicU64, Ordering}; 8 | use std::sync::{Arc, Mutex}; 9 | use std::time::SystemTime; 10 | 11 | use super::output_writer::OutputWriter; 12 | 13 | // All the high level information we want to know about the fuzzer. 14 | // TODO: This should be merged with the monitor. 15 | #[allow(dead_code)] 16 | pub struct FuzzerInfo { 17 | // Dynamic info 18 | start_time: SystemTime, 19 | last_print_time: AtomicU64, 20 | last_num_exec: AtomicU64, 21 | num_exec: AtomicU64, 22 | cur_exec_speed: AtomicU64, 23 | num_timeout_exec: AtomicU64, 24 | num_crashes: AtomicU64, 25 | cycle_done: AtomicU64, 26 | cur_coverage: AtomicU64, 27 | 28 | // CPU usage util past print 29 | cpu_usage: AtomicU64, 30 | all_cpu_usage: AtomicU64, 31 | all_cpu_record_times: AtomicU64, 32 | // Configuration 33 | total_coverage: u64, 34 | max_exec: u64, 35 | max_fuzzing_time: u64, 36 | cmd: String, 37 | num_executor: u64, 38 | num_mutator: u64, 39 | num_feedback_collector: u64, 40 | num_queue_manager: u64, 41 | num_core: u64, 42 | monitor_addr: SocketAddr, 43 | output_writer: Option, 44 | } 45 | 46 | lazy_static! { 47 | static ref FUZZER_INFO: Arc> = Arc::new(Mutex::new(FuzzerInfo::new())); 48 | static ref CPU_PERCENT_COLLECTOR: Arc> = 49 | Arc::new(Mutex::new(CpuPercentCollector::new().unwrap())); 50 | } 51 | 52 | pub fn get_fuzzer_info() -> Arc> { 53 | FUZZER_INFO.clone() 54 | } 55 | 56 | pub fn get_cpu_percentage_collector() -> Arc> { 57 | CPU_PERCENT_COLLECTOR.clone() 58 | } 59 | 60 | impl FuzzerInfo { 61 | pub fn new() -> Self { 62 | FuzzerInfo { 63 | start_time: SystemTime::now(), 64 | last_print_time: AtomicU64::new(0), 65 | last_num_exec: AtomicU64::new(0), 66 | num_exec: AtomicU64::new(0), 67 | cur_exec_speed: AtomicU64::new(0), 68 | num_timeout_exec: AtomicU64::new(0), 69 | num_crashes: AtomicU64::new(0), 70 | cycle_done: AtomicU64::new(0), 71 | total_coverage: 0, 72 | cur_coverage: AtomicU64::new(0), 73 | cmd: "".to_string(), 74 | num_executor: 0, 75 | num_mutator: 0, 76 | num_feedback_collector: 0, 77 | num_queue_manager: 0, 78 | num_core: 1, 79 | max_exec: u64::MAX, 80 | max_fuzzing_time: u64::MAX, 81 | monitor_addr: "127.0.0.1:12345".parse().unwrap(), 82 | output_writer: None, 83 | cpu_usage: AtomicU64::new(0), 84 | all_cpu_usage: AtomicU64::new(0), 85 | all_cpu_record_times: AtomicU64::new(0), 86 | } 87 | } 88 | 89 | pub fn reset_start_time(&mut self) { 90 | self.start_time = SystemTime::now(); 91 | } 92 | 93 | pub fn get_output_writer(&self) -> Option<&OutputWriter> { 94 | if self.output_writer.is_some() { 95 | Some(self.output_writer.as_ref().unwrap()) 96 | } else { 97 | None 98 | } 99 | } 100 | 101 | fn simple_calculate(&self) { 102 | let elapsed_time = (self.start_time.elapsed().unwrap().as_millis() as u64).max(1); 103 | let last_print_time = self.last_print_time.swap(elapsed_time, Ordering::Relaxed); 104 | let millis = elapsed_time - last_print_time; 105 | 106 | let num_exec = self.num_exec.load(Ordering::Relaxed); 107 | let last_num_exec = self.last_num_exec.swap(num_exec, Ordering::Relaxed); 108 | 109 | self.cur_exec_speed.store( 110 | (num_exec - last_num_exec) * 1000 / millis.max(1), 111 | Ordering::Relaxed, 112 | ); 113 | 114 | let cpu_percent: u64 = { 115 | (get_cpu_percentage_collector() 116 | .lock() 117 | .unwrap() 118 | .cpu_percent_percpu() 119 | .unwrap() 120 | .iter() 121 | .take(self.num_core as usize) 122 | .sum::() 123 | / self.num_core as f32) as u64 124 | }; 125 | self.cpu_usage.store(cpu_percent, Ordering::Relaxed); 126 | self.all_cpu_usage.fetch_add(cpu_percent, Ordering::Relaxed); 127 | if cpu_percent != 0 { 128 | self.all_cpu_record_times.fetch_add(1, Ordering::Relaxed); 129 | } 130 | } 131 | 132 | pub fn simple_show(&self) { 133 | let elapsed_time = (self.start_time.elapsed().unwrap().as_millis() as u64).max(1); 134 | self.simple_calculate(); 135 | let num_exec = self.num_exec.load(Ordering::Relaxed); 136 | let total_exec = pretty_format_num(num_exec); 137 | let cur_exec_speed = pretty_format_num(self.cur_exec_speed.load(Ordering::Relaxed)); 138 | let average_speed = pretty_format_num(num_exec * 1000 / elapsed_time); 139 | let average_speed_per_core = 140 | pretty_format_num(num_exec * 1000 / elapsed_time.max(1) / self.num_core.max(1)); 141 | let timeout_exec = self.num_timeout_exec.load(Ordering::Relaxed); 142 | let timeout_exec_str = pretty_format_num(timeout_exec); 143 | let current_coverage = pretty_format_num(self.cur_coverage.load(Ordering::Relaxed)); 144 | let time = pretty_format_num(elapsed_time / 1000); 145 | let crash_num = pretty_format_num(self.num_crashes.load(Ordering::Relaxed)); 146 | let timeout_rate = timeout_exec as f64 / num_exec.max(1) as f64 * 100f64; 147 | 148 | let cpu_percent = self.cpu_usage.load(Ordering::Relaxed); 149 | let all_cpu_usage = self.all_cpu_usage.load(Ordering::Relaxed); 150 | let all_cpu_usage_record_times = self.all_cpu_record_times.load(Ordering::Relaxed); 151 | let cpu_average = all_cpu_usage as f64 / all_cpu_usage_record_times.max(1) as f64; 152 | // Assume the cpu are occupied from lowest id to highest. Otherwise this number is wrong. 153 | println!( 154 | "Time: {time}, \ 155 | Total exec: {total_exec}, current speed: {cur_exec_speed}/s, \ 156 | average speed: {average_speed}/s, per core: {average_speed_per_core}/s, \ 157 | timeout exec: {timeout_exec_str}, crash: {crash_num}, \ 158 | Interesting inputs {current_coverage}, timeout rate: {timeout_rate:.3}%, cpu usage: {cpu_percent}%, average cpu usage {cpu_average:.1}%" 159 | ); 160 | 161 | self.save_plot_data(); 162 | } 163 | 164 | pub fn save_plot_data(&self) { 165 | if let Some(output_writer) = self.output_writer.as_ref() { 166 | output_writer.write_plot_data(self).unwrap(); 167 | } 168 | } 169 | 170 | pub fn save_fuzzer_stat(&self) { 171 | if let Some(output_writer) = self.output_writer.as_ref() { 172 | output_writer.write_fuzzer_stat(self).unwrap(); 173 | } 174 | } 175 | 176 | pub fn sync_with_json_config(&mut self, json_file_path: String) { 177 | let json_str = fs::read_to_string(json_file_path).unwrap(); 178 | let config: Value = serde_json::from_str(&json_str).unwrap(); 179 | 180 | if let Some(v) = config.get("monitor_addr") { 181 | self.set_monitor_addr(v.as_str().unwrap().parse().unwrap()); 182 | } 183 | 184 | if let Some(v) = config.get("max_exec") { 185 | self.set_max_exec(v.as_u64().unwrap()); 186 | } 187 | 188 | if let Some(v) = config.get("core") { 189 | self.set_core_num(v.as_u64().unwrap()); 190 | } 191 | } 192 | 193 | fn set_monitor_addr(&mut self, addr: SocketAddr) { 194 | self.monitor_addr = addr; 195 | } 196 | 197 | pub fn sync_with_fuzzer_config(&mut self, fuzzer: &crate::FuzzerConfig) { 198 | self.set_cmd(fuzzer.cmd.clone()) 199 | .set_total_coverage(fuzzer.map_size as u64) 200 | .set_core_num(fuzzer.core as u64) 201 | .set_max_exec(fuzzer.max_exec) 202 | .set_max_fuzzing_time(fuzzer.max_fuzzing_time) 203 | .set_queue_manager_num(fuzzer.queuemgr as u64) 204 | .set_executor_num(fuzzer.executor as u64) 205 | .set_feedback_collector_num(fuzzer.feedback as u64); 206 | if let Some(output_dir) = fuzzer.output.as_ref() { 207 | self.output_writer = Some(OutputWriter::new(output_dir.clone()).unwrap()); 208 | } 209 | } 210 | 211 | pub fn get_output_dir(&self) -> Option { 212 | if self.output_writer.is_some() { 213 | Some(self.output_writer.as_ref().unwrap().get_root_path()) 214 | } else { 215 | None 216 | } 217 | } 218 | 219 | pub fn set_max_exec(&mut self, max_exec: u64) -> &mut Self { 220 | self.max_exec = max_exec; 221 | self 222 | } 223 | 224 | pub fn get_max_exec(&self) -> u64 { 225 | self.max_exec 226 | } 227 | pub fn set_executor_num(&mut self, executor_num: u64) -> &mut Self { 228 | self.num_executor = executor_num; 229 | self 230 | } 231 | 232 | pub fn get_executor_num(&self) -> u64 { 233 | self.num_executor 234 | } 235 | 236 | pub fn set_feedback_collector_num(&mut self, feedback_collector_num: u64) -> &mut Self { 237 | self.num_feedback_collector = feedback_collector_num; 238 | self 239 | } 240 | 241 | pub fn get_feebback_collector_num(&self) -> u64 { 242 | self.num_feedback_collector 243 | } 244 | 245 | pub fn set_queue_manager_num(&mut self, queue_manager_num: u64) -> &mut Self { 246 | self.num_queue_manager = queue_manager_num; 247 | self 248 | } 249 | 250 | pub fn get_queue_manager_num(&self) -> u64 { 251 | self.num_queue_manager 252 | } 253 | 254 | pub fn add_exec(&self, exec: u64) { 255 | self.num_exec.fetch_add(exec, Ordering::Relaxed); 256 | } 257 | 258 | pub fn get_exec(&self) -> u64 { 259 | self.num_exec.load(Ordering::Relaxed) 260 | } 261 | 262 | pub fn get_crash(&self) -> u64 { 263 | self.num_crashes.load(Ordering::Relaxed) 264 | } 265 | 266 | pub fn add_crash(&self, n_crash: u64) { 267 | self.num_crashes.fetch_add(n_crash, Ordering::Relaxed); 268 | } 269 | 270 | pub fn get_timeout_exec(&self) -> u64 { 271 | self.num_timeout_exec.load(Ordering::Relaxed) 272 | } 273 | 274 | pub fn add_timeout_exec(&self, n_timeout_exec: u64) { 275 | self.num_timeout_exec 276 | .fetch_add(n_timeout_exec, Ordering::Relaxed); 277 | } 278 | 279 | pub fn add_coverage(&self, coverage: u64) { 280 | self.cur_coverage.fetch_add(coverage, Ordering::Relaxed); 281 | } 282 | 283 | pub fn set_cmd(&mut self, cmd: String) -> &mut Self { 284 | self.cmd = cmd; 285 | self 286 | } 287 | 288 | pub fn set_core_num(&mut self, core_num: u64) -> &mut Self { 289 | self.num_core = core_num.min(*crate::util::AFFINITY_NUM as u64); 290 | self 291 | } 292 | 293 | pub fn get_core_num(&self) -> u64 { 294 | self.num_core 295 | } 296 | 297 | pub fn set_total_coverage(&mut self, total_coverage: u64) -> &mut Self { 298 | self.total_coverage = total_coverage; 299 | self 300 | } 301 | 302 | // Set max fuzzing time in seconds. The fuzzer will stop after fuzzing that period of time. 303 | pub fn set_max_fuzzing_time(&mut self, max_fuzzing_time: u64) -> &mut Self { 304 | self.max_fuzzing_time = max_fuzzing_time; 305 | self 306 | } 307 | 308 | pub fn get_max_fuzzing_time(&self) -> u64 { 309 | self.max_fuzzing_time 310 | } 311 | 312 | pub fn get_fuzzing_time(&self) -> u64 { 313 | self.start_time.elapsed().unwrap().as_secs() 314 | } 315 | 316 | // If the exec num or fuzzing time reach limit, then we should stop the fuzzer. 317 | pub fn stop_if_done(&self) { 318 | if self.get_exec() >= self.get_max_exec() 319 | || self.get_fuzzing_time() >= self.get_max_fuzzing_time() 320 | { 321 | self.save_fuzzer_stat(); 322 | println!("Fuzzing done. Have a good day!"); 323 | std::process::exit(0); 324 | } 325 | } 326 | } 327 | 328 | impl Default for FuzzerInfo { 329 | fn default() -> Self { 330 | Self::new() 331 | } 332 | } 333 | 334 | #[cfg(test)] 335 | mod tests { 336 | use super::*; 337 | use crate::util; 338 | 339 | #[test] 340 | fn read_fuzzer_info_from_config() { 341 | let mut fuzzer_info = FuzzerInfo::new(); 342 | fuzzer_info 343 | .sync_with_json_config(format!("{}/test_config.json", util::get_test_bin_path())); 344 | 345 | assert_eq!(fuzzer_info.get_max_exec(), 10); 346 | //assert_eq!(fuzzer_info.get_core_num(), 4); 347 | } 348 | } 349 | -------------------------------------------------------------------------------- /src/queue/manager.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::{self, ExecutionStatus, FeedbackData, TestCase, TestCaseID}; 2 | use crate::minimizer::TestCaseMinimizer; 3 | use crate::mutator::afl_mutator::DeterministicMutationPass; 4 | use crate::queue::QueueManager; 5 | use itertools::sorted; 6 | use rand::rngs::StdRng; 7 | use rand::Rng; 8 | use rand::SeedableRng; 9 | use std::cmp; 10 | use std::collections::HashMap; 11 | use std::collections::VecDeque; 12 | 13 | pub const DEFAULT_SCORE: i64 = 100000000; 14 | const SCORE_INTERESTING: i64 = 50; 15 | const SCORE_TIMEOUT: i64 = -700; 16 | const SCORE_CRASH: i64 = -50; // We assume that it is not typical for a input to find two different bugs. 17 | const SCORE_UNINTERESTING: i64 = -7; 18 | 19 | const MAX_SELECT_COUNT: u64 = 0xFF; 20 | const TRIGGER_DETERMINISTIC_THRESHOLD: u64 = 6; 21 | const TRIGGER_DETERMINISTIC_INTERVAL: u64 = 3; 22 | 23 | #[derive(Debug, Clone, Default)] 24 | pub struct TestCaseWithMeta { 25 | test_case: TestCase, 26 | // Indicate the quality of the test case. The higher it is, the more likely it should be picked for mutation. 27 | score: i64, 28 | select_count: u64, // How many times is the test case selected for mutation 29 | } 30 | 31 | impl PartialEq for TestCaseWithMeta { 32 | fn eq(&self, other: &Self) -> bool { 33 | self.score == other.score 34 | } 35 | } 36 | 37 | impl Eq for TestCaseWithMeta {} 38 | 39 | impl Ord for TestCaseWithMeta { 40 | fn cmp(&self, other: &Self) -> cmp::Ordering { 41 | other.score.cmp(&self.score) 42 | } 43 | } 44 | 45 | impl PartialOrd for TestCaseWithMeta { 46 | fn partial_cmp(&self, other: &Self) -> Option { 47 | Some(self.cmp(other)) 48 | } 49 | } 50 | 51 | impl TestCaseWithMeta { 52 | pub fn new(test_case: TestCase, score: i64) -> Self { 53 | Self { 54 | test_case, 55 | score, 56 | select_count: 0, 57 | } 58 | } 59 | 60 | pub fn get_test_case(&self) -> &TestCase { 61 | &self.test_case 62 | } 63 | 64 | pub fn add_score(&mut self, delta: i64) { 65 | self.score += delta; 66 | assert!(self.score > 0); 67 | } 68 | 69 | pub fn set_score(&mut self, delta: i64) { 70 | self.score = delta; 71 | assert!(self.score > 0); 72 | } 73 | 74 | pub fn select(&mut self) -> TestCase { 75 | self.increase_select_count(); 76 | 77 | self.select_by_count(self.select_count) 78 | //self.test_case.clone() 79 | } 80 | 81 | pub fn get_select_counter(&self) -> u64 { 82 | self.select_count 83 | } 84 | 85 | fn increase_select_count(&mut self) { 86 | if self.select_count < MAX_SELECT_COUNT { 87 | self.select_count += 1; 88 | } 89 | } 90 | 91 | fn select_by_count(&mut self, select_count: u64) -> TestCase { 92 | if select_count >= TRIGGER_DETERMINISTIC_THRESHOLD 93 | && select_count % TRIGGER_DETERMINISTIC_INTERVAL == 0 94 | { 95 | let mut test_case = self.test_case.clone(); 96 | let mut meta = test_case.take_meta().unwrap_or_default(); 97 | for flag in DeterministicMutationPass::all().iter() { 98 | if !meta.get_done_pass().contains(flag) { 99 | meta.set_done_pass(meta.get_done_pass() | flag); 100 | self.test_case.set_meta(meta.clone()); 101 | meta.set_todo_pass(flag); 102 | test_case.set_meta(meta); 103 | return test_case; 104 | } 105 | } 106 | } 107 | self.test_case.clone_without_meta() 108 | } 109 | } 110 | 111 | pub struct SimpleQueueManager { 112 | // Test cases that have been picked and are waiting for mutation. 113 | to_mutate_queue: VecDeque, 114 | // Interesting queues that we should save. 115 | interesting_queue: HashMap, 116 | rng: StdRng, 117 | minimizer: Option, 118 | } 119 | 120 | impl SimpleQueueManager { 121 | pub fn new(test_case_minimizer: Option) -> Self { 122 | SimpleQueueManager { 123 | to_mutate_queue: VecDeque::default(), 124 | interesting_queue: HashMap::default(), 125 | rng: StdRng::seed_from_u64(crate::datatype::get_id() as u64), 126 | minimizer: test_case_minimizer, 127 | } 128 | } 129 | 130 | pub fn set_test_case_score(&mut self, id: TestCaseID, score: i64) { 131 | if self.interesting_queue.contains_key(&id) { 132 | self.interesting_queue 133 | .get_mut(&id) 134 | .unwrap() 135 | .set_score(score); 136 | } 137 | } 138 | 139 | pub fn add_score_to_test_case(&mut self, id: TestCaseID, score: i64) { 140 | if self.interesting_queue.contains_key(&id) { 141 | self.interesting_queue 142 | .get_mut(&id) 143 | .unwrap() 144 | .add_score(score); 145 | } 146 | } 147 | 148 | fn add_new_test_case(&mut self, test_case: TestCase) { 149 | let mut test_case = test_case; 150 | let new_id = test_case.get_id(); 151 | if let Some(minimizer) = self.minimizer.as_mut() { 152 | test_case = minimizer.minimize(test_case); 153 | assert_eq!(test_case.get_id(), new_id); 154 | } 155 | assert!(!self.interesting_queue.contains_key(&new_id)); 156 | self.interesting_queue 157 | .insert(new_id, TestCaseWithMeta::new(test_case, DEFAULT_SCORE)); 158 | } 159 | } 160 | 161 | impl Default for SimpleQueueManager { 162 | fn default() -> Self { 163 | Self::new(None) 164 | } 165 | } 166 | 167 | impl QueueManager for SimpleQueueManager { 168 | fn receive_interesting_test_cases(&mut self, test_cases: Vec) { 169 | for mut test_case in test_cases.into_iter() { 170 | // Generate a unique id for the new test case 171 | test_case.gen_id(); 172 | self.add_new_test_case(test_case); 173 | } 174 | } 175 | 176 | fn process_test_case_feedbacks(&mut self, feedbacks: &[datatype::Feedback]) { 177 | let mut score_changes: HashMap = HashMap::default(); 178 | for feedback in feedbacks.iter() { 179 | let mutation_info = feedback.get_mutation_info().unwrap(); 180 | let pid = mutation_info.get_pid(); 181 | 182 | let counter = match feedback.borrow_data() { 183 | Some(FeedbackData::Counter(c)) => *c as i64, 184 | _ => { 185 | unreachable!(); 186 | } 187 | }; 188 | 189 | let score_change = counter 190 | * match feedback.get_status() { 191 | ExecutionStatus::Ok => SCORE_UNINTERESTING, 192 | ExecutionStatus::Interesting => SCORE_INTERESTING, 193 | ExecutionStatus::Timeout => SCORE_TIMEOUT, 194 | 195 | ExecutionStatus::Crash => SCORE_CRASH, 196 | }; 197 | self.add_score_to_test_case(pid, score_change); 198 | score_changes 199 | .entry(pid) 200 | .and_modify(|e| *e += score_change) 201 | .or_insert(score_change); 202 | } 203 | 204 | for (pid, score_change) in score_changes { 205 | self.add_score_to_test_case(pid, score_change); 206 | } 207 | } 208 | 209 | fn select_interesting_inputs(&mut self, num: usize) -> Vec { 210 | if self.interesting_queue.is_empty() { 211 | return Vec::default(); 212 | } 213 | 214 | while self.to_mutate_queue.len() < num { 215 | let len = self.interesting_queue.len(); 216 | self.to_mutate_queue.extend( 217 | sorted(self.interesting_queue.values_mut()) 218 | .filter_map(|test_case_with_meta| { 219 | // Make every test case have at least 50% chance of being selected 220 | if self.rng.gen_range(1..MAX_SELECT_COUNT * 2) 221 | > test_case_with_meta.get_select_counter() 222 | { 223 | Some(test_case_with_meta.select()) 224 | } else { 225 | None 226 | } 227 | }) 228 | .take((len / 2).max(1)), 229 | ); 230 | } 231 | self.to_mutate_queue.drain(0..num).collect() 232 | } 233 | } 234 | 235 | #[cfg(test)] 236 | mod tests { 237 | use super::*; 238 | use datatype::Feedback; 239 | use datatype::MutationInfo; 240 | 241 | #[test] 242 | #[serial_test::serial] 243 | fn simple_manager_correctly_add_new_interesting_inputs() { 244 | let mut simple_manager = SimpleQueueManager::default(); 245 | 246 | // Create 0x100 executed inputs which wait for feedback. 247 | simple_manager.receive_interesting_test_cases( 248 | (0..100) 249 | .map(|i| TestCase::new(vec![i; 1], i as u32)) 250 | .collect(), 251 | ); 252 | (0..10).for_each(|_| { 253 | assert_eq!(simple_manager.select_interesting_inputs(10).len(), 10); 254 | }) 255 | } 256 | 257 | #[test] 258 | #[serial_test::serial] 259 | fn queue_manager_select_inputs_with_highest_scores() { 260 | let mut simple_manager = SimpleQueueManager::default(); 261 | for i in 0..20 { 262 | simple_manager.add_new_test_case(TestCase::new(Vec::default(), i)); 263 | } 264 | 265 | let mut feedbacks = Vec::default(); 266 | 267 | // Make the test case from 5 to 14 most interesting. 268 | for i in 5..15 { 269 | feedbacks.push( 270 | Feedback::new(ExecutionStatus::Interesting) 271 | .set_mutation_info(MutationInfo::new(i, 1)) 272 | .set_data(FeedbackData::Counter(1)), 273 | ); 274 | } 275 | 276 | simple_manager.process_test_case_feedbacks(&feedbacks); 277 | let interesing_inputs = simple_manager.select_interesting_inputs(20); 278 | assert_eq!(interesing_inputs.len(), 20); 279 | let mut ids = interesing_inputs 280 | .iter() 281 | .take(10) 282 | .map(|a| a.get_id()) 283 | .collect::>(); 284 | 285 | ids.sort_unstable(); 286 | println!("{:?}", ids); 287 | assert!((5..15).eq(ids.into_iter())); 288 | } 289 | 290 | #[test] 291 | fn every_test_case_have_chances_of_being_selected_even_with_very_low_score() { 292 | let mut simple_manager = SimpleQueueManager::default(); 293 | for i in 0..10 { 294 | simple_manager.add_new_test_case(TestCase::new(Vec::default(), i)); 295 | } 296 | 297 | // Make the test case 0 very uninteresting. 298 | let feedbacks = vec![Feedback::new(ExecutionStatus::Timeout) 299 | .set_mutation_info(MutationInfo::new(0, 1)) 300 | .set_data(FeedbackData::Counter(100))]; 301 | 302 | simple_manager.process_test_case_feedbacks(&feedbacks); 303 | loop { 304 | if simple_manager 305 | .select_interesting_inputs(20) 306 | .iter() 307 | .any(|test_case| test_case.get_id() == 0) 308 | { 309 | break; 310 | } 311 | } 312 | } 313 | 314 | #[test] 315 | fn test_case_will_try_to_run_deterministic_pass_if_selected_too_many_times() { 316 | let mut simple_manager = SimpleQueueManager::default(); 317 | simple_manager.add_new_test_case(TestCase::new(Vec::default(), 0)); 318 | let test_cases = 319 | simple_manager.select_interesting_inputs(TRIGGER_DETERMINISTIC_THRESHOLD as usize - 1); 320 | assert_eq!( 321 | test_cases 322 | .iter() 323 | .filter(|test_case| test_case.has_meta()) 324 | .count(), 325 | 0, 326 | ); 327 | let test_case = simple_manager 328 | .select_interesting_inputs(1) 329 | .first() 330 | .unwrap() 331 | .clone(); 332 | assert!(test_case.has_meta()); 333 | assert!(test_case 334 | .borrow_meta() 335 | .unwrap() 336 | .get_todo_pass() 337 | .contains(DeterministicMutationPass::BITFLIP1)); 338 | let test_case = simple_manager 339 | .select_interesting_inputs(TRIGGER_DETERMINISTIC_INTERVAL as usize) 340 | .last() 341 | .unwrap() 342 | .clone(); 343 | assert!(test_case.has_meta()); 344 | assert!(test_case 345 | .borrow_meta() 346 | .unwrap() 347 | .get_todo_pass() 348 | .contains(DeterministicMutationPass::BITFLIP2)); 349 | } 350 | } 351 | -------------------------------------------------------------------------------- /src/feedback/newbit_filter.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::NewBit; 2 | 3 | const VIRGIN_BYTE: u8 = 0xff; 4 | 5 | static COUNT_CLASS_LOOKUP: [u8; 256] = [ 6 | 0, 1, 2, 4, 8, 8, 8, 8, 16, 16, 16, 16, 16, 16, 16, 16, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 7 | 32, 32, 32, 32, 32, 32, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 8 | 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 9 | 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 10 | 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 11 | 64, 64, 64, 64, 64, 64, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 12 | 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 13 | 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 14 | 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 15 | 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 16 | 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 17 | 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 18 | ]; 19 | 20 | const CLASSIFY_COUNT: bool = true; 21 | 22 | // Deduplicate visited bits. 23 | pub struct NewBitFilter { 24 | virgin_bits: Vec, 25 | } 26 | 27 | impl NewBitFilter { 28 | pub fn new(map_size: usize) -> Self { 29 | Self { 30 | virgin_bits: vec![VIRGIN_BYTE; map_size], 31 | } 32 | } 33 | 34 | pub fn byte_count(&self) -> usize { 35 | self.virgin_bits 36 | .iter() 37 | .filter(|&b| *b != VIRGIN_BYTE) 38 | .count() 39 | } 40 | 41 | fn has_new_bit(&self, bit: &NewBit) -> bool { 42 | self.has_new_bit_at(bit.index, bit.val) 43 | } 44 | 45 | fn has_new_bit_at(&self, index: usize, val: u8) -> bool { 46 | self.virgin_bits[index] & val != 0 47 | } 48 | 49 | // Try adding a new bit. If the bit has been visited, return false. Otherwise update the 50 | // virgin map and return true. 51 | fn try_update_bit(&mut self, bit: &NewBit) -> bool { 52 | if self.has_new_bit(bit) { 53 | self.virgin_bits[bit.index] &= !bit.val; 54 | true 55 | } else { 56 | false 57 | } 58 | } 59 | 60 | pub fn update_bits(&mut self, bits: &Vec) { 61 | for bit in bits { 62 | self.update_bit_at(bit.index, bit.val); 63 | } 64 | } 65 | 66 | fn update_bit_at(&mut self, idx: usize, val: u8) { 67 | self.virgin_bits[idx] &= !val; 68 | } 69 | 70 | // Try finding new bits and update the virgin map. If any, return true. 71 | pub fn try_update_bits(&mut self, bits: &[NewBit]) -> bool { 72 | let mut result = false; 73 | for bit in bits { 74 | result |= self.try_update_bit(bit); 75 | } 76 | result 77 | } 78 | 79 | // Filter all the visited bits and return the new ones. 80 | pub fn filter_old_bits(&self, new_bits: Vec) -> Option> { 81 | let result = new_bits 82 | .into_iter() 83 | .filter(|new_bit| self.has_new_bit(new_bit)) 84 | .collect::>(); 85 | 86 | if result.is_empty() { 87 | None 88 | } else { 89 | Some(result) 90 | } 91 | } 92 | 93 | // Filter all the visited bits, record and return the new ones. 94 | pub fn filter_old_bits_mut(&mut self, new_bits: Vec) -> Option> { 95 | let mut result = Vec::new(); 96 | for new_bit in new_bits.into_iter() { 97 | if self.try_update_bit(&new_bit) { 98 | result.push(new_bit); 99 | } 100 | } 101 | if result.is_empty() { 102 | None 103 | } else { 104 | Some(result) 105 | } 106 | } 107 | 108 | // Filter all the visited bits in raw bitmap and return the new ones. 109 | // Notice: This is only used for criterion and will be removed in the future. 110 | pub fn filter_old_version(&mut self, raw_info: &[u8]) -> Option> { 111 | assert_eq!(raw_info.len(), self.virgin_bits.len()); 112 | 113 | let mut result = Vec::new(); 114 | for (idx, b) in raw_info.iter().enumerate() { 115 | let new_bits = self.virgin_bits[idx] & *b; 116 | if new_bits != 0 { 117 | self.update_bit_at(idx, new_bits); 118 | result.push(NewBit::new(idx, new_bits)); 119 | } 120 | } 121 | if result.is_empty() { 122 | None 123 | } else { 124 | Some(result) 125 | } 126 | } 127 | 128 | // TODO: Why this is slower than filter + clear? 129 | pub fn filter_new_version(&mut self, raw_info: &mut [u8]) -> Option> { 130 | assert_eq!(raw_info.len(), self.virgin_bits.len()); 131 | 132 | //TODO: This is ugly. How to fix this. 133 | #[cfg(target_pointer_width = "64")] 134 | const WORD_SIZE: usize = 8; 135 | #[cfg(target_pointer_width = "64")] 136 | let mut raw_info_ptr: *mut u64 = raw_info.as_ptr() as *mut u64; 137 | #[cfg(target_pointer_width = "64")] 138 | let mut virgin_bits_ptr: *const u64 = self.virgin_bits.as_ptr() as *const u64; 139 | 140 | #[cfg(target_pointer_width = "32")] 141 | const WORD_SIZE: usize = 4; 142 | #[cfg(target_pointer_width = "32")] 143 | let mut raw_info_ptr: *mut u32 = raw_info.as_ptr() as *mut u32; 144 | #[cfg(target_pointer_width = "32")] 145 | let mut virgin_bits_ptr: *const u32 = self.virgin_bits.as_ptr() as *const u32; 146 | 147 | let mut result = Vec::new(); 148 | for i in (0..self.virgin_bits.len()).step_by(WORD_SIZE) { 149 | unsafe { 150 | if *raw_info_ptr != 0 { 151 | if CLASSIFY_COUNT { 152 | for j in 0..WORD_SIZE { 153 | if raw_info[i + j] != 0 { 154 | raw_info[i + j] = COUNT_CLASS_LOOKUP[raw_info[i + j] as usize]; 155 | } 156 | } 157 | } 158 | if (*raw_info_ptr & *virgin_bits_ptr) != 0 { 159 | for (j, v) in raw_info.iter().enumerate().skip(i).take(WORD_SIZE) { 160 | let new_bit = v & self.virgin_bits[j]; 161 | if new_bit != 0 { 162 | self.update_bit_at(j, new_bit); 163 | result.push(NewBit::new(j, new_bit)); 164 | } 165 | } 166 | } 167 | for j in 0..WORD_SIZE { 168 | raw_info[i + j] = 0; 169 | } 170 | } 171 | raw_info_ptr = raw_info_ptr.add(1); 172 | virgin_bits_ptr = virgin_bits_ptr.add(1); 173 | } 174 | } 175 | 176 | if result.is_empty() { 177 | None 178 | } else { 179 | Some(result) 180 | } 181 | } 182 | 183 | // Filter all the visited bits in raw bitmap and return the new ones. 184 | pub fn filter(&mut self, raw_info: &mut [u8]) -> Option> { 185 | assert_eq!(raw_info.len(), self.virgin_bits.len()); 186 | 187 | //TODO: This is ugly. How to fix this. 188 | #[cfg(target_pointer_width = "64")] 189 | const WORD_SIZE: usize = 8; 190 | #[cfg(target_pointer_width = "64")] 191 | let mut raw_info_ptr: *const u64 = raw_info.as_ptr() as *const u64; 192 | #[cfg(target_pointer_width = "64")] 193 | let mut virgin_bits_ptr: *const u64 = self.virgin_bits.as_ptr() as *const u64; 194 | 195 | #[cfg(target_pointer_width = "32")] 196 | const WORD_SIZE: usize = 4; 197 | #[cfg(target_pointer_width = "32")] 198 | let mut raw_info_ptr: *const u32 = raw_info.as_ptr() as *const u32; 199 | #[cfg(target_pointer_width = "32")] 200 | let mut virgin_bits_ptr: *const u32 = self.virgin_bits.as_ptr() as *const u32; 201 | 202 | let mut result = Vec::new(); 203 | for i in (0..self.virgin_bits.len()).step_by(WORD_SIZE) { 204 | unsafe { 205 | if *raw_info_ptr != 0 { 206 | if CLASSIFY_COUNT { 207 | for j in 0..WORD_SIZE { 208 | if raw_info[i + j] != 0 { 209 | raw_info[i + j] = COUNT_CLASS_LOOKUP[raw_info[i + j] as usize]; 210 | } 211 | } 212 | } 213 | if (*raw_info_ptr & *virgin_bits_ptr) != 0 { 214 | for (j, v) in raw_info.iter().enumerate().skip(i).take(WORD_SIZE) { 215 | let new_bit = v & self.virgin_bits[j]; 216 | if new_bit != 0 { 217 | self.update_bit_at(j, new_bit); 218 | result.push(NewBit::new(j, new_bit)); 219 | } 220 | } 221 | } 222 | } 223 | raw_info_ptr = raw_info_ptr.add(1); 224 | virgin_bits_ptr = virgin_bits_ptr.add(1); 225 | } 226 | } 227 | 228 | if result.is_empty() { 229 | None 230 | } else { 231 | Some(result) 232 | } 233 | } 234 | 235 | // Filter all the visited bits in raw bitmap and return the new ones. 236 | pub fn filter_without_chaning_virgin(&mut self, raw_info: &mut [u8]) -> Option> { 237 | assert_eq!(raw_info.len(), self.virgin_bits.len()); 238 | 239 | //TODO: This is ugly. How to fix this. 240 | #[cfg(target_pointer_width = "64")] 241 | const WORD_SIZE: usize = 8; 242 | #[cfg(target_pointer_width = "64")] 243 | let mut raw_info_ptr: *const u64 = raw_info.as_ptr() as *const u64; 244 | #[cfg(target_pointer_width = "64")] 245 | let mut virgin_bits_ptr: *const u64 = self.virgin_bits.as_ptr() as *const u64; 246 | 247 | #[cfg(target_pointer_width = "32")] 248 | const WORD_SIZE: usize = 4; 249 | #[cfg(target_pointer_width = "32")] 250 | let mut raw_info_ptr: *const u32 = raw_info.as_ptr() as *const u32; 251 | #[cfg(target_pointer_width = "32")] 252 | let mut virgin_bits_ptr: *const u32 = self.virgin_bits.as_ptr() as *const u32; 253 | 254 | let mut result = Vec::new(); 255 | for i in (0..self.virgin_bits.len()).step_by(WORD_SIZE) { 256 | unsafe { 257 | if *raw_info_ptr != 0 { 258 | if CLASSIFY_COUNT { 259 | for j in 0..WORD_SIZE { 260 | if raw_info[i + j] != 0 { 261 | raw_info[i + j] = COUNT_CLASS_LOOKUP[raw_info[i + j] as usize]; 262 | } 263 | } 264 | } 265 | if (*raw_info_ptr & *virgin_bits_ptr) != 0 { 266 | for (j, v) in raw_info.iter().enumerate().skip(i).take(WORD_SIZE) { 267 | let new_bit = v & self.virgin_bits[j]; 268 | if new_bit != 0 { 269 | result.push(NewBit::new(j, new_bit)); 270 | } 271 | } 272 | } 273 | } 274 | raw_info_ptr = raw_info_ptr.add(1); 275 | virgin_bits_ptr = virgin_bits_ptr.add(1); 276 | } 277 | } 278 | 279 | if result.is_empty() { 280 | None 281 | } else { 282 | Some(result) 283 | } 284 | } 285 | 286 | pub fn visited_bits_num(&self) -> u32 { 287 | let mut result = 0; 288 | for i in &self.virgin_bits { 289 | if *i != VIRGIN_BYTE { 290 | result += 1; 291 | } 292 | } 293 | result 294 | } 295 | } 296 | 297 | #[test] 298 | fn filter_correctly_classify_hit_count() { 299 | let map_size = 65536; 300 | let mut nbfilter = NewBitFilter::new(map_size); 301 | let mut raw_info = vec![0; map_size]; 302 | raw_info[0x1] = 0xe0; 303 | raw_info[0x10] = 0x4; 304 | raw_info[0x100] = 0xe1; 305 | 306 | let result = nbfilter.filter(&mut raw_info).unwrap(); 307 | assert_eq!(result.len(), 3); 308 | 309 | if CLASSIFY_COUNT { 310 | assert_eq!(result[0], NewBit::new(0x1, COUNT_CLASS_LOOKUP[0xe0])); 311 | assert_eq!(result[1], NewBit::new(0x10, COUNT_CLASS_LOOKUP[0x4])); 312 | assert_eq!(result[2], NewBit::new(0x100, COUNT_CLASS_LOOKUP[0xe1])); 313 | } else { 314 | assert_eq!(result[0], NewBit::new(0x1, 0xe0)); 315 | assert_eq!(result[1], NewBit::new(0x10, 0x4)); 316 | assert_eq!(result[2], NewBit::new(0x100, 0xe1)); 317 | } 318 | 319 | raw_info[0x1] = 0xf0; 320 | raw_info[0x10] = 0x0f; 321 | raw_info[0x100] = 0x80; 322 | let result = nbfilter.filter(&mut raw_info).unwrap(); 323 | if CLASSIFY_COUNT { 324 | assert_eq!(result.len(), 1); 325 | assert_eq!(result[0], NewBit::new(0x10, COUNT_CLASS_LOOKUP[0xf])); 326 | } else { 327 | assert_eq!(result.len(), 2); 328 | assert_eq!(result[0], NewBit::new(0x1, 0x10)); 329 | assert_eq!(result[1], NewBit::new(0x10, 0xB)); 330 | } 331 | } 332 | -------------------------------------------------------------------------------- /src/feedback/bitmap_collector.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype; 2 | use crate::datatype::ExecutionStatus; 3 | use crate::datatype::{Feedback, FeedbackData, MutationInfo, NewBit}; 4 | use crate::feedback::FeedbackCollector; 5 | use crate::feedback::NewBitFilter; 6 | use crate::util; 7 | use serde_json::{json, Value}; 8 | use std::collections::HashMap; 9 | 10 | //pub type BitmapFeedbackWithTestCase = crate::datatype::Feedback>; 11 | 12 | /* 13 | #[derive(Clone)] 14 | pub struct BitmapFeedbackWithTestCase { 15 | pub id: u32, 16 | pub new_bits: Vec<(usize, u8)>, 17 | pub test_case: TestCase, 18 | } 19 | */ 20 | 21 | impl Feedback { 22 | pub fn create_coverage_feedback( 23 | status: datatype::ExecutionStatus, 24 | test_case: Option, 25 | data: Option>, 26 | ) -> Self { 27 | let mut result = Self::new(status); 28 | if let Some(test_case) = test_case { 29 | result = result.set_test_case(test_case); 30 | } 31 | if let Some(data) = data { 32 | result = result.set_data(FeedbackData::new_coverage(data)); 33 | } 34 | result 35 | } 36 | } 37 | 38 | pub struct BitmapCollector { 39 | filter: NewBitFilter, 40 | interesting_test_cases: Vec, 41 | crash_test_cases: Vec, 42 | mutation_feedbacks: Vec, 43 | test_case_feedbacks: Vec, 44 | monitor_data: Vec, 45 | } 46 | 47 | impl BitmapCollector { 48 | pub fn new(n: usize) -> Self { 49 | BitmapCollector { 50 | filter: NewBitFilter::new(n), //interesing_inputs: Vec::default(), 51 | interesting_test_cases: Vec::default(), 52 | crash_test_cases: Vec::default(), 53 | mutation_feedbacks: Vec::default(), 54 | test_case_feedbacks: Vec::default(), 55 | monitor_data: Vec::default(), 56 | } 57 | } 58 | 59 | #[cfg(test)] 60 | fn visited_bytes_num(&self) -> u32 { 61 | self.filter.visited_bits_num() 62 | } 63 | 64 | pub fn retrieve_monitor_data(&mut self) -> Vec { 65 | self.monitor_data.split_off(0) 66 | } 67 | } 68 | 69 | impl FeedbackCollector for BitmapCollector { 70 | fn process_feedback(&mut self, feedbacks: Vec) { 71 | let mut uninteresting_counter: HashMap = HashMap::new(); 72 | let mut interesting_counter: HashMap = HashMap::new(); 73 | let mut timeout_counter: HashMap = HashMap::new(); 74 | let mut crash_counter: HashMap = HashMap::new(); 75 | let mut interesting_test_cases = Vec::default(); 76 | let mut crash_test_cases = Vec::default(); 77 | let mut crashes_vec = Vec::new(); 78 | 79 | let mut total_timeout: u64 = 0; 80 | let mut total_crash: u64 = 0; 81 | 82 | for mut feedback in feedbacks.into_iter() { 83 | assert!(feedback.is_valid()); 84 | let mutation_info = *feedback.get_mutation_info().unwrap(); 85 | match feedback.get_status() { 86 | ExecutionStatus::Ok => { 87 | *uninteresting_counter.entry(mutation_info).or_insert(1) += 1; 88 | } 89 | ExecutionStatus::Interesting => { 90 | let coverage = feedback.take_data().unwrap().into_new_coverage().unwrap(); 91 | if let Some(new_bits) = self.filter.filter_old_bits_mut(coverage) { 92 | assert!(feedback.contain_test_case()); 93 | interesting_test_cases 94 | .push(feedback.set_data(FeedbackData::NewCoverage(new_bits))); 95 | 96 | *interesting_counter.entry(mutation_info).or_insert(1) += 1; 97 | } else { 98 | // This interesting test case can be coverage by some others. So we skip it. 99 | *uninteresting_counter.entry(mutation_info).or_insert(1) += 1; 100 | } 101 | } 102 | ExecutionStatus::Crash => { 103 | assert!(feedback.contain_test_case()); 104 | let test_case_str = 105 | serde_json::to_string(feedback.borrow_test_case().unwrap()).unwrap(); 106 | crashes_vec.push(test_case_str); 107 | crash_test_cases.push(feedback); 108 | *crash_counter.entry(mutation_info).or_insert(1) += 1; 109 | total_crash += 1; 110 | } 111 | ExecutionStatus::Timeout => { 112 | *timeout_counter.entry(mutation_info).or_insert(1) += 1; 113 | total_timeout += 1; 114 | } 115 | }; 116 | } 117 | 118 | // generate monitor data if any 119 | if total_crash != 0 { 120 | self.monitor_data 121 | .push(json!({ "crash": total_crash, "testcases": crashes_vec })); 122 | } 123 | if total_timeout != 0 { 124 | self.monitor_data.push(json!({ "timeout": total_timeout })); 125 | } 126 | 127 | // Further separate them based on mutation info and parent id. 128 | for (status, counter) in [ 129 | (ExecutionStatus::Ok, uninteresting_counter), 130 | (ExecutionStatus::Interesting, interesting_counter), 131 | (ExecutionStatus::Timeout, timeout_counter), 132 | (ExecutionStatus::Crash, crash_counter), 133 | ] { 134 | let mut mutation_counter = HashMap::new(); 135 | let mut test_case_counter = HashMap::new(); 136 | for (mutation_info, count) in counter.into_iter() { 137 | if count == 0 { 138 | continue; 139 | } 140 | *mutation_counter 141 | .entry(mutation_info.get_mutation_id()) 142 | .or_insert(count) += count; 143 | *test_case_counter 144 | .entry(mutation_info.get_pid()) 145 | .or_insert(count) += count; 146 | } 147 | 148 | self.mutation_feedbacks 149 | .extend(mutation_counter.into_iter().map(|(k, v)| { 150 | Feedback::new(status) 151 | .set_mutation_info(MutationInfo::new(0, k)) 152 | .set_data(FeedbackData::Counter(v)) 153 | })); 154 | self.test_case_feedbacks 155 | .extend(test_case_counter.into_iter().map(|(k, v)| { 156 | Feedback::new(status) 157 | .set_mutation_info(MutationInfo::new(k, 0)) 158 | .set_data(FeedbackData::Counter(v)) 159 | })); 160 | } 161 | 162 | self.interesting_test_cases.extend(interesting_test_cases); 163 | self.crash_test_cases.extend(crash_test_cases); 164 | } 165 | 166 | fn get_interesting_test_cases(&mut self, num: Option) -> Vec { 167 | util::take_n_elements_from_vector(&mut self.interesting_test_cases, num) 168 | } 169 | 170 | fn get_crash_test_cases(&mut self, num: Option) -> Vec { 171 | util::take_n_elements_from_vector(&mut self.crash_test_cases, num) 172 | } 173 | 174 | fn get_mutation_feedbacks(&mut self, num: Option) -> Vec { 175 | util::take_n_elements_from_vector(&mut self.mutation_feedbacks, num) 176 | } 177 | 178 | fn get_test_case_feedbacks(&mut self, num: Option) -> Vec { 179 | util::take_n_elements_from_vector(&mut self.test_case_feedbacks, num) 180 | } 181 | } 182 | 183 | #[cfg(test)] 184 | mod tests { 185 | use super::*; 186 | use crate::datatype::ExecutionStatus; 187 | use crate::datatype::TestCase; 188 | 189 | #[test] 190 | fn virgin_map_tract_new_bit() { 191 | let mut feedback_collector = BitmapCollector::new(0x100); 192 | let mut feedbacks = Vec::new(); 193 | for i in 0..0x10usize { 194 | feedbacks.push( 195 | Feedback::create_coverage_feedback( 196 | ExecutionStatus::Interesting, 197 | Some(TestCase::new(vec![i as u8; 1], 0)), 198 | Some(vec![NewBit::new(i, 1), NewBit::new(i + 1, 1)]), 199 | ) 200 | .set_mutation_info(datatype::MutationInfo::new(0, 0)), 201 | ); 202 | } 203 | 204 | feedback_collector.process_feedback(feedbacks); 205 | assert_eq!(feedback_collector.visited_bytes_num(), 0x11); 206 | } 207 | 208 | #[test] 209 | fn feedback_collector_categorize_different_types_feedback() { 210 | let mut feedback_collector = BitmapCollector::new(0x100); 211 | let mut feedbacks = Vec::new(); 212 | 213 | // Add 10 interesting feedback of new coverage, this should produce 1 mutation feedback, 10 interesting test case feedback, 1 test case feedback. 214 | for i in 0..10usize { 215 | feedbacks.push( 216 | Feedback::create_coverage_feedback( 217 | ExecutionStatus::Interesting, 218 | Some(TestCase::new(vec![i as u8; 1], 0)), 219 | Some(vec![NewBit::new(i, 1), NewBit::new(i + 1, 1)]), 220 | ) 221 | .set_mutation_info(datatype::MutationInfo::new(0, 0)), 222 | ); 223 | } 224 | 225 | // Add 10 feedback of uninteresting from 2 mutator, this should produce two mutation feedback, 1 test case feedback. 226 | for i in 0..10usize { 227 | feedbacks.push( 228 | Feedback::create_coverage_feedback(ExecutionStatus::Ok, None, None) 229 | .set_mutation_info(datatype::MutationInfo::new(0, (i as u32) & 1)), 230 | ); 231 | } 232 | 233 | // Add 10 feedback of timeout from 2 mutator, this should produce two timeout feedback, 1 test case feedback. 234 | for i in 0..10usize { 235 | feedbacks.push( 236 | Feedback::create_coverage_feedback(ExecutionStatus::Timeout, None, None) 237 | .set_mutation_info(datatype::MutationInfo::new(0, (i as u32) & 1)), 238 | ); 239 | } 240 | // Add 10 feedback of timeout from 2 mutator, this should produce two crash feedback, 1 test case feedback, 2 crash test case feedback. 241 | for i in 0..10usize { 242 | feedbacks.push( 243 | Feedback::create_coverage_feedback( 244 | ExecutionStatus::Crash, 245 | Some(TestCase::new(vec![i as u8; 1], 0)), 246 | None, 247 | ) 248 | .set_mutation_info(datatype::MutationInfo::new(0, (i as u32) & 1)), 249 | ); 250 | } 251 | 252 | feedback_collector.process_feedback(feedbacks); 253 | assert_eq!(feedback_collector.get_mutation_feedbacks(None).len(), 7); 254 | assert_eq!(feedback_collector.get_test_case_feedbacks(None).len(), 4); 255 | assert_eq!(feedback_collector.get_crash_test_cases(None).len(), 10); 256 | assert_eq!( 257 | feedback_collector.get_interesting_test_cases(None).len(), 258 | 10 259 | ); 260 | } 261 | 262 | #[test] 263 | fn test_virgin_map_filter_seen_bit() { 264 | let mut feedback_collector = BitmapCollector::new(0x100); 265 | let mut feedbacks = Vec::new(); 266 | for i in 0..0x10usize { 267 | feedbacks.push( 268 | Feedback::create_coverage_feedback( 269 | ExecutionStatus::Interesting, 270 | Some(TestCase::new(vec![i as u8; 1], 0)), 271 | Some(vec![NewBit::new(i, 1), NewBit::new(i + 1, 1)]), 272 | ) 273 | .set_mutation_info(datatype::MutationInfo::new(0, 0)), 274 | ); 275 | } 276 | 277 | feedback_collector.process_feedback(feedbacks); 278 | let len1 = feedback_collector.visited_bytes_num(); 279 | 280 | let mut all_visited_new_bits = Vec::new(); 281 | for i in 0..0x11 { 282 | all_visited_new_bits.push(NewBit::new(i, 1)); 283 | } 284 | feedback_collector.process_feedback(vec![Feedback::create_coverage_feedback( 285 | ExecutionStatus::Interesting, 286 | Some(TestCase::new(vec![0u8; 1], 0)), 287 | Some(all_visited_new_bits.clone()), 288 | ) 289 | .set_mutation_info(datatype::MutationInfo::new(0, 0))]); 290 | assert_eq!(len1, feedback_collector.visited_bytes_num()); 291 | 292 | let len2 = feedback_collector.visited_bytes_num(); 293 | all_visited_new_bits.push(NewBit::new(0x11, 1)); 294 | feedback_collector.process_feedback(vec![Feedback::create_coverage_feedback( 295 | ExecutionStatus::Interesting, 296 | Some(TestCase::new(vec![0u8; 1], 0)), 297 | Some(all_visited_new_bits.clone()), 298 | ) 299 | .set_mutation_info(datatype::MutationInfo::new(0, 0))]); 300 | assert_eq!(feedback_collector.visited_bytes_num(), len2 + 1); 301 | } 302 | 303 | #[test] 304 | fn process_feedback_generate_monitor_data() { 305 | let mut feedback_collector = BitmapCollector::new(0x100); 306 | let mut feedbacks = Vec::new(); 307 | 308 | // Add 10 interesting feedback of new coverage, this should produce 1 mutation feedback, 10 interesting test case feedback, 1 test case feedback. 309 | for i in 0..10usize { 310 | feedbacks.push( 311 | Feedback::create_coverage_feedback( 312 | ExecutionStatus::Crash, 313 | Some(TestCase::new(vec![i as u8; 1], 0)), 314 | Some(vec![NewBit::new(i, 1), NewBit::new(i + 1, 1)]), 315 | ) 316 | .set_mutation_info(datatype::MutationInfo::new(0, 0)), 317 | ); 318 | } 319 | 320 | for i in 0..10usize { 321 | feedbacks.push( 322 | Feedback::create_coverage_feedback(ExecutionStatus::Timeout, None, None) 323 | .set_mutation_info(datatype::MutationInfo::new(0, (i as u32) & 1)), 324 | ); 325 | } 326 | 327 | feedback_collector.process_feedback(feedbacks); 328 | 329 | let monitor_data = feedback_collector.retrieve_monitor_data(); 330 | assert_eq!(monitor_data.len(), 2); 331 | for v in ["crash", "timeout"] { 332 | assert_eq!( 333 | monitor_data 334 | .iter() 335 | .find(|a| a.get(v).is_some()) 336 | .unwrap() 337 | .get(v) 338 | .unwrap(), 339 | 10 340 | ); 341 | } 342 | } 343 | } 344 | -------------------------------------------------------------------------------- /src/executor/frontend.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::{Feedback, TestCase}; 2 | use crate::executor::shmem::ShMem; 3 | use crate::executor::*; 4 | use crate::frontend::simple_frontend::{ 5 | AsyncFrontend, AsyncFrontendChannel, AsyncReceiver, AsyncSender, BasicFrontend, 6 | }; 7 | use crate::frontend::Worker; 8 | use crate::monitor::Monitor; 9 | use async_trait::async_trait; 10 | use std::sync::{Arc, Mutex}; 11 | 12 | pub type ForkServerExecutorResultPool = Arc>>; 13 | pub type ForkServerExecutorWorkerMutex = Arc>; 14 | pub type ForkServerExecutorWorkerPool = Vec; 15 | 16 | pub struct ForkServerExecutorWorker { 17 | executor: ForkServerExecutor, 18 | } 19 | 20 | pub struct ForkServerExecutorFrontend { 21 | worker_pool: ForkServerExecutorWorkerPool, // TODO(yongheng): Lock this with read/write lock 22 | bitmap_size: usize, 23 | args: Vec, 24 | sender: Option>>, 25 | output_receiver: Option>>, 26 | output_sender: Option>>, 27 | working_dir: Option, 28 | timeout: u64, 29 | } 30 | 31 | impl ForkServerExecutorWorker { 32 | pub fn new( 33 | bitmap_size: usize, 34 | args: Vec, 35 | timeout: u64, 36 | working_dir: Option, 37 | ) -> Self { 38 | let bitmap_tracer = BitmapTracer::new(bitmap_size); 39 | bitmap_tracer.bitmap.write_to_env("__AFL_SHM_ID").unwrap(); 40 | match ForkServerExecutor::new(bitmap_tracer, args, working_dir, timeout, true) { 41 | Ok(executor) => ForkServerExecutorWorker { executor }, 42 | Err(e) => { 43 | panic!("ForkServerExecutorWorker::new failed: {}", e); 44 | } 45 | } 46 | } 47 | } 48 | 49 | impl Worker for ForkServerExecutorWorker { 50 | type Input = TestCase; 51 | type Output = Feedback; 52 | 53 | fn handle_one_input(&mut self, input: Vec) -> Vec { 54 | let output_len = input.len(); 55 | let feedbacks = input 56 | .into_iter() 57 | .map(|single_input| self.executor.execute(single_input)) 58 | .collect::>(); 59 | crate::monitor::get_monitor() 60 | .read() 61 | .unwrap() 62 | .receive_statistics(serde_json::json!({ "exec": output_len as u64 })); 63 | feedbacks 64 | } 65 | } 66 | 67 | impl ForkServerExecutorFrontend { 68 | #[allow(dead_code)] 69 | pub fn new( 70 | bitmap_size: usize, 71 | args: Vec, 72 | timeout: u64, 73 | working_dir: Option, 74 | ) -> Self { 75 | ForkServerExecutorFrontend { 76 | worker_pool: Vec::default(), 77 | bitmap_size, 78 | args, 79 | sender: None, 80 | output_receiver: None, 81 | output_sender: None, 82 | working_dir, 83 | timeout, 84 | } 85 | } 86 | } 87 | 88 | impl BasicFrontend for ForkServerExecutorFrontend { 89 | type Worker = ForkServerExecutorWorker; 90 | type Output = Feedback; 91 | 92 | fn create_worker(&mut self, num: u32) { 93 | for _i in 0..num { 94 | self.worker_pool 95 | .push(Arc::new(Mutex::new(ForkServerExecutorWorker::new( 96 | self.bitmap_size, 97 | self.args.clone(), 98 | self.timeout, 99 | self.working_dir.clone(), 100 | )))); 101 | } 102 | } 103 | 104 | fn output_transfrom(&self, input: Vec) -> Vec { 105 | input 106 | } 107 | crate::frontend_default!(); 108 | } 109 | 110 | impl AsyncFrontendChannel for ForkServerExecutorFrontend { 111 | crate::async_frontend_default!(); 112 | } 113 | 114 | #[async_trait] 115 | impl AsyncFrontend for ForkServerExecutorFrontend {} 116 | 117 | #[cfg(test)] 118 | mod tests { 119 | use super::*; 120 | use crate::util; 121 | 122 | #[test] 123 | #[serial_test::serial] 124 | fn forkserver_executor_frontend_add_delete_worker_correctly() { 125 | const MAP_SIZE: usize = 65536; 126 | let args = vec![util::get_test_bin_by_name("simple_prog").unwrap()]; 127 | let mut toy_frontend = ForkServerExecutorFrontend::new(MAP_SIZE, args, 20, None); 128 | toy_frontend.create_worker(15); 129 | 130 | assert!(toy_frontend.worker_pool.len() == 15); 131 | toy_frontend.delete_worker(5); 132 | assert!(toy_frontend.worker_pool.len() == 10); 133 | } 134 | 135 | #[tokio::test] 136 | #[serial_test::serial] 137 | #[should_panic] 138 | async fn forkserver_executor_fronted_panic_when_no_worker_is_created() { 139 | const MAP_SIZE: usize = 65536; 140 | let args = vec![util::get_test_bin_by_name("simple_prog").unwrap()]; 141 | let toy_frontend = ForkServerExecutorFrontend::new(MAP_SIZE, args, 20, None); 142 | toy_frontend 143 | .handle_inputs(vec![TestCase::new(vec![0, 1], 0)]) 144 | .await; 145 | } 146 | } 147 | 148 | pub mod new_executor_frontend { 149 | use crate::executor::shmem::ShMem; 150 | use crate::executor::*; 151 | use crate::frontend::load_balance_frontend::*; 152 | use crate::frontend::FuzzerIO; 153 | 154 | pub struct ExecutorWorker { 155 | executor: super::ForkServerExecutor, 156 | execute_len: usize, 157 | } 158 | 159 | impl Worker for ExecutorWorker { 160 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 161 | match input { 162 | FuzzerIO::TestCase(data) => { 163 | let len = data.len(); 164 | let feedbacks = data 165 | .into_iter() 166 | .map(|test_case| self.executor.execute(test_case)) 167 | .collect(); 168 | let mut result = vec![FuzzerIO::Feedback(feedbacks)]; 169 | //if self.execute_len >= 1000 { 170 | result.push(FuzzerIO::MonitorData(vec![serde_json::json!({ 171 | "exec": len as u64 172 | })])); 173 | self.execute_len = 0; 174 | //} 175 | Some(result) 176 | } 177 | _ => { 178 | println!("{:?}", input.get_type()); 179 | unreachable!(); 180 | } 181 | } 182 | } 183 | } 184 | impl ExecutorWorker { 185 | pub fn new( 186 | bitmap_size: usize, 187 | args: Vec, 188 | timeout: u64, 189 | working_dir: Option, 190 | ) -> Self { 191 | let bitmap_tracer = super::BitmapTracer::new(bitmap_size); 192 | bitmap_tracer.bitmap.write_to_env("__AFL_SHM_ID").unwrap(); 193 | 194 | let executor = 195 | super::ForkServerExecutor::new(bitmap_tracer, args, working_dir, timeout, true) 196 | .ok() 197 | .unwrap(); 198 | Self { 199 | executor, 200 | execute_len: 0, 201 | } 202 | } 203 | } 204 | 205 | pub struct ExecutorTransformer {} 206 | 207 | impl Transformer for ExecutorTransformer {} 208 | 209 | pub type ExecutorFrontend = Frontend; 210 | 211 | #[cfg(test)] 212 | mod tests { 213 | use super::*; 214 | use crate::datatype::TestCase; 215 | use crate::frontend::FuzzerIOType; 216 | use crate::util; 217 | 218 | #[tokio::test(flavor = "multi_thread", worker_threads = 5)] 219 | #[serial_test::serial] 220 | async fn do_testing() { 221 | let bitmap_size: usize = 65536; 222 | let args = vec![util::get_test_bin_by_name("simple_prog").unwrap()]; 223 | let timeout = 20; 224 | let working_dir: Option = None; 225 | let closure = move || { 226 | ExecutorWorker::new(bitmap_size, args.clone(), timeout, working_dir.clone()) 227 | }; 228 | let mut frontend: ExecutorFrontend = ExecutorFrontend::new(Box::new(closure)); 229 | frontend.create_worker(5); 230 | let (input_sender, input_rx) = frontend_channel(100); 231 | let (output_sender, mut output_rx) = frontend_channel(100); 232 | let (monitor_output_sender, _monitor_output_rx) = frontend_channel(100); 233 | frontend.register_input_handler(Receiver::AsyncExclusive(input_rx)); 234 | frontend.register_output_handler( 235 | FuzzerIOType::Feedback, 236 | Sender::AsyncExclusive(output_sender), 237 | ); 238 | frontend.register_output_handler( 239 | FuzzerIOType::MonitorData, 240 | Sender::AsyncExclusive(monitor_output_sender), 241 | ); 242 | 243 | frontend.run(); 244 | 245 | let test_case = TestCase::new(String::from("select 1;").as_bytes().to_vec(), 0); 246 | let test_cases = (0..100) 247 | .map(|_x| test_case.clone()) 248 | .collect::>(); 249 | for _i in 0..10 { 250 | input_sender 251 | .send(FuzzerIO::TestCase(test_cases.clone())) 252 | .await 253 | .unwrap(); 254 | tokio::time::sleep(tokio::time::Duration::from_millis(5)).await; 255 | let output = output_rx.recv().await.unwrap(); 256 | match output.get_type() { 257 | FuzzerIOType::Feedback => { 258 | println!("Receive output"); 259 | } 260 | _ => { 261 | unreachable!(); 262 | } 263 | } 264 | } 265 | } 266 | } 267 | } 268 | 269 | pub mod work_stealing_executor_frontend { 270 | use crate::{ 271 | executor::{shmem::ShMem, Executor}, 272 | frontend::{work_stealing_frontend::*, FuzzerIO}, 273 | }; 274 | 275 | pub struct ExecutorWorker { 276 | executor: super::ForkServerExecutor, 277 | execute_len: usize, 278 | } 279 | 280 | impl WorkerImpl for ExecutorWorker { 281 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 282 | match input { 283 | FuzzerIO::TestCase(data) => { 284 | let len = data.len(); 285 | let feedbacks = data 286 | .into_iter() 287 | .map(|test_case| self.executor.execute(test_case)) 288 | .collect(); 289 | let mut result = vec![FuzzerIO::Feedback(feedbacks)]; 290 | //if self.execute_len >= 1000 { 291 | result.push(FuzzerIO::MonitorData(vec![serde_json::json!({ 292 | "exec": len as u64 293 | })])); 294 | self.execute_len = 0; 295 | //} 296 | Some(result) 297 | } 298 | _ => { 299 | println!("{:?}", input.get_type()); 300 | unreachable!(); 301 | } 302 | } 303 | } 304 | } 305 | 306 | impl ExecutorWorker { 307 | pub fn new( 308 | bitmap_size: usize, 309 | args: Vec, 310 | timeout: u64, 311 | working_dir: Option, 312 | ) -> Self { 313 | let bitmap_tracer = super::BitmapTracer::new(bitmap_size); 314 | bitmap_tracer.bitmap.write_to_env("__AFL_SHM_ID").unwrap(); 315 | 316 | let executor = 317 | super::ForkServerExecutor::new(bitmap_tracer, args, working_dir, timeout, true) 318 | .ok() 319 | .unwrap(); 320 | Self { 321 | executor, 322 | execute_len: 0, 323 | } 324 | } 325 | } 326 | 327 | #[cfg(test)] 328 | mod tests { 329 | use super::*; 330 | use crate::datatype::TestCase; 331 | use crate::frontend::FuzzerIOType; 332 | use crate::util; 333 | 334 | #[tokio::test(flavor = "multi_thread", worker_threads = 5)] 335 | #[serial_test::serial] 336 | async fn do_testing() { 337 | let bitmap_size: usize = 65536; 338 | let args = vec![util::get_test_bin_by_name("simple_prog").unwrap()]; 339 | let timeout = 20; 340 | let working_dir: Option = None; 341 | let closure = move || { 342 | Worker::new( 343 | Box::new(ExecutorWorker::new( 344 | bitmap_size, 345 | args.clone(), 346 | timeout, 347 | working_dir.clone(), 348 | )), 349 | 0, 350 | ) 351 | }; 352 | let frontend: Frontend = FrontendBuilder::new() 353 | .worker_creator(Box::new(closure)) 354 | .worker_num(2) 355 | .add_input_type(FuzzerIOType::TestCase) 356 | .add_output_type(FuzzerIOType::Feedback) 357 | .add_output_type(FuzzerIOType::MonitorData) 358 | .build() 359 | .unwrap(); 360 | //WorkStealingFrontend::new(Box::new(closure)); 361 | // frontend.create_worker(5); 362 | let (all_senders, all_receivers) = frontend.self_setup(); 363 | 364 | assert_eq!(all_senders.len(), 1); 365 | assert_eq!(all_senders[0].len(), 2); 366 | assert_eq!(all_receivers.len(), 2); 367 | frontend.run(); 368 | 369 | let test_case = TestCase::new(String::from("select 1;").as_bytes().to_vec(), 0); 370 | let test_cases = (0..100) 371 | .map(|_x| test_case.clone()) 372 | .collect::>(); 373 | 374 | let test_case_sender = all_senders[0][0].clone(); 375 | let feedback_receiver = all_receivers[0][0].clone(); 376 | let monitor_receiver = all_receivers[1][0].clone(); 377 | assert!(!feedback_receiver.is_closed()); 378 | assert!(!monitor_receiver.is_closed()); 379 | assert!(!test_case_sender.is_closed()); 380 | 381 | test_case_sender 382 | .send(FuzzerIO::TestCase(test_cases.clone())) 383 | .await 384 | .unwrap(); 385 | assert!(!test_case_sender.is_closed()); 386 | 387 | let feedback = feedback_receiver.recv().await.unwrap(); 388 | let monitor_data = monitor_receiver.recv().await.unwrap(); 389 | assert_eq!(feedback.get_type(), FuzzerIOType::Feedback); 390 | assert_eq!(monitor_data.get_type(), FuzzerIOType::MonitorData); 391 | tokio::time::sleep(tokio::time::Duration::from_millis(5)).await; 392 | test_case_sender 393 | .send(FuzzerIO::TestCase(test_cases.clone())) 394 | .await 395 | .unwrap(); 396 | let feedback = feedback_receiver.recv().await.unwrap(); 397 | let monitor_data = monitor_receiver.recv().await.unwrap(); 398 | assert_eq!(feedback.get_type(), FuzzerIOType::Feedback); 399 | assert_eq!(monitor_data.get_type(), FuzzerIOType::MonitorData); 400 | tokio::time::sleep(tokio::time::Duration::from_millis(5)).await; 401 | } 402 | } 403 | } 404 | -------------------------------------------------------------------------------- /src/queue/frontend.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::{Feedback, TestCase}; 2 | use crate::frontend::simple_frontend::{ 3 | AsyncFrontend, AsyncFrontendChannel, AsyncReceiver, AsyncSender, BasicFrontend, 4 | }; 5 | use crate::frontend::Worker; 6 | use crate::frontend_default; 7 | use crate::queue::manager::SimpleQueueManager; 8 | use crate::queue::QueueManager; 9 | use async_trait::async_trait; 10 | use rayon::prelude::*; 11 | use std::sync::{Arc, Mutex}; 12 | 13 | type SimpleQueueManagerWorkerPool = Vec>>; 14 | type SimpleManagerT = SimpleQueueManager; 15 | 16 | #[derive(Default)] 17 | pub struct SimpleQueueManagerWorker { 18 | manager: SimpleManagerT, 19 | } 20 | 21 | pub struct SimpleQueueManagerFrontend { 22 | worker_pool: SimpleQueueManagerWorkerPool, 23 | sender: Option>>, 24 | output_receiver: Option>>, 25 | output_sender: Option>>, 26 | } 27 | 28 | impl SimpleQueueManagerWorker { 29 | pub fn new() -> Self { 30 | Self::default() 31 | } 32 | 33 | pub fn process_feedbacks(&mut self, feedbacks: &[Feedback]) { 34 | self.manager.process_test_case_feedbacks(feedbacks); 35 | } 36 | 37 | pub fn select_interesting_inputs(&mut self, num: usize) -> Vec { 38 | self.manager.select_interesting_inputs(num) 39 | } 40 | } 41 | 42 | impl Worker for SimpleQueueManagerWorker { 43 | type Input = TestCase; 44 | type Output = TestCase; 45 | 46 | // TODO: This should be non-blocking frontend 47 | fn handle_one_input(&mut self, input: Vec) -> Vec { 48 | self.manager.receive_interesting_test_cases(input); 49 | Vec::default() 50 | } 51 | } 52 | 53 | impl SimpleQueueManagerFrontend { 54 | #[allow(dead_code)] 55 | pub fn new() -> Self { 56 | SimpleQueueManagerFrontend { 57 | worker_pool: Vec::default(), 58 | sender: None, 59 | output_receiver: None, 60 | output_sender: None, 61 | } 62 | } 63 | 64 | fn collect_interesting_inputs(&self, num: usize) -> Vec { 65 | let mut result = Vec::new(); 66 | let average = 67 | num / self.get_worker_pool().len() + 1.min(num % self.get_worker_pool().len()); 68 | assert_ne!(average, 0); 69 | while result.len() < num { 70 | for worker in self.get_worker_pool() { 71 | let mut worker = worker.lock().unwrap(); 72 | result.extend( 73 | worker 74 | .select_interesting_inputs(average.min(num - result.len())) 75 | .into_iter(), 76 | ); 77 | } 78 | } 79 | assert_eq!(result.len(), num); 80 | result 81 | } 82 | 83 | pub async fn process_feedbacks(&self, inputs: Vec) { 84 | if inputs.is_empty() { 85 | return; 86 | } 87 | 88 | self.get_worker_pool().par_iter().for_each(|worker| { 89 | let worker = worker.clone(); 90 | let mut worker = worker.lock().unwrap(); 91 | worker.process_feedbacks(&inputs); 92 | }); 93 | } 94 | } 95 | 96 | impl Default for SimpleQueueManagerFrontend { 97 | fn default() -> Self { 98 | Self::new() 99 | } 100 | } 101 | 102 | impl BasicFrontend for SimpleQueueManagerFrontend { 103 | type Worker = SimpleQueueManagerWorker; 104 | type Output = TestCase; 105 | fn create_worker(&mut self, num: u32) { 106 | self.worker_pool 107 | .extend((0..num).map(|_| Arc::new(Mutex::new(SimpleQueueManagerWorker::new())))); 108 | } 109 | 110 | fn output_transfrom(&self, input: Vec) -> Vec { 111 | input 112 | } 113 | frontend_default!(); 114 | } 115 | 116 | impl AsyncFrontendChannel for SimpleQueueManagerFrontend { 117 | crate::async_frontend_default!(); 118 | } 119 | 120 | #[async_trait] 121 | impl AsyncFrontend for SimpleQueueManagerFrontend { 122 | async fn get_results(&self, num: Option) -> Vec { 123 | // TODO: FIX this pure magic number. 124 | let max_num = self.get_worker_pool().len() * 2; 125 | let num = match num { 126 | Some(num) => (num as usize).min(max_num), 127 | None => max_num, 128 | }; 129 | self.collect_interesting_inputs(num) 130 | } 131 | } 132 | 133 | #[cfg(test)] 134 | mod tests { 135 | use super::*; 136 | 137 | #[test] 138 | #[serial_test::serial] 139 | fn simple_queue_manager_frontend_add_delete_worker_correctly() { 140 | let mut simple_queue_manager_frontend = SimpleQueueManagerFrontend::new(); 141 | simple_queue_manager_frontend.create_worker(15); 142 | assert!(simple_queue_manager_frontend.worker_pool.len() == 15); 143 | simple_queue_manager_frontend.delete_worker(5); 144 | assert!(simple_queue_manager_frontend.worker_pool.len() == 10); 145 | } 146 | } 147 | 148 | pub mod new_queue_frontend { 149 | use crate::datatype::{ExecutionStatus, FeedbackData}; 150 | use crate::datatype::{TestCase, TestCaseID}; 151 | use crate::frontend::load_balance_frontend::*; 152 | use crate::frontend::{FuzzerIO, FuzzerIOType}; 153 | use crate::minimizer::TestCaseMinimizer; 154 | use crate::queue::QueueManager; 155 | use std::collections::HashMap; 156 | 157 | pub struct QueueManagerWorker { 158 | manager: super::SimpleQueueManager, 159 | } 160 | 161 | impl Worker for QueueManagerWorker { 162 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 163 | match input { 164 | FuzzerIO::TestCase(data) => { 165 | let len = data.len(); 166 | self.manager.receive_interesting_test_cases(data); 167 | Some(vec![FuzzerIO::MonitorData(vec![serde_json::json!({ 168 | "interesting_test_case": len as u64 169 | })])]) 170 | } 171 | FuzzerIO::TestCaseScoreChange(data) => { 172 | for (id, score_change) in data.into_iter() { 173 | self.manager.set_test_case_score(id, score_change); 174 | } 175 | None 176 | } 177 | _ => { 178 | unreachable!(); 179 | } 180 | } 181 | } 182 | 183 | fn non_blocking_generate_output(&mut self, gen_type: FuzzerIOType) -> Option { 184 | match gen_type { 185 | FuzzerIOType::TestCase => { 186 | let outputs = self.manager.select_interesting_inputs(20); 187 | if outputs.is_empty() { 188 | None 189 | } else { 190 | Some(FuzzerIO::TestCase(outputs)) 191 | } 192 | } 193 | _ => { 194 | unreachable!(); 195 | } 196 | } 197 | } 198 | } 199 | 200 | impl QueueManagerWorker { 201 | pub fn new(test_case_minimizer: Option) -> Self { 202 | Self { 203 | manager: super::SimpleQueueManager::new(test_case_minimizer), 204 | } 205 | } 206 | 207 | pub fn new_with_seeds( 208 | initial_corpus: Option>, 209 | test_case_minimizer: Option, 210 | ) -> Self { 211 | let mut manager = super::SimpleQueueManager::new(test_case_minimizer); 212 | if let Some(initial_corpus) = initial_corpus { 213 | manager.receive_interesting_test_cases(initial_corpus); 214 | } 215 | Self { manager } 216 | } 217 | } 218 | 219 | impl Default for QueueManagerWorker { 220 | fn default() -> Self { 221 | Self::new(None) 222 | } 223 | } 224 | 225 | #[derive(Default, Clone)] 226 | pub struct QueueManagerTransformer {} 227 | 228 | impl Transformer for QueueManagerTransformer {} 229 | 230 | #[derive(Default)] 231 | pub struct QueueSyncSource { 232 | queue_scores: HashMap, 233 | scores_to_send_out: Vec<(TestCaseID, i64)>, 234 | } 235 | 236 | impl SyncSource for QueueSyncSource { 237 | fn update_status(&mut self, input: FuzzerIO) { 238 | if let FuzzerIO::TestCaseFeedback(feedbacks) = input { 239 | const DEFAULT_SCORE: i64 = 100000000; 240 | const SCORE_INTERESTING: i64 = 20; 241 | const SCORE_TIMEOUT: i64 = -1000; 242 | const SCORE_CRASH: i64 = -50; 243 | const SCORE_UNINTERESTING: i64 = -1; 244 | 245 | for feedback in feedbacks.iter() { 246 | let mutation_info = feedback.get_mutation_info().unwrap(); 247 | let pid = mutation_info.get_pid(); 248 | 249 | let counter = match feedback.borrow_data() { 250 | Some(FeedbackData::Counter(c)) => *c as i64, 251 | _ => { 252 | unreachable!(); 253 | } 254 | }; 255 | 256 | let score_change = counter 257 | * match feedback.get_status() { 258 | ExecutionStatus::Ok => SCORE_UNINTERESTING, 259 | ExecutionStatus::Interesting => SCORE_INTERESTING, 260 | ExecutionStatus::Timeout => SCORE_TIMEOUT, 261 | 262 | ExecutionStatus::Crash => SCORE_CRASH, 263 | }; 264 | self.queue_scores 265 | .entry(pid) 266 | .and_modify(|e| *e += score_change) 267 | .or_insert(score_change + DEFAULT_SCORE); 268 | } 269 | } else { 270 | unreachable!(); 271 | } 272 | } 273 | fn get_status(&mut self) -> FuzzerIO { 274 | if self.scores_to_send_out.len() < 1000 { 275 | self.scores_to_send_out 276 | .extend(self.queue_scores.iter().map(|(id, score)| (*id, *score))); 277 | } 278 | /* 279 | FuzzerIO::TestCaseScoreChange( 280 | self.queue_scores 281 | .iter() 282 | .map(|(id, score)| (*id, *score)) 283 | .collect(), 284 | ) 285 | */ 286 | FuzzerIO::TestCaseScoreChange( 287 | self.scores_to_send_out 288 | .split_off(self.scores_to_send_out.len().saturating_sub(1000)), 289 | ) 290 | } 291 | } 292 | 293 | pub type QueueManagerFrontend = 294 | Frontend; 295 | 296 | #[cfg(test)] 297 | mod tests { 298 | use super::*; 299 | use crate::datatype::TestCase; 300 | 301 | #[tokio::test(flavor = "multi_thread", worker_threads = 20)] 302 | #[serial_test::serial] 303 | async fn queue_manager_load_balance_frontend_test() { 304 | let closure = move || QueueManagerWorker::new(None); 305 | let mut frontend = QueueManagerFrontend::new(Box::new(closure)); 306 | let worker_num = 5; 307 | frontend.create_worker(worker_num); 308 | let (test_case_feedback_sender, test_case_feedback_rx) = frontend_channel(100); 309 | let (input_sender, input_rx) = frontend_channel(100); 310 | let (output_sender, mut output_rx) = frontend_channel(100); 311 | frontend.register_input_handler(Receiver::AsyncExclusive(input_rx)); 312 | 313 | frontend.register_watch_input_handler( 314 | Receiver::AsyncExclusive(test_case_feedback_rx), 315 | QueueSyncSource::default(), 316 | ); 317 | 318 | frontend.register_non_blocking_output_handler( 319 | FuzzerIOType::TestCase, 320 | Sender::AsyncExclusive(output_sender), 321 | ); 322 | 323 | frontend.run(); 324 | 325 | let mut feedbacks = Vec::default(); 326 | for i in 5..15 { 327 | feedbacks.push( 328 | crate::datatype::Feedback::new(crate::datatype::ExecutionStatus::Interesting) 329 | .set_mutation_info(crate::datatype::MutationInfo::new(i % 2, 1)) 330 | .set_data(crate::datatype::FeedbackData::Counter(1)), 331 | ); 332 | } 333 | 334 | let test_cases = (0..100) 335 | .map(|id| TestCase::new(String::from("select 1;").as_bytes().to_vec(), id)) 336 | .collect::>(); 337 | input_sender 338 | .send(FuzzerIO::TestCase(test_cases.clone())) 339 | .await 340 | .unwrap(); 341 | for _i in 0..10 { 342 | test_case_feedback_sender 343 | .send(FuzzerIO::TestCaseFeedback(feedbacks.clone())) 344 | .await 345 | .unwrap(); 346 | } 347 | for _i in 0..10 { 348 | let output = output_rx.recv().await.unwrap(); 349 | match output.get_type() { 350 | FuzzerIOType::TestCase => { 351 | //println!("Receive output {} ", _i); 352 | } 353 | _ => { 354 | unreachable!(); 355 | } 356 | } 357 | } 358 | } 359 | } 360 | } 361 | 362 | pub mod work_stealing_queue_manager { 363 | use crate::datatype::TestCase; 364 | use crate::frontend::load_balance_frontend::SyncSource; 365 | use crate::frontend::work_stealing_frontend::WorkerImpl; 366 | use crate::frontend::{FuzzerIO, FuzzerIOType}; 367 | use crate::minimizer::TestCaseMinimizer; 368 | use crate::queue::QueueManager; 369 | 370 | pub struct QueueManagerSyncSource { 371 | inner: super::new_queue_frontend::QueueSyncSource, 372 | counter: usize, 373 | } 374 | 375 | impl QueueManagerSyncSource { 376 | pub fn new() -> Self { 377 | Self { 378 | inner: super::new_queue_frontend::QueueSyncSource::default(), 379 | counter: 0, 380 | } 381 | } 382 | } 383 | 384 | impl Default for QueueManagerSyncSource { 385 | fn default() -> Self { 386 | Self::new() 387 | } 388 | } 389 | 390 | impl WorkerImpl for QueueManagerSyncSource { 391 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 392 | self.counter += 1; 393 | self.inner.update_status(input); 394 | if self.counter % 20 == 0 { 395 | self.counter = 0; 396 | Some(vec![self.inner.get_status()]) 397 | } else { 398 | None 399 | } 400 | } 401 | } 402 | pub struct QueueManagerWorker { 403 | manager: super::SimpleQueueManager, 404 | } 405 | 406 | impl WorkerImpl for QueueManagerWorker { 407 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 408 | match input { 409 | FuzzerIO::TestCase(data) => { 410 | let len = data.len(); 411 | self.manager.receive_interesting_test_cases(data); 412 | Some(vec![FuzzerIO::MonitorData(vec![serde_json::json!({ 413 | "interesting_test_case": len as u64 414 | })])]) 415 | } 416 | FuzzerIO::TestCaseScoreChange(data) => { 417 | for (id, score_change) in data.into_iter() { 418 | self.manager.set_test_case_score(id, score_change); 419 | } 420 | None 421 | } 422 | _ => { 423 | unreachable!(); 424 | } 425 | } 426 | } 427 | 428 | fn generate_outputs(&mut self, gen_type: FuzzerIOType) -> Option { 429 | match gen_type { 430 | FuzzerIOType::TestCase => { 431 | let outputs = self.manager.select_interesting_inputs(20); 432 | if outputs.is_empty() { 433 | None 434 | } else { 435 | Some(FuzzerIO::TestCase(outputs)) 436 | } 437 | } 438 | _ => { 439 | println!("Generate outputs for {:?}", gen_type); 440 | unreachable!(); 441 | } 442 | } 443 | } 444 | } 445 | 446 | impl QueueManagerWorker { 447 | pub fn new( 448 | initial_corpus: Option>, 449 | test_case_minimizer: Option, 450 | ) -> Self { 451 | let mut manager = super::SimpleQueueManager::new(test_case_minimizer); 452 | if let Some(initial_corpus) = initial_corpus { 453 | manager.receive_interesting_test_cases(initial_corpus); 454 | } 455 | Self { manager } 456 | } 457 | } 458 | 459 | impl Default for QueueManagerWorker { 460 | fn default() -> Self { 461 | Self::new(None, None) 462 | } 463 | } 464 | } 465 | -------------------------------------------------------------------------------- /src/mutator/frontend.rs: -------------------------------------------------------------------------------- 1 | use crate::datatype::Feedback; 2 | use crate::datatype::TestCase; 3 | use crate::frontend::simple_frontend::{ 4 | AsyncFrontend, AsyncFrontendChannel, AsyncReceiver, AsyncSender, BasicFrontend, 5 | }; 6 | use crate::frontend::Worker; 7 | use crate::mutator::afl_mutator::BitFlipMutator; 8 | use crate::mutator::Mutator; 9 | use async_trait::async_trait; 10 | use dashmap::DashSet; 11 | use rayon::prelude::*; 12 | use std::collections::hash_map::DefaultHasher; 13 | use std::hash::{Hash, Hasher}; 14 | use std::sync::atomic::{AtomicU64, Ordering}; 15 | use std::sync::{Arc, Mutex}; 16 | 17 | pub type BitFlipMutatorResultPool = Arc>>; 18 | pub type BitFlipMutatorWorkerPool = Vec>>; 19 | 20 | #[derive(Default)] 21 | pub struct BitFlipMutatorWorker { 22 | mutator: BitFlipMutator, 23 | mutate_round: usize, 24 | } 25 | 26 | #[derive(Default)] 27 | pub struct BitFlipMutatorFrontend { 28 | worker_pool: BitFlipMutatorWorkerPool, 29 | unique_hashes: DashSet, 30 | sender: Option>>, 31 | output_receiver: Option>>, 32 | output_sender: Option>>, 33 | total_mutation: AtomicU64, 34 | } 35 | 36 | const DEFAULT_ROUND: usize = 50; 37 | impl BitFlipMutatorWorker { 38 | pub fn new() -> Self { 39 | Self { 40 | mutator: BitFlipMutator::new(), 41 | mutate_round: DEFAULT_ROUND, 42 | } 43 | } 44 | 45 | pub fn process_feedbacks(&mut self, feedbacks: &[Feedback]) { 46 | self.mutator.process_mutation_feedback(feedbacks); 47 | } 48 | } 49 | 50 | impl Worker for BitFlipMutatorWorker { 51 | type Input = TestCase; 52 | type Output = TestCase; 53 | fn handle_one_input(&mut self, input: Vec) -> Vec { 54 | let mut new_test_cases = Vec::with_capacity(self.mutate_round * input.len()); 55 | for single_input in input { 56 | self.mutator.mutate(&single_input, self.mutate_round); 57 | new_test_cases.extend(self.mutator.get_mutated_test_cases(None).into_iter()); 58 | } 59 | new_test_cases 60 | } 61 | } 62 | 63 | impl BitFlipMutatorFrontend { 64 | #[allow(dead_code)] 65 | pub fn new() -> Self { 66 | BitFlipMutatorFrontend { 67 | worker_pool: Vec::default(), 68 | unique_hashes: DashSet::new(), 69 | sender: None, 70 | output_receiver: None, 71 | output_sender: None, 72 | total_mutation: AtomicU64::new(0), 73 | } 74 | } 75 | 76 | // Deduplicate inputs by hashing the data buffer in TestCase. 77 | fn deep_filter(&self, test_cases: Vec) -> Vec { 78 | fn calculate_test_case_hash(t: &TestCase) -> u64 { 79 | let mut s = DefaultHasher::new(); 80 | t.get_buffer().hash(&mut s); 81 | s.finish() 82 | } 83 | 84 | let total_mutation = self.total_mutation.load(Ordering::SeqCst) + test_cases.len() as u64; 85 | self.total_mutation.store(total_mutation, Ordering::SeqCst); 86 | 87 | let unique_hashes = &self.unique_hashes; 88 | test_cases 89 | .into_par_iter() 90 | .filter(|test_case| { 91 | let hash_sum = calculate_test_case_hash(test_case); 92 | unique_hashes.insert(hash_sum) 93 | }) 94 | .collect() 95 | /* 96 | println!( 97 | "Total mutation: {}, unique mutation {}, effective rate {}", 98 | total_mutation, 99 | unique_hashes.len(), 100 | unique_hashes.len() as f64 / total_mutation as f64 101 | ); 102 | */ 103 | } 104 | 105 | pub async fn process_feedbacks(&self, inputs: Vec) { 106 | if inputs.is_empty() { 107 | return; 108 | } 109 | 110 | self.get_worker_pool().par_iter().for_each(|worker| { 111 | let worker = worker.clone(); 112 | let mut worker = worker.lock().unwrap(); 113 | worker.process_feedbacks(&inputs); 114 | }); 115 | } 116 | 117 | #[cfg(test)] 118 | fn get_output_sender(&self) -> &AsyncSender> { 119 | self.output_sender.as_ref().unwrap() 120 | } 121 | } 122 | 123 | impl BasicFrontend for BitFlipMutatorFrontend { 124 | type Worker = BitFlipMutatorWorker; 125 | type Output = TestCase; 126 | 127 | // TODO: In the future we should be able to add different mutator. 128 | // We need to make this more flexible. 129 | fn create_worker(&mut self, num: u32) { 130 | self.worker_pool 131 | .extend((0..num).map(|_| Arc::new(Mutex::new(BitFlipMutatorWorker::new())))); 132 | } 133 | 134 | crate::frontend_default!(); 135 | 136 | fn output_transfrom(&self, input: Vec) -> Vec { 137 | self.deep_filter(input) 138 | } 139 | } 140 | 141 | impl AsyncFrontendChannel for BitFlipMutatorFrontend { 142 | crate::async_frontend_default!(); 143 | } 144 | 145 | #[async_trait] 146 | impl AsyncFrontend for BitFlipMutatorFrontend {} 147 | 148 | #[cfg(test)] 149 | mod tests { 150 | use super::*; 151 | 152 | #[tokio::test] 153 | async fn mutator_fronted_deduplicate_mutated_test_cases() { 154 | let mut mutator_frontend = BitFlipMutatorFrontend::new(); 155 | mutator_frontend.create_worker(10); 156 | println!("Go check1234 "); 157 | mutator_frontend.run().await; 158 | println!("Go check123 "); 159 | let sender = mutator_frontend.get_output_sender(); 160 | sender 161 | .send((0..10).map(|i| TestCase::new(vec![i; 1], 0)).collect()) 162 | .await 163 | .unwrap(); 164 | println!("Go check12 "); 165 | sender 166 | .send((0..10).map(|i| TestCase::new(vec![i; 1], 0)).collect()) 167 | .await 168 | .unwrap(); 169 | println!("Go check"); 170 | let result = mutator_frontend.get_results(None).await; 171 | assert_eq!(result.len(), 10); 172 | println!("Go check 33"); 173 | 174 | sender 175 | .send((0..100).map(|_i| TestCase::new(vec![0; 1], 0)).collect()) 176 | .await 177 | .unwrap(); 178 | 179 | let result = mutator_frontend.get_results(Some(100)).await; 180 | assert_eq!(result.len(), 0); 181 | } 182 | } 183 | 184 | pub mod new_mutator_frontend { 185 | use super::DashSet; 186 | use super::TestCase; 187 | use super::DEFAULT_ROUND; 188 | use crate::datatype::{ExecutionStatus, FeedbackData}; 189 | use crate::frontend::load_balance_frontend::*; 190 | use crate::frontend::FuzzerIO; 191 | use crate::mutator::Mutator; 192 | use rayon::prelude::*; 193 | use std::collections::hash_map::DefaultHasher; 194 | use std::collections::HashMap; 195 | use std::hash::{Hash, Hasher}; 196 | 197 | pub struct BitFlipMutatorWorker { 198 | mutator: super::BitFlipMutator, 199 | 200 | inputs_all_bytes: usize, 201 | inputs_num: usize, 202 | acc_counter: usize, 203 | } 204 | 205 | impl Worker for BitFlipMutatorWorker { 206 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 207 | match input { 208 | FuzzerIO::TestCase(data) => { 209 | let mut new_test_cases = Vec::with_capacity(DEFAULT_ROUND * data.len()); 210 | for single_input in data { 211 | self.mutator.mutate(&single_input, DEFAULT_ROUND); 212 | new_test_cases 213 | .extend(self.mutator.get_mutated_test_cases(None).into_iter()); 214 | } 215 | 216 | self.inputs_num += new_test_cases.len(); 217 | for input in new_test_cases.iter() { 218 | self.inputs_all_bytes += input.get_buffer().len(); 219 | } 220 | 221 | self.acc_counter += 1; 222 | if self.acc_counter % 100 == 0 { 223 | println!( 224 | "Average len: {} for {} inputs.", 225 | self.inputs_all_bytes / self.inputs_num, 226 | self.inputs_num 227 | ); 228 | self.inputs_all_bytes = 0; 229 | self.inputs_num = 0; 230 | } 231 | 232 | Some(vec![FuzzerIO::TestCase(new_test_cases)]) 233 | } 234 | 235 | FuzzerIO::MutatorScoreChange(data) => { 236 | self.mutator.update_mutator_func_score(data); 237 | None 238 | } 239 | _ => { 240 | unreachable!(); 241 | } 242 | } 243 | } 244 | } 245 | 246 | impl BitFlipMutatorWorker { 247 | pub fn new() -> Self { 248 | Self { 249 | mutator: super::BitFlipMutator::new(), 250 | 251 | inputs_all_bytes: 0, 252 | inputs_num: 0, 253 | acc_counter: 0, 254 | } 255 | } 256 | } 257 | 258 | impl Default for BitFlipMutatorWorker { 259 | fn default() -> Self { 260 | Self::new() 261 | } 262 | } 263 | 264 | pub struct MutatorTransformer { 265 | unique_hashes: DashSet, 266 | } 267 | 268 | impl MutatorTransformer { 269 | pub fn new() -> Self { 270 | Self { 271 | unique_hashes: DashSet::default(), 272 | } 273 | } 274 | } 275 | 276 | impl Default for MutatorTransformer { 277 | fn default() -> Self { 278 | Self::new() 279 | } 280 | } 281 | 282 | impl Transformer for MutatorTransformer { 283 | fn transform(&self, input: FuzzerIO) -> FuzzerIO { 284 | fn calculate_test_case_hash(t: &TestCase) -> u64 { 285 | let mut s = DefaultHasher::new(); 286 | t.get_buffer().hash(&mut s); 287 | s.finish() 288 | } 289 | 290 | match input { 291 | FuzzerIO::TestCase(test_cases) => { 292 | // Disable filtering 293 | if !test_cases.is_empty() { 294 | return FuzzerIO::TestCase(test_cases); 295 | } 296 | 297 | let unique_hashes = &self.unique_hashes; 298 | FuzzerIO::TestCase( 299 | test_cases 300 | .into_par_iter() 301 | .filter(|test_case| { 302 | let hash_sum = calculate_test_case_hash(test_case); 303 | unique_hashes.insert(hash_sum) 304 | }) 305 | .collect(), 306 | ) 307 | } 308 | _ => { 309 | unreachable!(); 310 | } 311 | } 312 | } 313 | } 314 | 315 | #[derive(Default)] 316 | pub struct MutatorSyncSource { 317 | mutator_scores: HashMap, 318 | } 319 | 320 | impl SyncSource for MutatorSyncSource { 321 | fn update_status(&mut self, input: FuzzerIO) { 322 | const DEFAULT_SCORE: i64 = 1000000000000; 323 | const SCORE_UNINTERESTING: i64 = -1; 324 | const SCORE_INTERESTING: i64 = 1000; 325 | const SCORE_TIMEOUT: i64 = -100000; 326 | const SCORE_CRASH: i64 = 10000; 327 | 328 | if let FuzzerIO::MutationFeedback(feedbacks) = input { 329 | for feedback in feedbacks.iter() { 330 | let mutation_info = feedback.get_mutation_info().unwrap(); 331 | let counter = match feedback.borrow_data() { 332 | Some(FeedbackData::Counter(c)) => *c as i64, 333 | _ => { 334 | unreachable!(); 335 | } 336 | }; 337 | 338 | let score_change = counter 339 | * match feedback.get_status() { 340 | ExecutionStatus::Ok => SCORE_UNINTERESTING, 341 | ExecutionStatus::Timeout => SCORE_TIMEOUT, 342 | ExecutionStatus::Crash => SCORE_CRASH, 343 | ExecutionStatus::Interesting => SCORE_INTERESTING, 344 | }; 345 | 346 | self.mutator_scores 347 | .entry(mutation_info.get_mutation_id()) 348 | .and_modify(|e| *e += score_change) 349 | .or_insert(DEFAULT_SCORE + score_change); 350 | } 351 | } else { 352 | unreachable!(); 353 | } 354 | } 355 | 356 | fn get_status(&mut self) -> FuzzerIO { 357 | FuzzerIO::MutatorScoreChange( 358 | self.mutator_scores 359 | .iter() 360 | .map(|(id, score)| (*id, *score)) 361 | .collect(), 362 | ) 363 | } 364 | } 365 | 366 | pub type BitFlipMutatorFrontend = 367 | Frontend; 368 | 369 | #[cfg(test)] 370 | mod tests { 371 | use super::*; 372 | use crate::frontend::FuzzerIOType; 373 | 374 | #[tokio::test(flavor = "multi_thread", worker_threads = 5)] 375 | #[serial_test::serial] 376 | async fn mutator_load_balance_frontend_test() { 377 | let closure = BitFlipMutatorWorker::new; 378 | let mut frontend = BitFlipMutatorFrontend::new(Box::new(closure)); 379 | let transformer = std::sync::Arc::new(MutatorTransformer::new()); 380 | let worker_num = 10; 381 | frontend.create_worker(worker_num); 382 | frontend.set_transformer(Some(transformer)); 383 | let (input_sender, input_rx) = frontend_channel(100); 384 | let (output_sender, mut output_rx) = frontend_channel(100); 385 | let (mutation_feedback_sender, mutation_feedback_rx) = frontend_channel(100); 386 | frontend.register_input_handler(Receiver::AsyncExclusive(input_rx)); 387 | frontend.register_watch_input_handler( 388 | Receiver::AsyncExclusive(mutation_feedback_rx), 389 | MutatorSyncSource::default(), 390 | ); 391 | frontend.register_output_handler( 392 | FuzzerIOType::TestCase, 393 | Sender::AsyncExclusive(output_sender), 394 | ); 395 | 396 | frontend.run(); 397 | 398 | let mut feedbacks = Vec::default(); 399 | for i in 5..15 { 400 | feedbacks.push( 401 | crate::datatype::Feedback::new(crate::datatype::ExecutionStatus::Interesting) 402 | .set_mutation_info(crate::datatype::MutationInfo::new(i % 2, 1)) 403 | .set_data(crate::datatype::FeedbackData::Counter(1)), 404 | ); 405 | } 406 | 407 | let test_case = TestCase::new(String::from("select 1;").as_bytes().to_vec(), 0); 408 | let test_cases = (0..100) 409 | .map(|_x| test_case.clone()) 410 | .collect::>(); 411 | 412 | for _i in 0..10 { 413 | input_sender 414 | .send(FuzzerIO::TestCase(test_cases.clone())) 415 | .await 416 | .unwrap(); 417 | let output = output_rx.recv().await.unwrap(); 418 | match output.get_type() { 419 | FuzzerIOType::TestCase => { 420 | //println!("Receive output"); 421 | } 422 | _ => { 423 | unreachable!(); 424 | } 425 | } 426 | mutation_feedback_sender 427 | .send(FuzzerIO::MutationFeedback(feedbacks.clone())) 428 | .await 429 | .unwrap(); 430 | } 431 | } 432 | } 433 | } 434 | 435 | pub mod work_stealing_mutator { 436 | use super::DEFAULT_ROUND; 437 | use crate::frontend::load_balance_frontend::SyncSource; 438 | use crate::frontend::work_stealing_frontend::WorkerImpl; 439 | use crate::frontend::{FuzzerIO, FuzzerIOType}; 440 | use crate::mutator::Mutator; 441 | 442 | pub struct MutatorSyncSource { 443 | sync_source: super::new_mutator_frontend::MutatorSyncSource, 444 | counter: usize, 445 | } 446 | 447 | impl MutatorSyncSource { 448 | pub fn new() -> Self { 449 | MutatorSyncSource { 450 | sync_source: super::new_mutator_frontend::MutatorSyncSource::default(), 451 | counter: 0, 452 | } 453 | } 454 | } 455 | 456 | impl Default for MutatorSyncSource { 457 | fn default() -> Self { 458 | Self::new() 459 | } 460 | } 461 | 462 | impl WorkerImpl for MutatorSyncSource { 463 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 464 | assert!(input.get_type() == FuzzerIOType::MutationFeedback); 465 | self.sync_source.update_status(input); 466 | self.counter += 1; 467 | if self.counter == 20 { 468 | self.counter = 0; 469 | Some(vec![self.sync_source.get_status()]) 470 | } else { 471 | None 472 | } 473 | } 474 | } 475 | 476 | pub struct MutatorWorker { 477 | mutator: super::BitFlipMutator, 478 | 479 | inputs_all_bytes: usize, 480 | inputs_num: usize, 481 | acc_counter: usize, 482 | } 483 | 484 | impl WorkerImpl for MutatorWorker { 485 | fn handle_inputs(&mut self, input: FuzzerIO) -> Option> { 486 | match input { 487 | FuzzerIO::TestCase(data) => { 488 | let mut new_test_cases = Vec::with_capacity(DEFAULT_ROUND * data.len()); 489 | for single_input in data { 490 | self.mutator.mutate(&single_input, DEFAULT_ROUND); 491 | new_test_cases 492 | .extend(self.mutator.get_mutated_test_cases(None).into_iter()); 493 | } 494 | 495 | self.inputs_num += new_test_cases.len(); 496 | for input in new_test_cases.iter() { 497 | self.inputs_all_bytes += input.get_buffer().len(); 498 | } 499 | 500 | self.acc_counter += 1; 501 | if self.acc_counter % 100 == 0 { 502 | println!( 503 | "Average len: {} for {} inputs.", 504 | self.inputs_all_bytes / self.inputs_num, 505 | self.inputs_num 506 | ); 507 | self.inputs_all_bytes = 0; 508 | self.inputs_num = 0; 509 | } 510 | 511 | Some(vec![FuzzerIO::TestCase(new_test_cases)]) 512 | } 513 | FuzzerIO::MutatorScoreChange(data) => { 514 | self.mutator.update_mutator_func_score(data); 515 | None 516 | } 517 | _ => { 518 | unreachable!(); 519 | } 520 | } 521 | } 522 | } 523 | 524 | impl MutatorWorker { 525 | pub fn new() -> Self { 526 | Self { 527 | mutator: super::BitFlipMutator::new(), 528 | 529 | inputs_all_bytes: 0, 530 | inputs_num: 0, 531 | acc_counter: 0, 532 | } 533 | } 534 | } 535 | 536 | impl Default for MutatorWorker { 537 | fn default() -> Self { 538 | Self::new() 539 | } 540 | } 541 | } 542 | --------------------------------------------------------------------------------