From cd6cf4add4d65207f17bda20569b28b527f2e2fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Romanowski?= <pawroman@pawroman.dev>
Date: Fri, 9 Oct 2020 14:54:50 +0200
Subject: [PATCH] Add Uri::host_ip method and tests

This lets extract the host IP address, if defined for a website.
Mail addresses are not supported.
---
 src/extract.rs | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/src/extract.rs b/src/extract.rs
index c1f7d72..9cc930c 100644
--- a/src/extract.rs
+++ b/src/extract.rs
@@ -1,5 +1,6 @@
 use linkify::LinkFinder;
 
+use std::net::IpAddr;
 use std::{collections::HashSet, fmt::Display};
 use url::Url;
 
@@ -23,6 +24,17 @@ impl Uri {
             Uri::Mail(_address) => None,
         }
     }
+
+    pub fn host_ip(&self) -> Option<IpAddr> {
+        match self {
+            Self::Website(url) => match url.host()? {
+                url::Host::Ipv4(v4_addr) => Some(v4_addr.into()),
+                url::Host::Ipv6(v6_addr) => Some(v6_addr.into()),
+                _ => None,
+            },
+            Self::Mail(_) => None,
+        }
+    }
 }
 
 impl Display for Uri {
@@ -56,6 +68,7 @@ pub(crate) fn extract_links(input: &str) -> HashSet<Uri> {
 mod test {
     use super::*;
     use std::iter::FromIterator;
+    use std::net::{Ipv4Addr, Ipv6Addr};
 
     #[test]
     fn test_extract_markdown_links() {
@@ -113,4 +126,30 @@ mod test {
         assert!(links.len() == 1);
         assert_eq!(links[0].as_str(), expected);
     }
+
+    #[test]
+    fn test_uri_host_ip_v4() {
+        let uri =
+            Uri::Website(Url::parse("http://127.0.0.1").expect("Expected URI with valid IPv4"));
+        let ip = uri.host_ip().expect("Expected a valid IPv4");
+        assert_eq!(ip, IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
+    }
+
+    #[test]
+    fn test_uri_host_ip_v6() {
+        let uri =
+            Uri::Website(Url::parse("https://[2020::0010]").expect("Expected URI with valid IPv6"));
+        let ip = uri.host_ip().expect("Expected a valid IPv6");
+        assert_eq!(
+            ip,
+            IpAddr::V6(Ipv6Addr::new(0x2020, 0, 0, 0, 0, 0, 0, 0x10))
+        );
+    }
+
+    #[test]
+    fn test_uri_host_ip_no_ip() {
+        let uri = Uri::Website(Url::parse("https://some.cryptic/url").expect("Expected valid URI"));
+        let ip = uri.host_ip();
+        assert!(ip.is_none());
+    }
 }