fix(bitfield): prevent construction of out-of-range bitfields

Stebalien · Stebalien · commit 6f7fce7d93a0 · 2022-03-29T16:24:45.000-04:00
1. Produce an error when collecting a bitfield with an invalid bit, or
   attempting to set an invalid bit.
2. Make sensitive functions crate-private.
diff --git a/ipld/bitfield/src/iter/mod.rs b/ipld/bitfield/src/iter/mod.rs
@@ -133,6 +133,9 @@ where
     I: Iterator<Item = Range<u64>>,
 {
     /// Creates a new `Ranges` instance.
+    ///
+    /// WARNING: This is asserting that the underlying iterator obeys the `RangeIterator`
+    /// constraints. Using this incorrectly could lead to panics, etc.
     pub fn new<II>(iter: II) -> Self
     where
         II: IntoIterator<IntoIter = I, Item = Range<u64>>,
@@ -155,16 +158,21 @@ where
 impl<I> RangeIterator for Ranges<I> where I: Iterator<Item = Range<u64>> {}
 
 /// Returns a `RangeIterator` which ranges contain the values from the provided iterator.
-/// The values need to be in ascending order — if not, the returned iterator may not satisfy
-/// all `RangeIterator` requirements.
-pub fn ranges_from_bits(bits: impl IntoIterator<Item = u64>) -> impl RangeIterator {
+/// The values need to be in ascending order and may not include u64::MAX. Otherwise, the iterator
+/// will panic.
+pub(crate) fn ranges_from_bits(bits: impl IntoIterator<Item = u64>) -> impl RangeIterator {
     let mut iter = bits.into_iter().peekable();
 
     Ranges::new(iter::from_fn(move || {
         let start = iter.next()?;
-        let mut end = start + 1;
-        while iter.peek() == Some(&end) {
-            end += 1;
+        let mut end = start.checked_add(1).expect("bitfield overflow");
+        while let Some(&next) = iter.peek() {
+            if next < end {
+                panic!("out of order bitfield")
+            } else if next > end {
+                break;
+            }
+            end = end.checked_add(1).expect("bitfield overflow");
             iter.next();
         }
         Some(start..end)
diff --git a/ipld/bitfield/src/lib.rs b/ipld/bitfield/src/lib.rs
@@ -1,22 +1,36 @@
 // Copyright 2019-2022 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
+// disable this lint because it can actually cause performance regressions, and usually leads to
+// hard to read code.
+#![allow(clippy::comparison_chain)]
+
 pub mod iter;
 mod range;
 mod rleplus;
 mod unvalidated;
 
 use std::collections::BTreeSet;
-use std::iter::FromIterator;
 use std::ops::{
     BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Range, Sub, SubAssign,
 };
 
 use iter::{ranges_from_bits, RangeIterator};
 pub(crate) use range::RangeSize;
 pub use rleplus::Error;
+use thiserror::Error;
 pub use unvalidated::{UnvalidatedBitField, Validate};
 
+#[derive(Clone, Error, Debug)]
+#[error("bitfields may not include u64::MAX")]
+pub struct OutOfRangeError;
+
+impl From<OutOfRangeError> for Error {
+    fn from(_: OutOfRangeError) -> Self {
+        Error::RLEOverflow
+    }
+}
+
 /// A bit field with buffered insertion/removal that serializes to/from RLE+. Similar to
 /// `HashSet<u64>`, but more memory-efficient when long runs of 1s and 0s are present.
 #[derive(Debug, Default, Clone)]
@@ -35,22 +49,79 @@ impl PartialEq for BitField {
     }
 }
 
-impl FromIterator<u64> for BitField {
-    fn from_iter<I: IntoIterator<Item = u64>>(iter: I) -> Self {
-        let mut vec: Vec<_> = iter.into_iter().collect();
-        if vec.is_empty() {
-            Self::new()
+/// Possibly a valid bitfield, or an out of bounds error. Ideally we'd just use a result, but we
+/// can't implement [`FromIterator`] on a [`Result`] due to coherence.
+///
+/// You probably want to call [`BitField::try_from_bits`] instead of using this directly.
+#[doc(hidden)]
+pub enum MaybeBitField {
+    /// A valid bitfield.
+    Ok(BitField),
+    /// Out of bounds.
+    OutOfBounds,
+}
+
+impl MaybeBitField {
+    pub fn unwrap(self) -> BitField {
+        use MaybeBitField::*;
+        match self {
+            Ok(bf) => bf,
+            OutOfBounds => panic!("bitfield bit out of bounds"),
+        }
+    }
+
+    pub fn expect(self, message: &str) -> BitField {
+        use MaybeBitField::*;
+        match self {
+            Ok(bf) => bf,
+            OutOfBounds => panic!("{}", message),
+        }
+    }
+}
+
+impl TryFrom<MaybeBitField> for BitField {
+    type Error = OutOfRangeError;
+
+    fn try_from(value: MaybeBitField) -> Result<Self, Self::Error> {
+        match value {
+            MaybeBitField::Ok(bf) => Ok(bf),
+            MaybeBitField::OutOfBounds => Err(OutOfRangeError),
+        }
+    }
+}
+
+impl FromIterator<bool> for MaybeBitField {
+    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> MaybeBitField {
+        let mut iter = iter.into_iter().fuse();
+        let bits = (0u64..u64::MAX)
+            .zip(&mut iter)
+            .filter(|&(_, b)| b)
+            .map(|(i, _)| i);
+        let bf = BitField::from_ranges(ranges_from_bits(bits));
+
+        // Now, if we have remaining bits, raise an error. Otherwise, we're good.
+        if iter.next().is_some() {
+            MaybeBitField::OutOfBounds
         } else {
-            vec.sort_unstable();
-            Self::from_ranges(ranges_from_bits(vec))
+            MaybeBitField::Ok(bf)
         }
     }
 }
 
-impl FromIterator<bool> for BitField {
-    fn from_iter<I: IntoIterator<Item = bool>>(iter: I) -> Self {
-        let bits = (0u64..).zip(iter).filter(|&(_, b)| b).map(|(i, _)| i);
-        Self::from_ranges(ranges_from_bits(bits))
+impl FromIterator<u64> for MaybeBitField {
+    fn from_iter<T: IntoIterator<Item = u64>>(iter: T) -> MaybeBitField {
+        let mut vec: Vec<_> = iter.into_iter().collect();
+        if vec.is_empty() {
+            MaybeBitField::Ok(BitField::new())
+        } else {
+            vec.sort_unstable();
+            vec.dedup();
+            if vec.last() == Some(&u64::MAX) {
+                MaybeBitField::OutOfBounds
+            } else {
+                MaybeBitField::Ok(BitField::from_ranges(ranges_from_bits(vec)))
+            }
+        }
     }
 }
 
@@ -68,10 +139,33 @@ impl BitField {
         }
     }
 
-    /// Adds the bit at a given index to the bit field.
+    /// Tries to create a new bitfield from a bit iterator. It fails if the resulting bitfield would
+    /// contain values not in the range `0..u64::MAX` (non-inclusive).
+    pub fn try_from_bits<I>(iter: I) -> Result<Self, OutOfRangeError>
+    where
+        I: IntoIterator,
+        MaybeBitField: FromIterator<I::Item>,
+    {
+        iter.into_iter().collect::<MaybeBitField>().try_into()
+    }
+
+    /// Adds the bit at a given index to the bit field, panicing if it's out of range.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `bit` is `u64::MAX`.
     pub fn set(&mut self, bit: u64) {
+        self.try_set(bit).unwrap()
+    }
+
+    /// Adds the bit at a given index to the bit field, returning an error if it's out of range.
+    pub fn try_set(&mut self, bit: u64) -> Result<(), OutOfRangeError> {
+        if bit == u64::MAX {
+            return Err(OutOfRangeError);
+        }
         self.unset.remove(&bit);
         self.set.insert(bit);
+        Ok(())
     }
 
     /// Removes the bit at a given index from the bit field.
@@ -328,7 +422,7 @@ macro_rules! bitfield {
         std::iter::once($head != 0_u32).chain(bitfield!(@iter $($tail),*))
     };
     ($($val:literal),* $(,)?) => {
-        bitfield!(@iter $($val),*).collect::<$crate::BitField>()
+        bitfield!(@iter $($val),*).collect::<$crate::MaybeBitField>().unwrap()
     };
 }
 
diff --git a/ipld/bitfield/tests/bitfield_tests.rs b/ipld/bitfield/tests/bitfield_tests.rs
@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0, MIT
 
 use std::collections::HashSet;
-use std::iter::FromIterator;
 
 use fvm_ipld_bitfield::{bitfield, BitField};
 use fvm_shared::encoding;
@@ -17,7 +16,7 @@ fn random_indices(range: u64, seed: u64) -> Vec<u64> {
 #[test]
 fn bitfield_slice() {
     let vals = random_indices(10000, 2);
-    let bf: BitField = vals.iter().copied().collect();
+    let bf = BitField::try_from_bits(vals.iter().copied()).unwrap();
 
     let slice = bf.slice(600, 500).unwrap();
     let out_vals: Vec<_> = slice.iter().collect();
@@ -38,7 +37,7 @@ fn bitfield_slice_small() {
     let vals = [1, 5, 6, 7, 10, 11, 12, 15];
 
     let test_permutations = |start: usize, count: usize| {
-        let bf: BitField = vals.iter().copied().collect();
+        let bf = BitField::try_from_bits(vals.iter().copied()).unwrap();
         let sl = bf.slice(start as u64, count as u64).unwrap();
         let exp = &vals[start..start + count];
         let out: Vec<_> = sl.iter().collect();
@@ -56,8 +55,8 @@ fn set_up_test_bitfields() -> (Vec<u64>, Vec<u64>, BitField, BitField) {
     let a = random_indices(100, 1);
     let b = random_indices(100, 2);
 
-    let bf_a: BitField = a.iter().copied().collect();
-    let bf_b: BitField = b.iter().copied().collect();
+    let bf_a = BitField::try_from_bits(a.iter().copied()).unwrap();
+    let bf_b = BitField::try_from_bits(b.iter().copied()).unwrap();
 
     (a, b, bf_a, bf_b)
 }
@@ -101,50 +100,62 @@ fn bitfield_difference() {
 // Ported test from go impl (specs-actors)
 #[test]
 fn subtract_more() {
-    let have = BitField::from_iter(vec![5, 6, 8, 10, 11, 13, 14, 17]);
-    let s1 = &BitField::from_iter(vec![5, 6]) - &have;
-    let s2 = &BitField::from_iter(vec![8, 10]) - &have;
-    let s3 = &BitField::from_iter(vec![11, 13]) - &have;
-    let s4 = &BitField::from_iter(vec![14, 17]) - &have;
+    let have = BitField::try_from_bits(vec![5, 6, 8, 10, 11, 13, 14, 17]).unwrap();
+    let s1 = &BitField::try_from_bits(vec![5, 6]).unwrap() - &have;
+    let s2 = &BitField::try_from_bits(vec![8, 10]).unwrap() - &have;
+    let s3 = &BitField::try_from_bits(vec![11, 13]).unwrap() - &have;
+    let s4 = &BitField::try_from_bits(vec![14, 17]).unwrap() - &have;
 
     let u = BitField::union(&[s1, s2, s3, s4]);
     assert_eq!(u.len(), 0);
 }
 
 #[test]
 fn contains_any() {
-    assert!(!BitField::from_iter(vec![0, 4]).contains_any(&BitField::from_iter(vec![1, 3, 5])));
+    assert!(!BitField::try_from_bits(vec![0, 4])
+        .unwrap()
+        .contains_any(&BitField::try_from_bits(vec![1, 3, 5]).unwrap()));
 
-    assert!(BitField::from_iter(vec![0, 2, 5, 6]).contains_any(&BitField::from_iter(vec![1, 3, 5])));
+    assert!(BitField::try_from_bits(vec![0, 2, 5, 6])
+        .unwrap()
+        .contains_any(&BitField::try_from_bits(vec![1, 3, 5]).unwrap()));
 
-    assert!(BitField::from_iter(vec![1, 2, 3]).contains_any(&BitField::from_iter(vec![1, 2, 3])));
+    assert!(BitField::try_from_bits(vec![1, 2, 3])
+        .unwrap()
+        .contains_any(&BitField::try_from_bits(vec![1, 2, 3]).unwrap()));
 }
 
 #[test]
 fn contains_all() {
-    assert!(
-        !BitField::from_iter(vec![0, 2, 4]).contains_all(&BitField::from_iter(vec![0, 2, 4, 5]))
-    );
+    assert!(!BitField::try_from_bits(vec![0, 2, 4])
+        .unwrap()
+        .contains_all(&BitField::try_from_bits(vec![0, 2, 4, 5]).unwrap()));
 
-    assert!(BitField::from_iter(vec![0, 2, 4, 5]).contains_all(&BitField::from_iter(vec![0, 2, 4])));
+    assert!(BitField::try_from_bits(vec![0, 2, 4, 5])
+        .unwrap()
+        .contains_all(&BitField::try_from_bits(vec![0, 2, 4]).unwrap()));
 
-    assert!(BitField::from_iter(vec![1, 2, 3]).contains_all(&BitField::from_iter(vec![1, 2, 3])));
+    assert!(BitField::try_from_bits(vec![1, 2, 3])
+        .unwrap()
+        .contains_all(&BitField::try_from_bits(vec![1, 2, 3]).unwrap()));
 }
 
 #[test]
 fn bit_ops() {
-    let a = &BitField::from_iter(vec![1, 2, 3]) & &BitField::from_iter(vec![1, 3, 4]);
+    let a = &BitField::try_from_bits(vec![1, 2, 3]).unwrap()
+        & &BitField::try_from_bits(vec![1, 3, 4]).unwrap();
     assert_eq!(a.iter().collect::<Vec<_>>(), &[1, 3]);
 
-    let mut a = BitField::from_iter(vec![1, 2, 3]);
-    a &= &BitField::from_iter(vec![1, 3, 4]);
+    let mut a = BitField::try_from_bits(vec![1, 2, 3]).unwrap();
+    a &= &BitField::try_from_bits(vec![1, 3, 4]).unwrap();
     assert_eq!(a.iter().collect::<Vec<_>>(), &[1, 3]);
 
-    let a = &BitField::from_iter(vec![1, 2, 3]) | &BitField::from_iter(vec![1, 3, 4]);
+    let a = &BitField::try_from_bits(vec![1, 2, 3]).unwrap()
+        | &BitField::try_from_bits(vec![1, 3, 4]).unwrap();
     assert_eq!(a.iter().collect::<Vec<_>>(), &[1, 2, 3, 4]);
 
-    let mut a = BitField::from_iter(vec![1, 2, 3]);
-    a |= &BitField::from_iter(vec![1, 3, 4]);
+    let mut a = BitField::try_from_bits(vec![1, 2, 3]).unwrap();
+    a |= &BitField::try_from_bits(vec![1, 3, 4]).unwrap();
     assert_eq!(a.iter().collect::<Vec<_>>(), &[1, 2, 3, 4]);
 }
 
@@ -209,3 +220,16 @@ fn padding() {
     let deserialized: BitField = encoding::from_slice(&cbor).unwrap();
     assert_eq!(deserialized, bf);
 }
+
+#[test]
+fn exceeds_bitfield_range() {
+    let mut bf = BitField::new();
+    bf.try_set(u64::MAX)
+        .expect_err("expected setting u64::MAX to fail");
+    bf.try_set(u64::MAX - 1)
+        .expect("expected setting u64::MAX-1 to succeed");
+    BitField::try_from_bits([0, 1, 4, 99, u64::MAX])
+        .expect_err("expected setting u64::MAX to fail");
+    BitField::try_from_bits([0, 1, 4, 99, u64::MAX - 1])
+        .expect("expected setting u64::MAX-1 to succeed");
+}
