use itertools::{iproduct, Product};
use std::net::{IpAddr, SocketAddr};
pub struct SocketIterator<'s> {
product_it:
Product<Box<std::slice::Iter<'s, u16>>, Box<std::slice::Iter<'s, std::net::IpAddr>>>,
}
impl<'s> SocketIterator<'s> {
pub fn new(ips: &'s [IpAddr], ports: &'s [u16]) -> Self {
let ports_it = Box::new(ports.iter());
let ips_it = Box::new(ips.iter());
Self {
product_it: iproduct!(ports_it, ips_it),
}
}
}
impl<'s> Iterator for SocketIterator<'s> {
type Item = SocketAddr;
fn next(&mut self) -> Option<Self::Item> {
match self.product_it.next() {
None => None,
Some((port, ip)) => Some(SocketAddr::new(*ip, *port)),
}
}
}
#[cfg(test)]
mod tests {
use super::SocketIterator;
use std::net::{IpAddr, SocketAddr};
#[test]
fn goes_through_every_ip_port_combination() {
let addrs = vec![
"127.0.0.1".parse::<IpAddr>().unwrap(),
"192.168.0.1".parse::<IpAddr>().unwrap(),
];
let ports: Vec<u16> = vec![22, 80, 443];
let mut it = SocketIterator::new(&addrs, &ports);
assert_eq!(Some(SocketAddr::new(addrs[0], ports[0])), it.next());
assert_eq!(Some(SocketAddr::new(addrs[1], ports[0])), it.next());
assert_eq!(Some(SocketAddr::new(addrs[0], ports[1])), it.next());
assert_eq!(Some(SocketAddr::new(addrs[1], ports[1])), it.next());
assert_eq!(Some(SocketAddr::new(addrs[0], ports[2])), it.next());
assert_eq!(Some(SocketAddr::new(addrs[1], ports[2])), it.next());
assert_eq!(None, it.next());
}
}