Skip to content

Commit

Permalink
req level ipv6 support
Browse files Browse the repository at this point in the history
  • Loading branch information
talhahwahla committed Jan 8, 2024
1 parent d97f769 commit 49a8826
Showing 1 changed file with 46 additions and 26 deletions.
72 changes: 46 additions & 26 deletions src/ipinfo.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -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
///
Expand All @@ -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(
Expand Down Expand Up @@ -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
}
}

Expand All @@ -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();

Expand Down Expand Up @@ -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);
}

Expand All @@ -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))
Expand Down Expand Up @@ -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(),
Expand All @@ -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()
Expand Down Expand Up @@ -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);

Expand Down

0 comments on commit 49a8826

Please sign in to comment.