Skip to content

Commit

Permalink
merge GH-784 into GH-468
Browse files Browse the repository at this point in the history
  • Loading branch information
czarte committed Oct 3, 2024
2 parents 6152a9e + a5e8cf2 commit 5bc34df
Show file tree
Hide file tree
Showing 13 changed files with 2,456 additions and 416 deletions.
2,261 changes: 2,210 additions & 51 deletions ip_country/Cargo.lock

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions ip_country/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ edition = "2021"
license = "GPL-3.0-only"
authors = ["Dan Wiebe <[email protected]>", "MASQ"]
description = "Handle embedding IP-address-to-country data in MASQ Node"
workspace = "../node"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

Expand Down
23 changes: 17 additions & 6 deletions ip_country/src/country_block_serde.rs
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,11 @@ impl CountryBlockSerializer {
}
}

struct VersionedIPSerializer<IPType: Debug, SegmentNumRep: Debug, const SEGMENTS_COUNT: usize> {
struct VersionedIPSerializer<IPType, SegmentNumRep, const SEGMENTS_COUNT: usize>
where
IPType: Debug,
SegmentNumRep: Debug,
{
prev_start: VersionedIP<IPType, SegmentNumRep, SEGMENTS_COUNT>,
prev_end: VersionedIP<IPType, SegmentNumRep, SEGMENTS_COUNT>,
block_count: usize,
Expand Down Expand Up @@ -310,7 +314,7 @@ where

// Rust forces public visibility on traits that come to be used as type boundaries in any public
// interface. This is how we can meet the requirements while the implementations of such
// traits becomes ineffective from farther than this file. It works as a form of prevention to
// traits become ineffective beyond this file. It works as a form of prevention to
// namespace pollution for such kind of trait to be implemented on massively common types,
// here namely Ipv4Addr or Ipv6Addr
mod semi_private_items {
Expand Down Expand Up @@ -424,7 +428,7 @@ pub trait DeserializerPublic {
fn next(&mut self) -> Option<CountryBlock>;
}

type Ipv4CountryBlockDeserializer = CountryBlockDeserializer<Ipv4Addr, u8, 4>;
pub type Ipv4CountryBlockDeserializer = CountryBlockDeserializer<Ipv4Addr, u8, 4>;

impl Ipv4CountryBlockDeserializer {
pub fn new(country_data: (Vec<u64>, usize)) -> Self {
Expand All @@ -439,7 +443,7 @@ impl Iterator for Ipv4CountryBlockDeserializer {
}
}

type Ipv6CountryBlockDeserializer = CountryBlockDeserializer<Ipv6Addr, u16, 8>;
pub type Ipv6CountryBlockDeserializer = CountryBlockDeserializer<Ipv6Addr, u16, 8>;

impl Ipv6CountryBlockDeserializer {
pub fn new(country_data: (Vec<u64>, usize)) -> Self {
Expand Down Expand Up @@ -594,8 +598,11 @@ where
segment_num_rep: PhantomData<SegmentNumRep>,
}

impl<IPType: Debug, SegmentNumRep: Debug, const SEGMENTS_COUNT: usize>
impl<IPType, SegmentNumRep, const SEGMENTS_COUNT: usize>
VersionedIP<IPType, SegmentNumRep, SEGMENTS_COUNT>
where
IPType: Debug,
SegmentNumRep: Debug,
{
fn new(ip: IPType) -> VersionedIP<IPType, SegmentNumRep, SEGMENTS_COUNT> {
let segment_num_rep = Default::default();
Expand All @@ -607,7 +614,11 @@ impl<IPType: Debug, SegmentNumRep: Debug, const SEGMENTS_COUNT: usize>
}

#[derive(Debug)]
struct StreamRecord<IPType: Debug, SegmentNumRep: Debug, const SEGMENTS_COUNT: usize> {
struct StreamRecord<IPType, SegmentNumRep, const SEGMENTS_COUNT: usize>
where
IPType: Debug,
SegmentNumRep: Debug,
{
start: VersionedIP<IPType, SegmentNumRep, SEGMENTS_COUNT>,
country_idx: usize,
}
Expand Down
168 changes: 55 additions & 113 deletions ip_country/src/country_finder.rs
Original file line number Diff line number Diff line change
@@ -1,28 +1,20 @@
use crate::country_block_serde::CountryBlockDeserializer;
use crate::country_block_stream::{Country, CountryBlock};
use crate::dbip_country;
use itertools::Itertools;
use lazy_static::lazy_static;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

#[cfg(not(test))]
lazy_static! {
pub static ref COUNTRY_CODE_FINDER: CountryCodeFinder = CountryCodeFinder::new(
crate::dbip_country::ipv4_country_data(),
crate::dbip_country::ipv6_country_data()
);
}

#[cfg(test)]
lazy_static! {
pub static ref COUNTRY_CODE_FINDER: CountryCodeFinder = CountryCodeFinder::new(
crate::test_dbip_country::ipv4_country_data(),
crate::test_dbip_country::ipv6_country_data()
dbip_country::ipv4_country_data(),
dbip_country::ipv6_country_data()
);
}

pub struct CountryCodeFinder {
pub ipv4: Vec<CountryBlock>,
pub ipv6: Vec<CountryBlock>,
ipv4: Vec<CountryBlock>,
ipv6: Vec<CountryBlock>,
}

impl CountryCodeFinder {
Expand Down Expand Up @@ -53,75 +45,26 @@ impl CountryCodeFinder {
_ => Some(country),
}
}

pub fn ensure_init(&self) {
//This should provoke lazy_static to perform the value initialization
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::country_block_serde::CountryBlockDeserializer;
use crate::country_block_serde::{
CountryBlockDeserializer, Ipv4CountryBlockDeserializer, Ipv6CountryBlockDeserializer,
};
use crate::dbip_country;
use lazy_static::lazy_static;
use std::str::FromStr;
use std::time::SystemTime;

lazy_static! {
static ref COUNTRY_CODE_FINDER_TEST: CountryCodeFinder =
CountryCodeFinder::new(ipv4_country_data(), ipv6_country_data());
static ref FULL_COUNTRY_CODE_FINDER: CountryCodeFinder = CountryCodeFinder::new(
crate::dbip_country::ipv4_country_data(),
crate::dbip_country::ipv6_country_data()
);
}

pub fn ipv4_country_data() -> (Vec<u64>, usize) {
(
vec![
0x0080000300801003,
0x82201C0902E01807,
0x28102E208388840B,
0x605C0100AB76020E,
0x0000000000000000,
],
271,
)
}

pub fn ipv6_country_data() -> (Vec<u64>, usize) {
(
vec![
0x3000040000400007,
0x00C0001400020000,
0xA80954B000000700,
0x4000000F0255604A,
0x0300004000040004,
0xE04AAC8380003800,
0x00018000A4000001,
0x2AB0003485C0001C,
0x0600089000000781,
0xC001D20700007000,
0x00424000001E04AA,
0x15485C0001C00018,
0xC90000007812AB00,
0x2388000700006002,
0x000001E04AAC00C5,
0xC0001C0001801924,
0x0007812AB0063485,
0x0070000600C89000,
0x1E04AAC049D23880,
0xC000180942400000,
0x12AB025549BA0001,
0x0040002580000078,
0xAC8B800038000300,
0x000000000001E04A,
],
1513,
)
}

#[test]
fn finds_ipv4_address_in_fourth_block() {
let result = CountryCodeFinder::find_country(
&COUNTRY_CODE_FINDER_TEST,
&COUNTRY_CODE_FINDER,
IpAddr::from_str("1.0.6.15").unwrap(),
);

Expand All @@ -130,17 +73,18 @@ mod tests {

#[test]
fn does_not_find_ipv4_address_in_zz_block() {
COUNTRY_CODE_FINDER.ensure_init();
let time_start = SystemTime::now();
let result = CountryCodeFinder::find_country(
&COUNTRY_CODE_FINDER_TEST,
&COUNTRY_CODE_FINDER,
IpAddr::from_str("0.0.5.0").unwrap(),
);
let time_end = SystemTime::now();

assert_eq!(result, None);
let duration = time_end.duration_since(time_start).unwrap();
assert!(
duration.as_secs() < 1,
duration.as_millis() <= 1,
"Duration of the search was too long: {} ms",
duration.as_millis()
);
Expand All @@ -149,17 +93,17 @@ mod tests {
#[test]
fn finds_ipv6_address_in_fourth_block() {
let result = CountryCodeFinder::find_country(
&COUNTRY_CODE_FINDER_TEST,
IpAddr::from_str("1:0:5:0:0:0:0:0").unwrap(),
&COUNTRY_CODE_FINDER,
IpAddr::from_str("2001:2::").unwrap(),
);

assert_eq!(result, Some(Country::try_from("AU").unwrap()))
assert_eq!(result, Some(Country::try_from("US").unwrap()))
}

#[test]
fn does_not_find_ipv6_address_in_zz_block() {
let result = CountryCodeFinder::find_country(
&COUNTRY_CODE_FINDER_TEST,
&COUNTRY_CODE_FINDER,
IpAddr::from_str("0:0:5:0:0:0:0:0").unwrap(),
);

Expand All @@ -168,11 +112,11 @@ mod tests {

#[test]
fn real_test_ipv4_with_google() {
let result = FULL_COUNTRY_CODE_FINDER
.find_country(
IpAddr::from_str("142.250.191.132").unwrap(), // dig www.google.com A
)
.unwrap();
let result = CountryCodeFinder::find_country(
&COUNTRY_CODE_FINDER,
IpAddr::from_str("142.250.191.132").unwrap(), // dig www.google.com A
)
.unwrap();

assert_eq!(result.free_world, true);
assert_eq!(result.iso3166, "US".to_string());
Expand All @@ -182,7 +126,7 @@ mod tests {
#[test]
fn real_test_ipv4_with_cz_ip() {
let result = CountryCodeFinder::find_country(
&FULL_COUNTRY_CODE_FINDER,
&COUNTRY_CODE_FINDER,
IpAddr::from_str("77.75.77.222").unwrap(), // dig www.seznam.cz A
)
.unwrap();
Expand All @@ -194,76 +138,77 @@ mod tests {

#[test]
fn real_test_ipv4_with_sk_ip() {
let _ = CountryCodeFinder::find_country(
&FULL_COUNTRY_CODE_FINDER,
IpAddr::from_str("213.81.185.100").unwrap(), // dig www.zoznam.sk A
)
.unwrap();
let time_start = SystemTime::now();

let result = CountryCodeFinder::find_country(
&FULL_COUNTRY_CODE_FINDER,
&COUNTRY_CODE_FINDER,
IpAddr::from_str("213.81.185.100").unwrap(), // dig www.zoznam.sk A
)
.unwrap();

let time_end = SystemTime::now();
let duration = time_end.duration_since(time_start).unwrap();

assert_eq!(result.free_world, true);
assert_eq!(result.iso3166, "SK".to_string());
assert_eq!(result.name, "Slovakia".to_string());
let duration = time_end.duration_since(time_start).unwrap();
assert!(
duration.as_secs() < 1,
"Duration of the search was too long: {} ms",
duration.as_millis() < 1,
"Duration of the search was too long: {} millisecond",
duration.as_millis()
);
}

#[test]
fn real_test_ipv6_with_google() {
let _ = CountryCodeFinder::find_country(
&FULL_COUNTRY_CODE_FINDER,
IpAddr::from_str("2607:f8b0:4009:814::2004").unwrap(), // dig www.google.com AAAA
)
.unwrap();
let time_start = SystemTime::now();

let result = CountryCodeFinder::find_country(
&FULL_COUNTRY_CODE_FINDER,
&COUNTRY_CODE_FINDER,
IpAddr::from_str("2607:f8b0:4009:814::2004").unwrap(), // dig www.google.com AAAA
)
.unwrap();

let time_end = SystemTime::now();
let duration = time_end.duration_since(time_start).unwrap();

assert_eq!(result.free_world, true);
assert_eq!(result.iso3166, "US".to_string());
assert_eq!(result.name, "United States".to_string());
let duration = time_end.duration_since(time_start).unwrap();
assert!(
duration.as_secs() < 1,
duration.as_millis() < 1,
"Duration of the search was too long: {} ms",
duration.as_millis()
);
}

#[test]
fn country_blocks_for_ipv4_and_ipv6_are_deserialized_filled_into_vecs() {
fn country_blocks_for_ipv4_and_ipv6_are_deserialized_and_inserted_into_vecs() {
let time_start = SystemTime::now();

let deserializer_ipv4 = CountryBlockDeserializer::<Ipv4Addr, u8, 4>::new(
crate::dbip_country::ipv4_country_data(),
);
let deserializer_ipv6 = CountryBlockDeserializer::<Ipv6Addr, u16, 8>::new(
crate::dbip_country::ipv6_country_data(),
);
let deserializer_ipv4 =
Ipv4CountryBlockDeserializer::new(dbip_country::ipv4_country_data());
let deserializer_ipv6 =
Ipv6CountryBlockDeserializer::new(dbip_country::ipv6_country_data());

let time_end = SystemTime::now();
let time_start_fill = SystemTime::now();
let _ = deserializer_ipv4.collect_vec();
let _ = deserializer_ipv6.collect_vec();

let country_block_finder_ipv4 = deserializer_ipv4.collect_vec();
let country_block_finder_ipv6 = deserializer_ipv6.collect_vec();

let time_end_fill = SystemTime::now();
let duration_deserialize = time_end.duration_since(time_start).unwrap();
let duration_fill = time_end_fill.duration_since(time_start_fill).unwrap();

assert_eq!(
country_block_finder_ipv4.len(),
dbip_country::ipv4_country_block_count()
);
assert_eq!(
country_block_finder_ipv6.len(),
dbip_country::ipv6_country_block_count()
);
assert!(
duration_deserialize.as_secs() < 15,
"Duration of the deserialization was too long: {} ms",
Expand All @@ -278,11 +223,8 @@ mod tests {

#[test]
fn check_ipv4_ipv6_country_blocks_length() {
let _result = FULL_COUNTRY_CODE_FINDER
.find_country(IpAddr::from_str("142.250.191.132").unwrap())
.unwrap();
let country_block_len_ipv4 = FULL_COUNTRY_CODE_FINDER.ipv4.len();
let country_block_len_ipv6 = FULL_COUNTRY_CODE_FINDER.ipv6.len();
let country_block_len_ipv4 = COUNTRY_CODE_FINDER.ipv4.len();
let country_block_len_ipv6 = COUNTRY_CODE_FINDER.ipv6.len();

assert_eq!(
country_block_len_ipv4,
Expand Down
Loading

0 comments on commit 5bc34df

Please sign in to comment.