From 0374e319079a10b8cb1d7e447fceece3f80ac5d3 Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Tue, 6 Aug 2019 11:05:35 +1000
Subject: [PATCH] Disallow extra bytes in Bitfield from_bytes (#488)

---
 eth2/utils/ssz_types/src/bitfield.rs | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/eth2/utils/ssz_types/src/bitfield.rs b/eth2/utils/ssz_types/src/bitfield.rs
index 592574e15..a95a5faa4 100644
--- a/eth2/utils/ssz_types/src/bitfield.rs
+++ b/eth2/utils/ssz_types/src/bitfield.rs
@@ -161,6 +161,7 @@ impl<N: Unsigned + Clone> Bitfield<Variable<N>> {
     ///
     /// Returns `None` if `bytes` are not a valid encoding.
     pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, Error> {
+        let bytes_len = bytes.len();
         let mut initial_bitfield: Bitfield<Variable<N>> = {
             let num_bits = bytes.len() * 8;
             Bitfield::from_raw_bytes(bytes, num_bits)?
@@ -170,6 +171,14 @@ impl<N: Unsigned + Clone> Bitfield<Variable<N>> {
             .highest_set_bit()
             .ok_or_else(|| Error::MissingLengthInformation)?;
 
+        // The length bit should be in the last byte, or else it means we have too many bytes.
+        if len / 8 + 1 != bytes_len {
+            return Err(Error::InvalidByteCount {
+                given: bytes_len,
+                expected: len / 8 + 1,
+            });
+        }
+
         if len <= Self::max_len() {
             initial_bitfield
                 .set(len, false)
@@ -825,6 +834,17 @@ mod bitlist {
         assert!(BitList8::from_ssz_bytes(&[0b0000_0001, 0b0000_0100]).is_err());
     }
 
+    #[test]
+    fn ssz_decode_extra_bytes() {
+        assert!(BitList0::from_ssz_bytes(&[0b0000_0001, 0b0000_0000]).is_err());
+        assert!(BitList1::from_ssz_bytes(&[0b0000_0001, 0b0000_0000]).is_err());
+        assert!(BitList8::from_ssz_bytes(&[0b0000_0001, 0b0000_0000]).is_err());
+        assert!(BitList16::from_ssz_bytes(&[0b0000_0001, 0b0000_0000]).is_err());
+        assert!(BitList1024::from_ssz_bytes(&[0b1000_0000, 0]).is_err());
+        assert!(BitList1024::from_ssz_bytes(&[0b1000_0000, 0, 0]).is_err());
+        assert!(BitList1024::from_ssz_bytes(&[0b1000_0000, 0, 0, 0, 0]).is_err());
+    }
+
     #[test]
     fn ssz_round_trip() {
         assert_round_trip(BitList0::with_capacity(0).unwrap());