Fix AllowedIps iteration with empty leaves (cloudflare#156)

Avoid right shift overflow when a leaf doesn't contain any bits.
sharjeel · Jan 15, 2021 · d513497 · d513497
1 parent e848dd1
commit d513497
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/src/device/allowed_ips.rs b/src/device/allowed_ips.rs
@@ -451,6 +451,9 @@ macro_rules! build_iter {
                         }
                         Some($node::Leaf(key, bits, data)) => {
                             let (cur_key, cur_bits) = self.key_hlp.pop().unwrap();
+                            if cur_bits == BITS && *bits == 0 {
+                                return Some((data, cur_key, cur_bits));
+                            }
                             return Some((data, cur_key ^ (key >> cur_bits), cur_bits + bits));
                         }
                         Some($node::Node {
@@ -776,4 +779,27 @@ mod tests {
             ]))
         );
     }
+
+    #[test]
+    fn test_allowed_ips_iter_zero_leaf_bits() {
+        let mut map: AllowedIps<char> = Default::default();
+        map.insert(IpAddr::from([10, 111, 0, 1]), 32, '1');
+        map.insert(IpAddr::from([10, 111, 0, 2]), 32, '2');
+        map.insert(IpAddr::from([10, 111, 0, 3]), 32, '3');
+
+        let mut map_iter = map.iter();
+        assert_eq!(
+            map_iter.next(),
+            Some((&'1', IpAddr::from([10, 111, 0, 1]), 32))
+        );
+        assert_eq!(
+            map_iter.next(),
+            Some((&'2', IpAddr::from([10, 111, 0, 2]), 32))
+        );
+        assert_eq!(
+            map_iter.next(),
+            Some((&'3', IpAddr::from([10, 111, 0, 3]), 32))
+        );
+        assert_eq!(map_iter.next(), None);
+    }
 }