From b2791c191a941208eff825a44440fbf605dcfb60 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Mon, 12 Feb 2024 00:38:10 +0000
Subject: [PATCH] Add a fuzzer of the proof building state machine

---
 fuzz/Cargo.toml          |  4 +++
 fuzz/src/fuzz_builder.rs | 68 ++++++++++++++++++++++++++++++++++++++++
 src/query.rs             | 19 +++++++++--
 3 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 fuzz/src/fuzz_builder.rs

diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml
index 56ed89c..b8869bd 100644
--- a/fuzz/Cargo.toml
+++ b/fuzz/Cargo.toml
@@ -48,3 +48,7 @@ path = "src/parse_response.rs"
 [[bin]]
 name = "parse_stream_validate"
 path = "src/parse_stream_validate.rs"
+
+[[bin]]
+name = "fuzz_builder"
+path = "src/fuzz_builder.rs"
diff --git a/fuzz/src/fuzz_builder.rs b/fuzz/src/fuzz_builder.rs
new file mode 100644
index 0000000..f4dc207
--- /dev/null
+++ b/fuzz/src/fuzz_builder.rs
@@ -0,0 +1,68 @@
+// This file is Copyright its original authors, visible in version control
+// history.
+//
+// This file is licensed under the Apache License, Version 2.0 <LICENSE-APACHE
+// or http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your option.
+// You may not use this file except in accordance with one or both of these
+// licenses.
+
+#![cfg_attr(feature = "libfuzzer_fuzz", no_main)]
+
+#[cfg(not(fuzzing))]
+compile_error!("Fuzz targets need cfg=fuzzing");
+
+extern crate dnssec_prover;
+use dnssec_prover::query::fuzz_proof_builder;
+
+#[cfg(feature = "afl")]
+#[macro_use] extern crate afl;
+#[cfg(feature = "afl")]
+fn main() {
+	fuzz!(|data| {
+		fuzz_proof_builder(data);
+	});
+}
+
+#[cfg(feature = "honggfuzz")]
+#[macro_use] extern crate honggfuzz;
+#[cfg(feature = "honggfuzz")]
+fn main() {
+	loop {
+		fuzz!(|data| {
+			fuzz_proof_builder(data);
+		});
+	}
+}
+
+#[cfg(feature = "libfuzzer_fuzz")]
+#[macro_use] extern crate libfuzzer_sys;
+#[cfg(feature = "libfuzzer_fuzz")]
+fuzz_target!(|data: &[u8]| {
+	fuzz_proof_builder(data);
+});
+
+#[cfg(feature = "stdin_fuzz")]
+fn main() {
+	use std::io::Read;
+
+	let mut data = Vec::with_capacity(8192);
+	std::io::stdin().read_to_end(&mut data).unwrap();
+	fuzz_proof_builder(&data);
+}
+
+#[test]
+fn run_test_cases() {
+	use std::fs;
+	use std::io::Read;
+
+	if let Ok(tests) = fs::read_dir("test_cases/parse_response") {
+		for test in tests {
+			let mut data: Vec<u8> = Vec::new();
+			let path = test.unwrap().path();
+			fs::File::open(&path).unwrap().read_to_end(&mut data).unwrap();
+
+			fuzz_proof_builder(&data);
+		}
+	}
+}
diff --git a/src/query.rs b/src/query.rs
index c2ea75e..37fbaef 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -21,9 +21,9 @@ use crate::ser::*;
 // In testing use a rather small buffer to ensure we hit the allocation paths sometimes. In
 // production, we should generally never actually need to go to heap as DNS messages are rarely
 // larger than a KiB or two.
-#[cfg(test)]
+#[cfg(any(test, fuzzing))]
 const STACK_BUF_LIMIT: u16 = 32;
-#[cfg(not(test))]
+#[cfg(not(any(test, fuzzing)))]
 const STACK_BUF_LIMIT: u16 = 2048;
 
 /// A buffer for storing queries and responses.
@@ -147,6 +147,21 @@ fn handle_response(resp: &[u8], proof: &mut Vec<u8>, rrsig_key_names: &mut Vec<N
 	Ok(min_ttl)
 }
 
+#[cfg(fuzzing)]
+/// Read a stream of responses and handle them it as if they came from a server, for fuzzing.
+pub fn fuzz_proof_builder(mut response_stream: &[u8]) {
+	let (mut builder, _) = ProofBuilder::new(&"example.com.".try_into().unwrap(), Txt::TYPE);
+	while builder.awaiting_responses() {
+		let len = if let Ok(len) = read_u16(&mut response_stream) { len } else { return };
+		let mut buf = QueryBuf::new_zeroed(len);
+		if response_stream.len() < len as usize { return; }
+		buf.copy_from_slice(&response_stream[..len as usize]);
+		response_stream = &response_stream[len as usize..];
+		let _ = builder.process_response(&buf);
+	}
+	let _ = builder.finish_proof();
+}
+
 const MAX_REQUESTS: usize = 10;
 /// A simple state machine which will generate a series of queries and process the responses until
 /// it has built a DNSSEC proof.
-- 
2.39.5