From 152f1a839333bf61d6230622314637883d81192a Mon Sep 17 00:00:00 2001
From: Luca BRUNO <luca.bruno@coreos.com>
Date: Wed, 27 Nov 2019 17:20:54 +0000
Subject: [PATCH] ipnetwork: add netmask constructors

This adds constructors to build all network types from a network
address and network mask pair.
---
 src/ipv4.rs | 35 +++++++++++++++++++++++++++++++++++
 src/ipv6.rs | 32 ++++++++++++++++++++++++++++++++
 src/lib.rs  |  8 ++++++++
 3 files changed, 75 insertions(+)

diff --git a/src/ipv4.rs b/src/ipv4.rs
index 6f88053..95e2a27 100644
--- a/src/ipv4.rs
+++ b/src/ipv4.rs
@@ -32,6 +32,7 @@ impl Serialize for Ipv4Network {
 
 impl Ipv4Network {
     /// Constructs a new `Ipv4Network` from any `Ipv4Addr` and a prefix denoting the network size.
+    ///
     /// If the prefix is larger than 32 this will return an `IpNetworkError::InvalidPrefix`.
     pub fn new(addr: Ipv4Addr, prefix: u8) -> Result<Ipv4Network, IpNetworkError> {
         if prefix > IPV4_BITS {
@@ -41,6 +42,21 @@ impl Ipv4Network {
         }
     }
 
+    /// Constructs a new `Ipv4Network` from a network address and a network mask.
+    ///
+    /// If the netmask is not valid this will return an `IpNetworkError::InvalidPrefix`.
+    pub fn with_netmask(
+        netaddr: Ipv4Addr,
+        netmask: Ipv4Addr,
+    ) -> Result<Ipv4Network, IpNetworkError> {
+        let prefix = ipv4_mask_to_prefix(netmask)?;
+        let net = Self {
+            addr: netaddr,
+            prefix,
+        };
+        Ok(net)
+    }
+
     /// Returns an iterator over `Ipv4Network`. Each call to `next` will return the next
     /// `Ipv4Addr` in the given network. `None` will be returned when there are no more
     /// addresses.
@@ -257,6 +273,7 @@ impl Iterator for Ipv4NetworkIterator {
 }
 
 /// Converts a `Ipv4Addr` network mask into a prefix.
+///
 /// If the mask is invalid this will return an `IpNetworkError::InvalidPrefix`.
 pub fn ipv4_mask_to_prefix(mask: Ipv4Addr) -> Result<u8, IpNetworkError> {
     let mask = u32::from(mask);
@@ -449,6 +466,24 @@ mod test {
         assert!(prefix.is_err());
     }
 
+    #[test]
+    fn ipv4network_with_netmask() {
+        {
+            // Positive test-case.
+            let addr = Ipv4Addr::new(127, 0, 0, 1);
+            let mask = Ipv4Addr::new(255, 0, 0, 0);
+            let net = Ipv4Network::with_netmask(addr, mask).unwrap();
+            let expected = Ipv4Network::new(Ipv4Addr::new(127, 0, 0, 1), 8).unwrap();
+            assert_eq!(net, expected);
+        }
+        {
+            // Negative test-case.
+            let addr = Ipv4Addr::new(127, 0, 0, 1);
+            let mask = Ipv4Addr::new(255, 0, 255, 0);
+            Ipv4Network::with_netmask(addr, mask).unwrap_err();
+        }
+    }
+
     #[test]
     fn ipv4network_from_ipv4addr() {
         let net = Ipv4Network::from(Ipv4Addr::new(127, 0, 0, 1));
diff --git a/src/ipv6.rs b/src/ipv6.rs
index d58c870..0790f93 100644
--- a/src/ipv6.rs
+++ b/src/ipv6.rs
@@ -33,6 +33,7 @@ impl Serialize for Ipv6Network {
 
 impl Ipv6Network {
     /// Constructs a new `Ipv6Network` from any `Ipv6Addr` and a prefix denoting the network size.
+    ///
     /// If the prefix is larger than 128 this will return an `IpNetworkError::InvalidPrefix`.
     pub fn new(addr: Ipv6Addr, prefix: u8) -> Result<Ipv6Network, IpNetworkError> {
         if prefix > IPV6_BITS {
@@ -42,6 +43,18 @@ impl Ipv6Network {
         }
     }
 
+    /// Constructs a new `Ipv6Network` from a network address and a network mask.
+    ///
+    /// If the netmask is not valid this will return an `IpNetworkError::InvalidPrefix`.
+    pub fn with_netmask(netaddr: Ipv6Addr, netmask: Ipv6Addr) -> Result<Self, IpNetworkError> {
+        let prefix = ipv6_mask_to_prefix(netmask)?;
+        let net = Self {
+            addr: netaddr,
+            prefix,
+        };
+        Ok(net)
+    }
+
     /// Returns an iterator over `Ipv6Network`. Each call to `next` will return the next
     /// `Ipv6Addr` in the given network. `None` will be returned when there are no more
     /// addresses.
@@ -366,6 +379,25 @@ mod test {
         assert!(prefix.is_err());
     }
 
+    #[test]
+    fn ipv6network_with_netmask() {
+        {
+            // Positive test-case.
+            let addr = Ipv6Addr::new(0xff01, 0, 0, 0x17, 0, 0, 0, 0x2);
+            let mask = Ipv6Addr::new(0xffff, 0xffff, 0xffff, 0, 0, 0, 0, 0);
+            let net = Ipv6Network::with_netmask(addr, mask).unwrap();
+            let expected =
+                Ipv6Network::new(Ipv6Addr::new(0xff01, 0, 0, 0x17, 0, 0, 0, 0x2), 48).unwrap();
+            assert_eq!(net, expected);
+        }
+        {
+            // Negative test-case.
+            let addr = Ipv6Addr::new(0xff01, 0, 0, 0x17, 0, 0, 0, 0x2);
+            let mask = Ipv6Addr::new(0, 0, 0xffff, 0xffff, 0, 0, 0, 0);
+            Ipv6Network::with_netmask(addr, mask).unwrap_err();
+        }
+    }
+
     #[test]
     fn iterator_v6() {
         let cidr: Ipv6Network = "2001:db8::/126".parse().unwrap();
diff --git a/src/lib.rs b/src/lib.rs
index 5401b73..24abf7a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -63,6 +63,14 @@ impl IpNetwork {
         }
     }
 
+    /// Constructs a new `IpNetwork` from a network address and a network mask.
+    ///
+    /// If the netmask is not valid this will return an `IpNetworkError::InvalidPrefix`.
+    pub fn with_netmask(netaddr: IpAddr, netmask: IpAddr) -> Result<Self, IpNetworkError> {
+        let prefix = ip_mask_to_prefix(netmask)?;
+        Self::new(netaddr, prefix)
+    }
+
     /// Returns the IP part of a given `IpNetwork`
     pub fn ip(&self) -> IpAddr {
         match *self {