Fix bug with overflow in iterator

This commit is contained in:
Michal 'vorner' Vaner
2019-12-07 21:24:19 +01:00
parent dbf366cfdf
commit 5b09d82d43
2 changed files with 17 additions and 17 deletions

View File

@ -62,8 +62,8 @@ impl Ipv4Network {
/// addresses. /// addresses.
pub fn iter(&self) -> Ipv4NetworkIterator { pub fn iter(&self) -> Ipv4NetworkIterator {
let start = u32::from(self.network()); let start = u32::from(self.network());
let end = start + self.size(); let end = start + (self.size() - 1);
Ipv4NetworkIterator { next: start, end } Ipv4NetworkIterator { next: Some(start), end }
} }
pub fn ip(&self) -> Ipv4Addr { pub fn ip(&self) -> Ipv4Addr {
@ -255,7 +255,7 @@ impl From<Ipv4Addr> for Ipv4Network {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct Ipv4NetworkIterator { pub struct Ipv4NetworkIterator {
next: u32, next: Option<u32>,
end: u32, end: u32,
} }
@ -263,13 +263,13 @@ impl Iterator for Ipv4NetworkIterator {
type Item = Ipv4Addr; type Item = Ipv4Addr;
fn next(&mut self) -> Option<Ipv4Addr> { fn next(&mut self) -> Option<Ipv4Addr> {
if self.next < self.end { let next = self.next?;
let next = Ipv4Addr::from(self.next as u32); self.next = if next == self.end {
self.next += 1;
Some(next)
} else {
None None
} } else {
Some(next + 1)
};
Some(next.into())
} }
} }

View File

@ -70,7 +70,7 @@ impl Ipv6Network {
let end: u128 = dec | mask; let end: u128 = dec | mask;
Ipv6NetworkIterator { Ipv6NetworkIterator {
next: start, next: Some(start),
end: end, end: end,
} }
} }
@ -229,7 +229,7 @@ impl From<Ipv6Addr> for Ipv6Network {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct Ipv6NetworkIterator { pub struct Ipv6NetworkIterator {
next: u128, next: Option<u128>,
end: u128, end: u128,
} }
@ -237,13 +237,13 @@ impl Iterator for Ipv6NetworkIterator {
type Item = Ipv6Addr; type Item = Ipv6Addr;
fn next(&mut self) -> Option<Ipv6Addr> { fn next(&mut self) -> Option<Ipv6Addr> {
if self.next <= self.end { let next = self.next?;
let next = Ipv6Addr::from(self.next); self.next = if next == self.end {
self.next += 1;
Some(next)
} else {
None None
} } else {
Some(next + 1)
};
Some(next.into())
} }
} }