From 90d557d9164d5fc44ae8d52a4669b3a7112f4a81 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Mon, 30 Nov 2020 12:18:10 -0800
Subject: [PATCH] Strengthen EpochSlots sanitization

---
 core/src/epoch_slots.rs | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/core/src/epoch_slots.rs b/core/src/epoch_slots.rs
index ed325d34c5..fdcf49f1dc 100644
--- a/core/src/epoch_slots.rs
+++ b/core/src/epoch_slots.rs
@@ -24,6 +24,15 @@ impl Sanitize for Uncompressed {
         if self.num >= MAX_SLOTS_PER_ENTRY {
             return Err(SanitizeError::ValueOutOfBounds);
         }
+        if self.slots.len() % 8 != 0 {
+            // Uncompressed::new() ensures the length is always a multiple of 8
+            return Err(SanitizeError::ValueOutOfBounds);
+        }
+        if self.slots.len() != self.slots.capacity() {
+            // A BitVec<u8> with a length that's a multiple of 8 will always have len() equal to
+            // capacity(), assuming no bit manipulation
+            return Err(SanitizeError::ValueOutOfBounds);
+        }
         Ok(())
     }
 }
@@ -132,7 +141,7 @@ impl Uncompressed {
             if *s < self.first_slot {
                 return i;
             }
-            if *s - self.first_slot >= self.slots.capacity() {
+            if *s - self.first_slot >= self.slots.len() {
                 return i;
             }
             self.slots.set(*s - self.first_slot, true);
@@ -393,6 +402,14 @@ mod tests {
         o.num = MAX_SLOTS_PER_ENTRY;
         assert_eq!(o.sanitize(), Err(SanitizeError::ValueOutOfBounds));
 
+        let mut o = slots.clone();
+        o.slots = BitVec::new_fill(false, 7); // Length not a multiple of 8
+        assert_eq!(o.sanitize(), Err(SanitizeError::ValueOutOfBounds));
+
+        let mut o = slots.clone();
+        o.slots = BitVec::with_capacity(8); // capacity() not equal to len()
+        assert_eq!(o.sanitize(), Err(SanitizeError::ValueOutOfBounds));
+
         let compressed = Flate2::deflate(slots).unwrap();
         assert!(compressed.sanitize().is_ok());