headless_lms_utils/
ip_to_country.rs

1use std::{collections::HashMap, io::Read, net::IpAddr, path::Path, time::Instant};
2
3use std::env;
4
5use anyhow;
6use anyhow::bail;
7use flate2::read::GzDecoder;
8use ipnet::IpNet;
9use walkdir::WalkDir;
10
11use crate::ApplicationConfiguration;
12
13pub struct IpToCountryMapper {
14    lists: HashMap<String, Vec<ipnet::IpNet>>,
15}
16
17impl IpToCountryMapper {
18    /**
19    Creates a new mapper by reading the folder from env IP_TO_COUNTRY_MAPPING_DIRECTORY
20
21    Expensive and synchronous.
22    */
23    pub fn new(app_conf: &ApplicationConfiguration) -> anyhow::Result<Self> {
24        let mut lists = HashMap::new();
25        if let Ok(ip_to_country_mapping_directory) = env::var("IP_TO_COUNTRY_MAPPING_DIRECTORY") {
26            info!("Loading country to ip mapping");
27            let start = Instant::now();
28            let path = Path::new(&ip_to_country_mapping_directory);
29            if !path.exists() {
30                bail!("The folder specified in IP_TO_COUNTRY_MAPPING_DIRECTORY does not exist.");
31            }
32            let walker = WalkDir::new(path)
33                .follow_links(false)
34                .max_open(10)
35                .contents_first(false)
36                .sort_by_file_name();
37            let mut some_lists_skipped = true;
38            for entry_result in walker {
39                let entry = entry_result?;
40                if !entry.file_type().is_file() {
41                    continue;
42                }
43                let file_name = entry.file_name().to_string_lossy();
44                let mut parts = if file_name.contains("v4") {
45                    file_name.split("v4")
46                } else {
47                    file_name.split("v6")
48                };
49
50                // Speed up loading only in bin/dev
51                if app_conf.test_mode
52                    && cfg!(debug_assertions)
53                    && lists.len() > 10
54                    && !file_name.to_lowercase().contains("fi")
55                {
56                    some_lists_skipped = true;
57                    continue;
58                }
59
60                if let Some(country_code) = parts.next() {
61                    let list = lists
62                        .entry(country_code.to_lowercase())
63                        .or_insert_with(Vec::new);
64                    let path = entry.path();
65                    let bytes = std::fs::read(path)?;
66                    // The file is gzipped in the dockerfile to save space.
67                    let mut gz = GzDecoder::new(&bytes[..]);
68                    let mut contents = String::new();
69                    gz.read_to_string(&mut contents)?;
70                    for line in contents.trim().lines() {
71                        if let Ok(ipnet) = line.parse::<IpNet>() {
72                            list.push(ipnet);
73
74                            if cfg!(debug_assertions) && (list.len() > 10) {
75                                some_lists_skipped = true;
76                                break;
77                            }
78                        }
79                    }
80                }
81            }
82            info!(
83                elapsed_time = ?start.elapsed(),
84                "Loaded ip to country mappings"
85            );
86            if some_lists_skipped {
87                warn!(
88                    "Some ip to country lists were skipped to speed up loading. This should not happen in production."
89                );
90            }
91        } else {
92            warn!(
93                "IP_TO_COUNTRY_MAPPING_DIRECTORY not specified, not loading ip to country mappings."
94            );
95            // Not failing to allow running the backend without the lists
96        }
97
98        Ok(Self { lists })
99    }
100
101    pub fn map_ip_to_country(&self, ip: &IpAddr) -> Option<&str> {
102        for (country, list) in &self.lists {
103            if list.iter().any(|ipnet| ipnet.contains(ip)) {
104                info!("Mapped ip {} to country {}.", ip, country);
105                return Some(country);
106            }
107        }
108        info!("Ip {} did not resolve to a country.", ip);
109        None
110    }
111}
headless_lms_utils/ip_to_country.rs

headless_lms_utils/
ip_to_country.rs
