From 49a88266d5f8363d50aead5910000d075581221a Mon Sep 17 00:00:00 2001
From: talhahwahla <mhameed.bee15seecs@seecs.edu.pk>
Date: Mon, 8 Jan 2024 18:14:49 +0500
Subject: [PATCH] req level ipv6 support

---
 src/ipinfo.rs | 72 ++++++++++++++++++++++++++++++++-------------------
 1 file changed, 46 insertions(+), 26 deletions(-)

diff --git a/src/ipinfo.rs b/src/ipinfo.rs
index 0a3efab..bc2d706 100644
--- a/src/ipinfo.rs
+++ b/src/ipinfo.rs
@@ -31,6 +31,10 @@ use tokio::time::timeout;
 
 const COUNTRY_FLAG_URL: &str =
     "https://cdn.ipinfo.io/static/images/countries-flags/";
+
+const BASE_URL: &str = "https://ipinfo.io";
+const BASE_URL_V6: &str = "https://v6.ipinfo.io";
+
 /// IpInfo structure configuration.
 pub struct IpInfoConfig {
     /// IPinfo access token.
@@ -75,7 +79,6 @@ impl Default for IpInfoConfig {
 
 /// IPinfo requests context structure.
 pub struct IpInfo {
-    url: String,
     token: Option<String>,
     client: reqwest::Client,
     cache: LruCache<String, IpDetails>,
@@ -103,7 +106,7 @@ impl Default for BatchReqOpts {
 }
 
 impl IpInfo {
-    /// Construct a new IpInfo structure with the default URL "https://ipinfo.io".
+    /// Construct a new IpInfo structure.
     ///
     /// # Examples
     ///
@@ -113,28 +116,10 @@ impl IpInfo {
     /// let ipinfo = IpInfo::new(Default::default()).expect("should construct");
     /// ```
     pub fn new(config: IpInfoConfig) -> Result<Self, IpError> {
-        Self::base_request(config, "https://ipinfo.io")
-    }
-
-    /// Construct a new IpInfo structure with the URL "https://v6.ipinfo.io/".
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// use ipinfo::IpInfo;
-    ///
-    /// let ipinfo = IpInfo::new_v6(Default::default()).expect("should construct");
-    /// ```
-    pub fn new_v6(config: IpInfoConfig) -> Result<Self, IpError> {
-        Self::base_request(config, "https://v6.ipinfo.io/")
-    }
-
-    fn base_request(config: IpInfoConfig, url: &str) -> Result<Self, IpError> {
         let client =
             reqwest::Client::builder().timeout(config.timeout).build()?;
 
         let mut ipinfo_obj = Self {
-            url: url.to_owned(),
             client,
             token: config.token,
             cache: LruCache::new(
@@ -196,17 +181,34 @@ impl IpInfo {
         &mut self,
         ips: &[&str],
         batch_config: BatchReqOpts,
+    ) -> Result<HashMap<String, IpDetails>, IpError> {
+        self.lookup_batch_internal(ips, batch_config, BASE_URL).await
+    }
+
+    pub async fn lookup_batch_v6(
+        &mut self,
+        ips: &[&str],
+        batch_config: BatchReqOpts,
+    ) -> Result<HashMap<String, IpDetails>, IpError> {
+        self.lookup_batch_internal(ips, batch_config, BASE_URL_V6).await
+    }
+
+    async fn lookup_batch_internal(
+        &mut self,
+        ips: &[&str],
+        batch_config: BatchReqOpts,
+        base_url: &'static str,
     ) -> Result<HashMap<String, IpDetails>, IpError> {
         // Handle the total timeout condition
         if let Some(total_timeout) = batch_config.timeout_total {
-            match timeout(total_timeout, self._lookup_batch(ips, batch_config))
+            match timeout(total_timeout, self._lookup_batch(ips, batch_config, base_url))
                 .await
             {
                 Ok(result) => result,
                 Err(_) => Err(err!(TimeOutError)),
             }
         } else {
-            self._lookup_batch(ips, batch_config).await
+            self._lookup_batch(ips, batch_config, base_url).await
         }
     }
 
@@ -215,6 +217,7 @@ impl IpInfo {
         &mut self,
         ips: &[&str],
         batch_config: BatchReqOpts,
+        base_url: &str
     ) -> Result<HashMap<String, IpDetails>, IpError> {
         let mut results: HashMap<String, IpDetails> = HashMap::new();
 
@@ -248,7 +251,7 @@ impl IpInfo {
 
         // Make batched requests
         for batch in work.chunks(batch_config.batch_size as usize) {
-            let response = self.batch_request(client.clone(), batch).await?;
+            let response = self.batch_request(client.clone(), batch, base_url).await?;
             results.extend(response);
         }
 
@@ -272,10 +275,11 @@ impl IpInfo {
         &self,
         client: reqwest::Client,
         ips: &[&str],
+        base_url: &str
     ) -> Result<HashMap<String, IpDetails>, IpError> {
         // Lookup cache misses which are not bogon
         let response = client
-            .post(&format!("{}/batch", self.url))
+            .post(&format!("{}/batch", base_url))
             .headers(Self::construct_headers())
             .bearer_auth(self.token.as_deref().unwrap_or_default())
             .json(&json!(ips))
@@ -318,6 +322,14 @@ impl IpInfo {
     /// }
     /// ```
     pub async fn lookup(&mut self, ip: &str) -> Result<IpDetails, IpError> {
+        self.lookup_internal(ip, BASE_URL).await
+    }
+
+    pub async fn lookup_v6(&mut self, ip: &str) -> Result<IpDetails, IpError> {
+        self.lookup_internal(ip, BASE_URL_V6).await
+    }
+
+    async fn lookup_internal(&mut self, ip: &str, base_url: &str) -> Result<IpDetails, IpError> {
         if is_bogon(ip) {
             return Ok(IpDetails {
                 ip: ip.to_string(),
@@ -336,7 +348,7 @@ impl IpInfo {
         // lookup in case of a cache miss
         let response = self
             .client
-            .get(&format!("{}/{}", self.url, ip))
+            .get(&format!("{}/{}", base_url, ip))
             .headers(Self::construct_headers())
             .bearer_auth(self.token.as_deref().unwrap_or_default())
             .send()
@@ -381,11 +393,19 @@ impl IpInfo {
     /// }
     /// ```
     pub async fn get_map(&self, ips: &[&str]) -> Result<String, IpError> {
+        self.get_map_internal(ips, BASE_URL).await
+    }
+
+    pub async fn get_map_v6(&self, ips: &[&str]) -> Result<String, IpError> {
+        self.get_map_internal(ips, BASE_URL_V6).await
+    }
+
+    async fn get_map_internal(&self, ips: &[&str], base_url: &str) -> Result<String, IpError> {
         if ips.len() > 500_000 {
             return Err(err!(MapLimitError));
         }
 
-        let map_url = &format!("{}/tools/map?cli=1", self.url);
+        let map_url = &format!("{}/tools/map?cli=1", base_url);
         let client = self.client.clone();
         let json_ips = serde_json::json!(ips);