From 8463a3fc5e15f051f1a3a10c8493f0c66f36afc6 Mon Sep 17 00:00:00 2001
From: Edward Barnard <eabarnard@gmail.com>
Date: Thu, 14 Jul 2016 21:37:39 +0200
Subject: [PATCH] Add ip_mask_to_prefix, ipv4_mask_to_prefix and
 ipv6_mask_to_prefix

---
 src/ipv4.rs | 26 ++++++++++++++++++++++++++
 src/ipv6.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs  | 13 +++++++++++--
 3 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/src/ipv4.rs b/src/ipv4.rs
index 451438b..6aa82e9 100644
--- a/src/ipv4.rs
+++ b/src/ipv4.rs
@@ -222,6 +222,18 @@ 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);
+
+    let prefix = (!mask).leading_zeros() as u8;
+    if ((mask as u64) << prefix) & 0xffffffff != 0 {
+        Err(IpNetworkError::InvalidPrefix)
+    } else {
+        Ok(prefix)
+    }
+}
 
 #[cfg(test)]
 mod test {
@@ -415,4 +427,18 @@ mod test {
         }
         assert_eq!(None, iter.next());
     }
+
+    #[test]
+    fn v4_mask_to_prefix() {
+        let mask = Ipv4Addr::new(255, 255, 255, 128);
+        let prefix = ipv4_mask_to_prefix(mask).unwrap();
+        assert_eq!(prefix, 25);
+    }
+
+    #[test]
+    fn invalid_v4_mask_to_prefix() {
+        let mask = Ipv4Addr::new(255, 0, 255, 0);
+        let prefix = ipv4_mask_to_prefix(mask);
+        assert!(prefix.is_err());
+    }
 }
diff --git a/src/ipv6.rs b/src/ipv6.rs
index 90503f1..95f12b8 100644
--- a/src/ipv6.rs
+++ b/src/ipv6.rs
@@ -6,6 +6,7 @@ use std::str::FromStr;
 use common::{IpNetworkError, cidr_parts, parse_prefix};
 
 const IPV6_BITS: u8 = 128;
+const IPV6_SEGMENT_BITS: u8 = 16;
 
 #[derive(Debug,Clone,Copy,Hash,PartialEq,Eq)]
 pub struct Ipv6Network {
@@ -94,6 +95,40 @@ impl fmt::Display for Ipv6Network {
     }
 }
 
+/// Converts a `Ipv6Addr` network mask into a prefix.
+/// If the mask is invalid this will return an `IpNetworkError::InvalidPrefix`.
+pub fn ipv6_mask_to_prefix(mask: Ipv6Addr) -> Result<u8, IpNetworkError> {
+    let mask = mask.segments();
+    let mut mask_iter = mask.into_iter();
+
+    // Count the number of set bits from the start of the address
+    let mut prefix = 0;
+    for &segment in &mut mask_iter {
+        if segment == 0xffff {
+            prefix += IPV6_SEGMENT_BITS;
+        } else if segment == 0 {
+            // Prefix finishes on a segment boundary
+            break;
+        } else {
+            let prefix_bits = (!segment).leading_zeros() as u8;
+            // Check that the remainder of the bits are all unset
+            if segment << prefix_bits != 0 {
+                return Err(IpNetworkError::InvalidPrefix);
+            }
+            prefix += prefix_bits;
+            break;
+        }
+    }
+
+    // Now check all the remaining bits are unset
+    for &segment in mask_iter {
+        if segment != 0 {
+            return Err(IpNetworkError::InvalidPrefix);
+        }
+    }
+
+    Ok(prefix)
+}
 
 #[cfg(test)]
 mod test {
@@ -158,4 +193,18 @@ mod test {
         let ip = Ipv6Addr::new(0xff01, 0, 0, 0x17, 0xffff, 0, 0, 0x2);
         assert!(!cidr.contains(ip));
     }
+
+    #[test]
+    fn v6_mask_to_prefix() {
+        let mask = Ipv6Addr::new(0xffff, 0xffff, 0xffff, 0, 0, 0, 0, 0);
+        let prefix = ipv6_mask_to_prefix(mask).unwrap();
+        assert_eq!(prefix, 48);
+    }
+
+    #[test]
+    fn invalid_v6_mask_to_prefix() {
+        let mask = Ipv6Addr::new(0, 0, 0xffff, 0xffff, 0, 0, 0, 0);
+        let prefix = ipv6_mask_to_prefix(mask);
+        assert!(prefix.is_err());
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 791d520..c223379 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,8 +10,8 @@ mod ipv4;
 mod ipv6;
 mod common;
 
-pub use ipv4::Ipv4Network;
-pub use ipv6::Ipv6Network;
+pub use ipv4::{Ipv4Network, ipv4_mask_to_prefix};
+pub use ipv6::{Ipv6Network, ipv6_mask_to_prefix};
 pub use common::IpNetworkError;
 
 // A network
@@ -43,3 +43,12 @@ impl IpNetwork {
         }
     }
 }
+
+/// Converts a `IpAddr` network mask into a prefix.
+/// If the mask is invalid this will return an `IpNetworkError::InvalidPrefix`.
+pub fn ip_mask_to_prefix(mask: IpAddr) -> Result<u8, IpNetworkError> {
+    match mask {
+        IpAddr::V4(mask) => ipv4_mask_to_prefix(mask),
+        IpAddr::V6(mask) => ipv6_mask_to_prefix(mask),
+    }
+}
\ No newline at end of file