Add checks to guard against overflow in iterator methods. (#76)

Peternator7 · Nov 17, 2019 · e981f24 · e981f24
1 parent 9310855
commit e981f24
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/strum_macros/src/macros/enum_iter.rs b/strum_macros/src/macros/enum_iter.rs
@@ -83,12 +83,14 @@ pub fn enum_iter_inner(ast: &syn::DeriveInput) -> TokenStream {
                     #(#arms),*
                 };
 
-                self.idx += 1;
+                if self.idx < #variant_count {
+                    self.idx += 1;
+                }
                 output
             }
 
             fn size_hint(&self) -> (usize, Option<usize>) {
-                let t = #variant_count - self.idx;
+                let t = if self.idx >= #variant_count { 0 } else { #variant_count - self.idx };
                 (t, Some(t))
             }
         }

diff --git a/strum_tests/tests/enum_iter.rs b/strum_tests/tests/enum_iter.rs
@@ -63,6 +63,9 @@ fn len_test() {
     i.next();
 
     assert_eq!(0, i.len());
+    i.next();
+
+    assert_eq!(0, i.size_hint().1.unwrap());
 }
 
 #[test]