Moved over to Geolocation-DB client

Cargo.lock

@@ -15,6 +15,7 @@ dependencies = [
 name = "always-online-stun"
 version = "0.1.0"
 dependencies = [
+ "async-trait",
  "futures",
  "log",
  "pretty_env_logger",
@@ -32,6 +33,17 @@ version = "1.0.45"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ee10e43ae4a853c0a3591d4e2ada1719e553be18199d9da9d4a83f5927c2f5c7"
 
+[[package]]
+name = "async-trait"
+version = "0.1.56"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96cf8829f67d2eab0b2dfa42c5d0ef737e0724e4a82b01b3e292456202b19716"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "atty"
 version = "0.2.14"
@@ -548,11 +560,11 @@ checksum = "bc881b2c22681370c6a780e47af9840ef841837bc98118431d4e1868bd0c1086"
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.32"
+version = "1.0.40"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ba508cc11742c0dc5c1659771673afbab7a0efab23aa17e854cbab0837ed0b43"
+checksum = "dd96a1e8ed2596c337f8eae5f24924ec83f5ad5ab21ea8e455d3566c69fbcaf7"
 dependencies = [
- "unicode-xid",
+ "unicode-ident",
 ]
 
 [[package]]
@@ -804,13 +816,13 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "1.0.81"
+version = "1.0.98"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f2afee18b8beb5a596ecb4a2dce128c719b4ba399d34126b9e4396e3f9860966"
+checksum = "c50aef8a904de4c23c788f104b7dddc7d6f79c647c7c8ce4cc8f73eb0ca773dd"
 dependencies = [
  "proc-macro2",
  "quote",
- "unicode-xid",
+ "unicode-ident",
 ]
 
 [[package]]
@@ -979,6 +991,12 @@ version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1a01404663e3db436ed2746d9fefef640d868edae3cceb81c3b8d5732fda678f"
 
+[[package]]
+name = "unicode-ident"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5bd2fe26506023ed7b5e1e315add59d6f584c621d037f9368fea9cfb988f368c"
+
 [[package]]
 name = "unicode-normalization"
 version = "0.1.19"
@@ -988,12 +1006,6 @@ dependencies = [
  "tinyvec",
 ]
 
-[[package]]
-name = "unicode-xid"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8ccb82d61f80a663efe1f787a51b16b5a51e3314d6ac365b08639f52387b33f3"
-
 [[package]]
 name = "untrusted"
 version = "0.7.1"

Cargo.toml

@@ -6,6 +6,7 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+async-trait = { version = "0.1.56", default-features = false }
 futures = { version = "0.3.17" }
 log = { version = "0.4.16" }
 pretty_env_logger = { version = "0.4.0" }

src/geoip.rs

@@ -4,6 +4,23 @@ use std::io::{ErrorKind};
 use std::io::ErrorKind::Other;
 use std::net::IpAddr;
 use std::time::Duration;
+use async_trait::async_trait;
+use serde_json::Value;
+
+type GeoIpData = (f32, f32);
+
+#[async_trait]
+pub(crate) trait GeoIpClient {
+    async fn get_hostname_geoip_info(&self, hostname: &str) -> io::Result<GeoIpData> {
+        self.get_geoip_info(hostname).await
+    }
+
+    async fn get_ip_geoip_info(&self, ip: IpAddr) -> io::Result<GeoIpData> {
+        self.get_geoip_info(ip.to_string().as_str()).await
+    }
+
+    async fn get_geoip_info(&self, hostname_or_ip: &str) -> io::Result<GeoIpData>;
+}
 
 pub(crate) struct IpGeolocationIoClient {
     api_key: String,
@@ -11,25 +28,80 @@ pub(crate) struct IpGeolocationIoClient {
     client: reqwest::Client,
 }
 
-type GeoIpData = (f32, f32);
-
 impl IpGeolocationIoClient {
-    pub(crate) fn new(api_key: String) -> IpGeolocationIoClient {
-        IpGeolocationIoClient {
+    pub(crate) fn new(api_key: String) -> Self {
+        Self {
             api_key,
             url: String::from("https://api.ipgeolocation.io/ipgeo"),
-            client: reqwest::Client::default()
+            client: reqwest::Client::builder()
+                .build().unwrap(),
         }
     }
+}
 
-    pub(crate) fn default() -> IpGeolocationIoClient {
-        IpGeolocationIoClient::new(std::env::var("IPGEOLOCATIONIO_API_KEY")
+impl Default for IpGeolocationIoClient {
+    fn default() -> Self {
+        Self::new(std::env::var("IPGEOLOCATIONIO_API_KEY")
             .expect("Env var IPGEOLOCATIONIO_API_KEY required. Get a free API key at https://ipgeolocation.io"))
     }
 }
 
-impl IpGeolocationIoClient {
-    pub(crate) async fn get_hostname_geoip_info(&self, hostname: &str) -> io::Result<GeoIpData> {
+#[async_trait]
+impl GeoIpClient for IpGeolocationIoClient {
+
+    async fn get_geoip_info(&self, hostname_or_ip: &str) -> io::Result<GeoIpData> {
+        let response = self.client.get(self.url.as_str())
+            .query(&[("apiKey", self.api_key.as_str())])
+            .query(&[("ip", hostname_or_ip)])
+            .query(&[("fields", "latitude,longitude")])
+            .timeout(Duration::from_secs(5))
+            .send()
+            .await
+            .map_err(|err| io::Error::new(ErrorKind::Other, err))?
+            .json::<Value>()
+            .await
+            .map_err(|err| io::Error::new(ErrorKind::Other, err))?;
+
+        let lat = response
+            .get("latitude")
+            .map(|v| v.as_str().unwrap_or("0"))
+            .map(|s| s.parse().unwrap_or(0 as f32))
+            .unwrap_or(0 as f32);
+
+        let lon = response
+            .get("longitude")
+            .map(|v| v.as_str().unwrap_or("0"))
+            .map(|s| s.parse().unwrap_or(0 as f32))
+            .unwrap_or(0 as f32);
+
+        Ok((lat, lon))
+    }
+}
+
+pub(crate) struct GeolocationDbClient {
+    url: String,
+    client: reqwest::Client,
+}
+
+impl Default for GeolocationDbClient {
+    fn default() -> Self {
+        Self {
+            url: String::from("https://geolocation-db.com/json"),
+            client: reqwest::Client::builder()
+                .build()
+                .unwrap(),
+        }
+    }
+}
+
+struct GeolocationDbClientResponse {
+    latitude: Option<f32>,
+    longitude: Option<f32>,
+}
+
+#[async_trait]
+impl GeoIpClient for GeolocationDbClient {
+    async fn get_hostname_geoip_info(&self, hostname: &str) -> io::Result<GeoIpData> {
         self.get_geoip_info(hostname).await
     }
 
@@ -38,56 +110,55 @@ impl IpGeolocationIoClient {
     }
 
     async fn get_geoip_info(&self, hostname_or_ip: &str) -> io::Result<GeoIpData> {
-        let response = self.client.get(self.url.as_str())
-            .query(&[("apiKey", self.api_key.as_str())])
-            .query(&[("ip", hostname_or_ip)])
-            .query(&[("fields", "latitude,longitude")])
-            .timeout(Duration::from_secs(1))
+        let url = self.url.clone() + "/" + hostname_or_ip;
+
+        let response = self.client.get(url)
+            .timeout(Duration::from_secs(5))
             .send()
             .await
             .map_err(|err| io::Error::new(ErrorKind::Other, err))?
-            .json::<HashMap<String, String>>()
+            .json::<Value>()
             .await
             .map_err(|err| io::Error::new(ErrorKind::Other, err))?;
 
-        let lat = response.get("latitude")
-            .cloned()
-            .map(|lat_str| lat_str.parse().unwrap())
-            .unwrap_or(0 as f32);
+        let lat = response
+            .get("latitude")
+            .map(|v| v.as_f64().unwrap_or(0 as f64))
+            .unwrap_or(0 as f64) as f32;
 
-        let lon = response.get("longitude")
-            .cloned()
-            .map(|lon_str| lon_str.parse().unwrap())
-            .unwrap_or(0 as f32);
+        let lon = response
+            .get("longitude")
+            .map(|v| v.as_f64().unwrap_or(0 as f64))
+            .unwrap_or(0 as f64) as f32;
 
         Ok((lat, lon))
     }
 }
 
-pub(crate) struct CachedIpGeolocationIpClient {
-    path: String,
-    client: IpGeolocationIoClient,
+pub(crate) struct CachedIpGeolocationIpClient<T: GeoIpClient + Default> {
+    cachefile_path: String,
+    client_impl: T,
     map: BTreeMap<String, GeoIpData>,
 }
 
-impl CachedIpGeolocationIpClient {
-    pub(crate) async fn new(filename: String) -> io::Result<CachedIpGeolocationIpClient> {
-        let cache_str = tokio::fs::read_to_string(filename.as_str()).await?;
+impl <T: GeoIpClient + Default> CachedIpGeolocationIpClient<T> {
+    pub(crate) async fn new(cachefile_path: String) -> io::Result<Self> {
+        let cache_str = tokio::fs::read_to_string(cachefile_path.as_str()).await?;
 
         let map: BTreeMap<String, GeoIpData> = serde_json::de::from_str(cache_str.as_str())
             .map_err(|err| io::Error::new(Other, err))?;
 
-        let client = CachedIpGeolocationIpClient {
-            path: filename,
-            client: IpGeolocationIoClient::default(),
-            map
-        };
+        let client_impl = T::default();
 
-        Ok(client)
+        Ok(Self {
+            cachefile_path,
+            client_impl,
+            map
+        })
     }
 
-    pub(crate) async fn default() -> io::Result<CachedIpGeolocationIpClient> {
-        CachedIpGeolocationIpClient::new(String::from("geoip_cache.txt")).await
+    pub(crate) async fn default() -> io::Result<Self> {
+        Self::new(String::from("geoip_cache.txt")).await
     }
 
     pub(crate) async fn get_hostname_geoip_info(&mut self, hostname: &str) -> io::Result<GeoIpData> {
@@ -101,15 +172,30 @@ impl CachedIpGeolocationIpClient {
     pub(crate) async fn save(&self) -> io::Result<()> {
         let str = serde_json::ser::to_string_pretty(&self.map)
             .map_err(|err| io::Error::new(Other, err))?;
-        tokio::fs::write(self.path.as_str(), str).await
+        tokio::fs::write(self.cachefile_path.as_str(), str).await
     }
 
     async fn get_geoip_info(&mut self, hostname_or_ip: &str) -> io::Result<GeoIpData> {
         if let Some(geo_ip_data) = self.map.get(hostname_or_ip).cloned() {
             return Ok(geo_ip_data)
         }
-        let geo_ip_data = self.client.get_geoip_info(hostname_or_ip).await?;
+        let geo_ip_data = self.client_impl.get_geoip_info(hostname_or_ip).await?;
         self.map.insert(String::from(hostname_or_ip), geo_ip_data.clone());
         Ok(geo_ip_data)
     }
 }
+
+
+#[cfg(test)]
+mod tests {
+    use std::net::Ipv4Addr;
+    use super::*;
+
+    #[tokio::test]
+    async fn geolocation_db_client() {
+        let (lat, lon) = GeolocationDbClient::default()
+            .get_ip_geoip_info(IpAddr::V4(Ipv4Addr::from([1,1,1,1]))).await.unwrap();
+        assert_ne!(0 as f32, lat);
+        assert_ne!(0 as f32, lon);
+    }
+}

src/main.rs

@@ -4,6 +4,7 @@ use std::rc::Rc;
 use std::time::Duration;
 use futures::StreamExt;
 use tokio::time::Instant;
+use crate::geoip::GeolocationDbClient;
 use crate::utils::join_all_with_semaphore;
 use crate::outputs::{ValidHosts, ValidIpV4s, ValidIpV6s};
 use crate::servers::StunServer;
@@ -29,7 +30,7 @@ async fn main() -> io::Result<()> {
         .parse()
         .expect("IS_BEHIND_NAT must be true or false");
 
-    let client = Rc::new(RefCell::new(geoip::CachedIpGeolocationIpClient::default().await?));
+    let client = Rc::new(RefCell::new(geoip::CachedIpGeolocationIpClient::<GeolocationDbClient>::default().await?));
 
     let stun_servers = servers::get_stun_servers().await?;