From c4fa53fa4014468a6d391432b3e2922f4ae8509b Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Thu, 15 Jul 2021 02:14:04 -0400 Subject: [PATCH] Add geo ip logging support --- .gitignore | 1 + README.md | 5 ++ src/config.rs | 6 +++ src/main.rs | 28 +++++++++++ src/metrics.rs | 132 ++++++++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 170 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 6fea2b0..0f19b72 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ perf.data* dhat.out.* settings.yaml tarpaulin-report.html +GeoLite2-Country.mmdb \ No newline at end of file diff --git a/README.md b/README.md index 570662d..43dd43a 100644 --- a/README.md +++ b/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. \ No newline at end of file diff --git a/src/config.rs b/src/config.rs index e5293c3..5c4fa2b 100644 --- a/src/config.rs +++ b/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]") diff --git a/src/main.rs b/src/main.rs index 3441c15..3f84c7f 100644 --- a/src/main.rs +++ b/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> { 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> { 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::(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) diff --git a/src/metrics.rs b/src/metrics.rs index def3f11..2ff2288 100644 --- a/src/metrics.rs +++ b/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>> = OnceCell::new(); + +static COUNTRY_VISIT_COUNTER: Lazy = 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) { + 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(); +}