From e25094eb4f5cbc03e741db80ca2852f9d512da14 Mon Sep 17 00:00:00 2001
From: "Subhojeet Mukherjee, PhD"
 <57270300+CowBoy4mH3LL@users.noreply.github.com>
Date: Tue, 12 Nov 2024 00:06:55 +0530
Subject: [PATCH] Lower capped RAND generators (#2671)

* Lower capped rand generators

* Updated all references to RAND generators

* Formatting updates

* New RAND bytes generator constructor

* Revert "Updated all references to RAND generators"

This reverts commit 9daad894b25ec3867daf93c4fe67c03abec1d8c6.

* Revert "Formatting updates"

This reverts commit ff2a61a366c48b3f313878f62409e51b1e1ed663.

* cargo nightly format

* Added must_use to with_min_size
---
 .../baby_fuzzer_custom_executor/src/main.rs   |  2 +-
 libafl/src/generators/mod.rs                  | 32 ++++++++++++++++---
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/fuzzers/baby/baby_fuzzer_custom_executor/src/main.rs b/fuzzers/baby/baby_fuzzer_custom_executor/src/main.rs
index 04b7993c6f..003839f84d 100644
--- a/fuzzers/baby/baby_fuzzer_custom_executor/src/main.rs
+++ b/fuzzers/baby/baby_fuzzer_custom_executor/src/main.rs
@@ -142,7 +142,7 @@ pub fn main() {
     let mut executor = WithObservers::new(executor, tuple_list!(observer));
 
     // Generator of printable bytearrays of max size 32
-    let mut generator = RandPrintablesGenerator::new(nonzero!(32));
+    let mut generator = RandPrintablesGenerator::with_min_size(nonzero!(1), nonzero!(32));
 
     // Generate 8 initial inputs
     state
diff --git a/libafl/src/generators/mod.rs b/libafl/src/generators/mod.rs
index 13fedb3cec..ce0ce4607f 100644
--- a/libafl/src/generators/mod.rs
+++ b/libafl/src/generators/mod.rs
@@ -74,6 +74,7 @@ where
 #[derive(Clone, Debug)]
 /// Generates random bytes
 pub struct RandBytesGenerator {
+    min_size: NonZeroUsize,
     max_size: NonZeroUsize,
 }
 
@@ -82,7 +83,9 @@ where
     S: HasRand,
 {
     fn generate(&mut self, state: &mut S) -> Result<BytesInput, Error> {
-        let mut size = state.rand_mut().below(self.max_size);
+        let mut size = state
+            .rand_mut()
+            .between(self.min_size.get(), self.max_size.get());
         size = max(size, 1);
         let random_bytes: Vec<u8> = (0..size)
             .map(|_| state.rand_mut().below(nonzero!(256)) as u8)
@@ -95,13 +98,23 @@ impl RandBytesGenerator {
     /// Returns a new [`RandBytesGenerator`], generating up to `max_size` random bytes.
     #[must_use]
     pub fn new(max_size: NonZeroUsize) -> Self {
-        Self { max_size }
+        Self {
+            min_size: nonzero!(1),
+            max_size,
+        }
+    }
+
+    /// Returns a new [`RandBytesGenerator`], generating from `min_size` up to `max_size` random bytes.
+    #[must_use]
+    pub fn with_min_size(min_size: NonZeroUsize, max_size: NonZeroUsize) -> Self {
+        Self { min_size, max_size }
     }
 }
 
 #[derive(Clone, Debug)]
 /// Generates random printable characters
 pub struct RandPrintablesGenerator {
+    min_size: NonZeroUsize,
     max_size: NonZeroUsize,
 }
 
@@ -110,7 +123,9 @@ where
     S: HasRand,
 {
     fn generate(&mut self, state: &mut S) -> Result<BytesInput, Error> {
-        let mut size = state.rand_mut().below(self.max_size);
+        let mut size = state
+            .rand_mut()
+            .between(self.min_size.get(), self.max_size.get());
         size = max(size, 1);
         let printables = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz \t\n!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~".as_bytes();
         let random_bytes: Vec<u8> = (0..size)
@@ -124,6 +139,15 @@ impl RandPrintablesGenerator {
     /// Returns a new [`RandBytesGenerator`], generating up to `max_size` random bytes.
     #[must_use]
     pub fn new(max_size: NonZeroUsize) -> Self {
-        Self { max_size }
+        Self {
+            min_size: nonzero!(1),
+            max_size,
+        }
+    }
+
+    /// Returns a new [`RandPrintablesGenerator`], generating from `min_size` up to `max_size` random bytes.
+    #[must_use]
+    pub fn with_min_size(min_size: NonZeroUsize, max_size: NonZeroUsize) -> Self {
+        Self { min_size, max_size }
     }
 }