Add geo ip logging support

.gitignore

@@ -6,3 +6,4 @@ perf.data*
 dhat.out.*
 settings.yaml
 tarpaulin-report.html
+GeoLite2-Country.mmdb

README.md

@@ -75,3 +75,8 @@ following:
 #### Beta testers
 
 - NigelVH#7162
+
+---
+
+If using the geo IP logging feature, then this product includes GeoLite2 data
+created by MaxMind, available from https://www.maxmind.com.

src/config.rs

@@ -200,6 +200,12 @@ impl Config {
 #[derive(Deserialize, Serialize, Clone)]
 pub struct ClientSecret(String);
 
+impl ClientSecret {
+    pub fn as_str(&self) -> &str {
+        self.0.as_ref()
+    }
+}
+
 impl std::fmt::Debug for ClientSecret {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "[client secret]")

src/main.rs

@@ -5,17 +5,21 @@
 use std::env::VarError;
 use std::error::Error;
 use std::fmt::Display;
+use std::net::SocketAddr;
 use std::num::ParseIntError;
+use std::str::FromStr;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
+use actix_web::dev::Service;
 use actix_web::rt::{spawn, time, System};
 use actix_web::web::{self, Data};
 use actix_web::{App, HttpResponse, HttpServer};
 use cache::{Cache, DiskCache};
 use chacha20::Key;
 use config::Config;
+use maxminddb::geoip2;
 use parking_lot::RwLock;
 use rustls::{NoClientAuth, ServerConfig};
 use sodiumoxide::crypto::stream::xchacha20::gen_key;
@@ -27,6 +31,7 @@ use tracing::{debug, error, info, warn};
 use crate::cache::mem::{Lfu, Lru};
 use crate::cache::{MemoryCache, ENCRYPTION_KEY};
 use crate::config::{CacheType, UnstableOptions, OFFLINE_MODE};
+use crate::metrics::{record_country_visit, GEOIP_DATABASE};
 use crate::state::DynamicServerCert;
 
 mod cache;
@@ -107,6 +112,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
         metrics::init();
     }
 
+    if let Some(key) = config.geoip_license_key.clone() {
+        if let Err(e) = metrics::load_geo_ip_data(key).await {
+            error!("Failed to initialize geo ip db: {}", e);
+        }
+    }
+
     // HTTP Server init
 
     let server = if OFFLINE_MODE.load(Ordering::Acquire) {
@@ -171,6 +182,23 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let server = HttpServer::new(move || {
         App::new()
             .wrap(actix_web::middleware::Compress::default())
+            .wrap_fn(|req, srv| {
+                if let Some(reader) = GEOIP_DATABASE.get() {
+                    let maybe_country = req
+                        .connection_info()
+                        .realip_remote_addr()
+                        .map(SocketAddr::from_str)
+                        .and_then(Result::ok)
+                        .as_ref()
+                        .map(SocketAddr::ip)
+                        .map(|ip| reader.lookup::<geoip2::Country>(ip))
+                        .and_then(Result::ok);
+
+                    record_country_visit(maybe_country);
+                }
+
+                srv.call(req)
+            })
             .service(routes::index)
             .service(routes::token_data)
             .service(routes::token_data_saver)

src/metrics.rs

@@ -1,5 +1,28 @@
-use once_cell::sync::Lazy;
-use prometheus::{register_int_counter, IntCounter};
+use std::fs::metadata;
+use std::hint::unreachable_unchecked;
+use std::time::SystemTime;
+
+use chrono::Duration;
+use flate2::read::GzDecoder;
+use maxminddb::geoip2::Country;
+use once_cell::sync::{Lazy, OnceCell};
+use prometheus::{register_int_counter, register_int_counter_vec, IntCounter, IntCounterVec};
+use tar::Archive;
+use thiserror::Error;
+use tracing::{debug, field::debug, info, warn};
+
+use crate::config::ClientSecret;
+
+pub static GEOIP_DATABASE: OnceCell<maxminddb::Reader<Vec<u8>>> = OnceCell::new();
+
+static COUNTRY_VISIT_COUNTER: Lazy<IntCounterVec> = Lazy::new(|| {
+    register_int_counter_vec!(
+        "country_visits_total",
+        "The number of visits from a country",
+        &["country"]
+    )
+    .unwrap()
+});
 
 macro_rules! init_counters {
     ($(($counter:ident, $ty:ty, $name:literal, $desc:literal),)*) => {
@@ -11,7 +34,11 @@ macro_rules! init_counters {
 
         #[allow(clippy::shadow_unrelated)]
         pub fn init() {
+            // These need to be called at least once, otherwise the macro never
+            // called and thus the metrics don't get logged
             $(let _a = $counter.get();)*
+
+            init_other();
         }
     };
 }
@@ -54,3 +81,104 @@ init_counters!(
         "The total number of request not served by primary endpoints."
     ),
 );
+
+// initialization for any other counters that aren't simple int counters
+fn init_other() {
+    let _a = COUNTRY_VISIT_COUNTER.local();
+}
+
+#[derive(Error, Debug)]
+pub enum DbLoadError {
+    #[error(transparent)]
+    Reqwest(#[from] reqwest::Error),
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+    #[error(transparent)]
+    MaxMindDb(#[from] maxminddb::MaxMindDBError),
+}
+
+pub async fn load_geo_ip_data(license_key: ClientSecret) -> Result<(), DbLoadError> {
+    const DB_PATH: &str = "./GeoLite2-Country.mmdb";
+
+    // Check date of db
+    let db_date_created = metadata(DB_PATH)
+        .ok()
+        .and_then(|metadata| match metadata.created() {
+            Ok(time) => Some(time),
+            Err(_) => {
+                debug("fs didn't report birth time, fall back to last modified instead");
+                metadata.modified().ok()
+            }
+        })
+        .unwrap_or(SystemTime::UNIX_EPOCH);
+    let duration = match SystemTime::now().duration_since(dbg!(db_date_created)) {
+        Ok(time) => Duration::from_std(time).expect("duration to fit"),
+        Err(_) => {
+            warn!("Clock may have gone backwards?");
+            Duration::max_value()
+        }
+    };
+
+    // DB expired, fetch a new one
+    if duration > Duration::weeks(1) {
+        fetch_db(license_key).await?;
+    } else {
+        info!("Geo IP database isn't old enough, not updating.");
+    }
+
+    // Result literally cannot panic here, buuuuuut if it does we'll panic
+    GEOIP_DATABASE
+        .set(maxminddb::Reader::open_readfile(DB_PATH)?)
+        .map_err(|_| ())
+        .expect("to set the geo ip db singleton");
+
+    Ok(())
+}
+
+async fn fetch_db(license_key: ClientSecret) -> Result<(), DbLoadError> {
+    let resp = reqwest::get(format!("https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country&license_key={}&suffix=tar.gz", license_key.as_str()))
+        .await?
+        .bytes()
+        .await?;
+    let mut decoder = Archive::new(GzDecoder::new(resp.as_ref()));
+    let mut decoded_paths: Vec<_> = decoder
+        .entries()?
+        .filter_map(Result::ok)
+        .filter_map(|mut entry| {
+            let path = entry.path().ok()?.to_path_buf();
+            let file_name = path.file_name()?;
+            if file_name != "GeoLite2-Country.mmdb" {
+                return None;
+            }
+            entry.unpack(file_name).ok()?;
+            Some(path)
+        })
+        .collect();
+
+    assert_eq!(decoded_paths.len(), 1);
+
+    let path = match decoded_paths.pop() {
+        Some(path) => path,
+        None => unsafe { unreachable_unchecked() },
+    };
+
+    debug!("Extracted {}", path.as_path().to_string_lossy());
+
+    Ok(())
+}
+
+pub fn record_country_visit(country: Option<Country>) {
+    let iso_code = if let Some(country) = country {
+        country
+            .country
+            .and_then(|c| c.iso_code)
+            .unwrap_or("unknown")
+    } else {
+        "unknown"
+    };
+
+    COUNTRY_VISIT_COUNTER
+        .get_metric_with_label_values(&[iso_code])
+        .unwrap()
+        .inc();
+}