Repository: ParkMyCar/compact_str
Pull request #161 (2022-07-24 21:09:37)

fix: Test case discovered by AFL
When we changed `CompactString` to have the same size as `Option<CompactString>`, we made the `HEAP_MASK` discriminant `254`, which made our `CAPACITY_IS_ON_THE_HEAP` flag `[254, 254, 254, 254]`.

When you encode `16711422` in little endian bytes, it's `[254, 254, 254, 0]`, and when we append our `HEAP_MASK` we get `[254, 254, 254, 254]`. So when we'd try to create a `CompactString` with size `16711422`, we'd incorrectly identify the capacity as being on the heap, when it's really inline.

To fix this I updated the `CAPACITY_IS_ON_THE_HEAP` `[255, 255, 255, ..., HEAP_MASK]`, which is the correct value. I added a test to assert all values for `Capacity` successfully round trip for 32-bit architectures.


Patch:
diff --git a/compact_str/src/repr/boxed/capacity.rs b/compact_str/src/repr/boxed/capacity.rs
index 1dc00605..1d4f60cd 100644
--- a/compact_str/src/repr/boxed/capacity.rs
+++ b/compact_str/src/repr/boxed/capacity.rs
@@ -2,8 +2,21 @@ use crate::repr::HEAP_MASK;
 
 // how many bytes a `usize` occupies
 const USIZE_SIZE: usize = core::mem::size_of::<usize>();
-// state that describes the capacity as being stored on the heap
-const CAPACITY_IS_ON_THE_HEAP: [u8; USIZE_SIZE] = [HEAP_MASK; USIZE_SIZE];
+
+/// Used to generate [`CAPACITY_IS_ON_THE_HEAP`]
+#[allow(non_snake_case)]
+const fn CAP_ON_HEAP_FLAG() -> [u8; USIZE_SIZE] {
+    // all bytes 255, with the last
+    let mut flag = [255; USIZE_SIZE];
+    flag[USIZE_SIZE - 1] = HEAP_MASK;
+    flag
+}
+
+/// State that describes the capacity as being stored on the heap.
+///
+/// All bytes `255`, with the last being [`HEAP_MASK`], using the same amount of bytes as `usize`
+/// Example (64-bit): `[255, 255, 255, 255, 255, 255, 255, 254]`
+const CAPACITY_IS_ON_THE_HEAP: [u8; USIZE_SIZE] = CAP_ON_HEAP_FLAG();
 
 // how many bytes we can use for capacity
 const SPACE_FOR_CAPACITY: usize = USIZE_SIZE - 1;
@@ -101,6 +114,8 @@ impl Capacity {
 
 #[cfg(test)]
 mod tests {
+    use rayon::prelude::*;
+
     use super::Capacity;
 
     #[test]
@@ -146,6 +161,27 @@ mod tests {
         let og = usize::MAX;
         assert!(Capacity::new(og).is_err());
     }
+
+    #[test]
+    #[cfg_attr(miri, ignore)]
+    fn test_all_valid_32bit_values() {
+        #[cfg(target_pointer_width = "32")]
+        assert_eq!(16_777_214, super::MAX_VALUE);
+
+        (0..=16_777_214)
+            .into_par_iter()
+            .for_each(|i| match Capacity::new(i) {
+                Ok(cap) => match cap.as_usize() {
+                    Ok(val) => assert_eq!(val, i, "value roundtriped to wrong value?"),
+                    Err(_) => panic!("value converted, but failed to roundtrip! val: {}", i),
+                },
+                Err(_) => panic!("failed to convert {}", i),
+            });
+
+        // one above the 32-bit max value
+        #[cfg(target_pointer_width = "32")]
+        assert!(Capacity::new(16_777_215).is_err());
+    }
 }
 
 crate::asserts::assert_size_eq!(Capacity, usize);
diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml
index 3d64d3a1..b188d6e4 100644
--- a/fuzz/Cargo.toml
+++ b/fuzz/Cargo.toml
@@ -32,6 +32,12 @@ test = false
 doc = false
 required-features = ["afl"]
 
+[[bin]]
+name = "debug"
+path = "src/bin/debug.rs"
+test = false
+doc = false
+
 [[bin]]
 name = "honggfuzz"
 path = "src/bin/honggfuzz.rs"
diff --git a/fuzz/src/bin/debug.rs b/fuzz/src/bin/debug.rs
new file mode 100644
index 00000000..4364fd33
--- /dev/null
+++ b/fuzz/src/bin/debug.rs
@@ -0,0 +1,26 @@
+//! Binary that allows you to pass an input file, it creates a [`compact_str_fuzz::Scenario`] and
+//! then runs it.
+//!
+//! This is helpful when AFL finds failures and we need to reproduce them.
+
+use std::path::PathBuf;
+
+use arbitrary::{
+    Arbitrary,
+    Unstructured,
+};
+use compact_str_fuzz::Scenario;
+
+pub fn main() {
+    let path = std::env::args()
+        .nth(1)
+        .map(PathBuf::from)
+        .expect("no path provided!");
+    let data = std::fs::read(path).expect("failed to read input file");
+    let mut unstructured = Unstructured::new(&data);
+    let scenario = Scenario::arbitrary(&mut unstructured).expect("failed to create Scenario");
+
+    println!("Scenario: {:?}", scenario);
+
+    scenario.run();
+}


Test patch:
diff --git a/compact_str/src/tests.rs b/compact_str/src/tests.rs
index 3e1af2fc..2ccfa3a1 100644
--- a/compact_str/src/tests.rs
+++ b/compact_str/src/tests.rs
@@ -1024,3 +1024,35 @@ fn test_insert() {
         "\u{ffff}\u{ffff}\u{ffff}\u{ffff}\u{ffff}\u{ffff}\u{ffff}\u{ffff}",
     );
 }
+
+#[test]
+fn test_with_capacity_16711422() {
+    // Fuzzing with AFL on a 32-bit ARM arch found this bug!
+    //
+    // We have our own heap implemenation called BoxString, which optionally stores the capacity
+    // on the heap, which is really only relevant for 32-bit architectures. The discriminant it used
+    // to determine if capacity was on the heap, was when the last `usize` number of bytes were all
+    // equal to our internal HEAP_MASK, which at the time was `255`. At the time this worked and was
+    // correct.
+    //
+    // When we released support to make the size of CompactString == Option<CompactString>, we
+    // changed the HEAP_MASK to `254`, which unintentionally made our discriminant for determining
+    // if our capacity was on the heap, all `254`s, yet our "max inline capacity value" was still
+    // based on the discriminant being all `255`s.
+    //
+    // When creating a BoxString with capacity 16711422, we'd correctly decide we could store the
+    // capacity inline, but this would create a capacity with an underlying value of
+    // [254, 254, 254, HEAP_MASK]. Once the HEAP_MASK changed to 254, this capacity was now the same
+    // as the discriminant to determine if the capacity was on the heap, so we'd incorrectly
+    // identify the capacity as being on the heap, when it was really inline.
+
+    assert_eq!(16711422_u32.to_le_bytes(), [254, 254, 254, 0]);
+
+    let compact = CompactString::with_capacity(16711422);
+    let std_str = String::with_capacity(16711422);
+
+    assert!(compact.is_heap_allocated());
+    assert_eq!(compact.capacity(), std_str.capacity());
+    assert_eq!(compact, "");
+    assert_eq!(compact, std_str);
+}