diff --git a/src/db.rs b/src/db.rs index 75073a3..49448ca 100644 --- a/src/db.rs +++ b/src/db.rs @@ -7,22 +7,32 @@ use maxminddb::Reader; use serde::Deserialize; use crate::models::{ - IpInfo, MmdbAsn, MmdbCity, MmdbCountry, MmdbProxy, MmdbRecord, ProxyInfo, from_mmdb_record, + from_mmdb_record, IpInfo, MmdbAsn, MmdbCity, MmdbCountry, MmdbProxy, MmdbRecord, ProxyInfo, }; pub type DbReader = Reader>; -pub type SharedDb = Arc>; +pub type SharedDb = Arc>>; pub fn load_db(path: &Path) -> Result { Reader::open_readfile(path) } pub fn new_shared(reader: DbReader) -> SharedDb { - Arc::new(ArcSwap::from_pointee(reader)) + Arc::new(ArcSwap::from_pointee(Some(reader))) } -pub fn lookup(db: &DbReader, ip: IpAddr) -> Option { - let result = db.lookup(ip).ok()?; +pub fn new_empty() -> SharedDb { + Arc::new(ArcSwap::from_pointee(None)) +} + +pub fn is_ready(db: &SharedDb) -> bool { + db.load().is_some() +} + +pub fn lookup(db: &SharedDb, ip: IpAddr) -> Option { + let guard = db.load(); + let reader = Option::as_ref(&guard)?; + let result = reader.lookup(ip).ok()?; let record: MmdbRecord<'_> = result.decode().ok()??; Some(from_mmdb_record(ip, &record)) } @@ -54,8 +64,10 @@ struct CityOnly<'a> { city: MmdbCity<'a>, } -pub fn lookup_proxy(db: &DbReader, ip: IpAddr) -> Option { - let result = db.lookup(ip).ok()?; +pub fn lookup_proxy(db: &SharedDb, ip: IpAddr) -> Option { + let guard = db.load(); + let reader = Option::as_ref(&guard)?; + let result = reader.lookup(ip).ok()?; let record: ProxyOnly = result.decode().ok()??; Some(ProxyInfo { is_proxy: record.proxy.is_proxy, @@ -68,26 +80,34 @@ pub fn lookup_proxy(db: &DbReader, ip: IpAddr) -> Option { }) } -pub fn lookup_asn_number(db: &DbReader, ip: IpAddr) -> Option { - let result = db.lookup(ip).ok()?; +pub fn lookup_asn_number(db: &SharedDb, ip: IpAddr) -> Option { + let guard = db.load(); + let reader = Option::as_ref(&guard)?; + let result = reader.lookup(ip).ok()?; let record: AsnOnly<'_> = result.decode().ok()??; record.asn.autonomous_system_number } -pub fn lookup_asn_org(db: &DbReader, ip: IpAddr) -> Option { - let result = db.lookup(ip).ok()?; +pub fn lookup_asn_org(db: &SharedDb, ip: IpAddr) -> Option { + let guard = db.load(); + let reader = Option::as_ref(&guard)?; + let result = reader.lookup(ip).ok()?; let record: AsnOnly<'_> = result.decode().ok()??; record.asn.autonomous_system_organization.map(str::to_owned) } -pub fn lookup_country_name(db: &DbReader, ip: IpAddr) -> Option { - let result = db.lookup(ip).ok()?; +pub fn lookup_country_name(db: &SharedDb, ip: IpAddr) -> Option { + let guard = db.load(); + let reader = Option::as_ref(&guard)?; + let result = reader.lookup(ip).ok()?; let record: CountryOnly<'_> = result.decode().ok()??; record.country.names.english.map(str::to_owned) } -pub fn lookup_city_name(db: &DbReader, ip: IpAddr) -> Option { - let result = db.lookup(ip).ok()?; +pub fn lookup_city_name(db: &SharedDb, ip: IpAddr) -> Option { + let guard = db.load(); + let reader = Option::as_ref(&guard)?; + let result = reader.lookup(ip).ok()?; let record: CityOnly<'_> = result.decode().ok()??; record.city.names.english.map(str::to_owned) } @@ -106,6 +126,10 @@ mod tests { load_db(Path::new(&test_db_path())).expect("failed to load test DB") } + fn test_shared_db() -> SharedDb { + new_shared(test_reader()) + } + #[test] fn load_db_with_valid_path() { let reader = load_db(Path::new(&test_db_path())); @@ -122,18 +146,25 @@ mod tests { fn new_shared_creates_valid_shared_db() { let reader = test_reader(); let shared = new_shared(reader); - let loaded = shared.load(); - // Verify the shared DB is functional by doing a simple lookup + assert!(is_ready(&shared)); + let guard = shared.load(); + let reader = Option::as_ref(&guard).unwrap(); let ip: IpAddr = "1.1.1.1".parse().unwrap(); - let result = loaded.lookup(ip); + let result = reader.lookup(ip); assert!(result.is_ok()); } + #[test] + fn new_empty_creates_unready_shared_db() { + let shared = new_empty(); + assert!(!is_ready(&shared)); + } + #[test] fn lookup_returns_some_for_known_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "1.1.1.1".parse().unwrap(); - let result = lookup(&reader, ip); + let result = lookup(&shared, ip); assert!(result.is_some()); let info = result.unwrap(); assert_eq!(info.ip, "1.1.1.1"); @@ -141,17 +172,24 @@ mod tests { #[test] fn lookup_returns_none_for_unroutable_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "0.0.0.0".parse().unwrap(); - let result = lookup(&reader, ip); + let result = lookup(&shared, ip); assert!(result.is_none()); } + #[test] + fn lookup_returns_none_when_db_not_loaded() { + let shared = new_empty(); + let ip: IpAddr = "1.1.1.1".parse().unwrap(); + assert!(lookup(&shared, ip).is_none()); + } + #[test] fn lookup_proxy_returns_data_for_known_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "45.77.77.77".parse().unwrap(); - let proxy = lookup_proxy(&reader, ip); + let proxy = lookup_proxy(&shared, ip); assert!(proxy.is_some()); let proxy = proxy.unwrap(); assert!(proxy.is_proxy); @@ -160,82 +198,82 @@ mod tests { #[test] fn lookup_proxy_returns_none_for_unroutable_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "0.0.0.0".parse().unwrap(); - let proxy = lookup_proxy(&reader, ip); + let proxy = lookup_proxy(&shared, ip); assert!(proxy.is_none()); } #[test] fn lookup_asn_number_for_known_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "1.1.1.1".parse().unwrap(); - let asn = lookup_asn_number(&reader, ip); + let asn = lookup_asn_number(&shared, ip); assert_eq!(asn, Some(13335)); } #[test] fn lookup_asn_number_returns_none_for_unroutable_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "0.0.0.0".parse().unwrap(); - let asn = lookup_asn_number(&reader, ip); + let asn = lookup_asn_number(&shared, ip); assert!(asn.is_none()); } #[test] fn lookup_asn_org_for_known_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "1.1.1.1".parse().unwrap(); - let org = lookup_asn_org(&reader, ip); + let org = lookup_asn_org(&shared, ip); assert!(org.is_some()); assert_eq!(org.unwrap(), "Cloudflare, Inc."); } #[test] fn lookup_asn_org_returns_none_for_unroutable_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "0.0.0.0".parse().unwrap(); - let org = lookup_asn_org(&reader, ip); + let org = lookup_asn_org(&shared, ip); assert!(org.is_none()); } #[test] fn lookup_country_name_for_known_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "45.77.77.77".parse().unwrap(); - let country = lookup_country_name(&reader, ip); + let country = lookup_country_name(&shared, ip); assert_eq!(country.as_deref(), Some("United States")); } #[test] fn lookup_country_name_returns_none_for_unroutable_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "0.0.0.0".parse().unwrap(); - let country = lookup_country_name(&reader, ip); + let country = lookup_country_name(&shared, ip); assert!(country.is_none()); } #[test] fn lookup_city_name_for_known_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "45.77.77.77".parse().unwrap(); - let city = lookup_city_name(&reader, ip); + let city = lookup_city_name(&shared, ip); assert_eq!(city.as_deref(), Some("Piscataway")); } #[test] fn lookup_city_name_returns_none_for_unroutable_ip() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "0.0.0.0".parse().unwrap(); - let city = lookup_city_name(&reader, ip); + let city = lookup_city_name(&shared, ip); assert!(city.is_none()); } #[test] fn lookup_works_with_ipv6() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "2606:4700:4700::1111".parse().unwrap(); - let result = lookup(&reader, ip); + let result = lookup(&shared, ip); assert!(result.is_some()); let info = result.unwrap(); assert_eq!(info.ip, "2606:4700:4700::1111"); @@ -248,10 +286,10 @@ mod tests { #[test] fn partial_lookups_work_with_ipv6() { - let reader = test_reader(); + let shared = test_shared_db(); let ip: IpAddr = "2606:4700:4700::1111".parse().unwrap(); - assert_eq!(lookup_asn_number(&reader, ip), Some(13335)); - assert!(lookup_asn_org(&reader, ip).is_some()); - assert!(lookup_proxy(&reader, ip).is_some()); + assert_eq!(lookup_asn_number(&shared, ip), Some(13335)); + assert!(lookup_asn_org(&shared, ip).is_some()); + assert!(lookup_proxy(&shared, ip).is_some()); } } diff --git a/src/error.rs b/src/error.rs index 2b8c5d2..5fb6b1c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -6,6 +6,7 @@ pub enum AppError { InvalidIp, NonPublicIp, MissingClientIp, + DbNotReady, DbLookupFailed, TemplateRenderFailed, } @@ -16,6 +17,7 @@ impl AppError { Self::IpNotFound => StatusCode::NOT_FOUND, Self::InvalidIp | Self::NonPublicIp => StatusCode::BAD_REQUEST, Self::MissingClientIp => StatusCode::BAD_REQUEST, + Self::DbNotReady => StatusCode::SERVICE_UNAVAILABLE, Self::DbLookupFailed => StatusCode::INTERNAL_SERVER_ERROR, Self::TemplateRenderFailed => StatusCode::INTERNAL_SERVER_ERROR, } @@ -27,6 +29,7 @@ impl AppError { Self::InvalidIp => "Invalid IP address", Self::NonPublicIp => "Only public IP addresses can be queried", Self::MissingClientIp => "Missing client IP address", + Self::DbNotReady => "Service is starting up, please try again shortly", Self::DbLookupFailed => "Database lookup failed", Self::TemplateRenderFailed => "Template render failed", } @@ -96,6 +99,16 @@ mod tests { .await; } + #[tokio::test] + async fn db_not_ready_returns_503() { + check_error_response( + AppError::DbNotReady, + StatusCode::SERVICE_UNAVAILABLE, + "Service is starting up, please try again shortly", + ) + .await; + } + #[tokio::test] async fn db_lookup_failed_returns_500() { check_error_response( diff --git a/src/handlers.rs b/src/handlers.rs index 6c5e1fa..29ff8f6 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -177,8 +177,10 @@ fn extract_client_ip(headers: &HeaderMap, dev_mode: bool) -> Result Result { - let reader = db.load(); - db::lookup(&reader, ip).ok_or(AppError::DbLookupFailed) + if !db::is_ready(db) { + return Err(AppError::DbNotReady); + } + db::lookup(db, ip).ok_or(AppError::DbLookupFailed) } fn wants_json(headers: &HeaderMap) -> bool { @@ -675,9 +677,11 @@ pub async fn asn_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let text = db::lookup_asn_number(&reader, ip) + let text = db::lookup_asn_number(&state.db, ip) .map(|n| format!("AS{n}")) .unwrap_or_default(); Ok(format!("{text}\n")) @@ -688,9 +692,11 @@ pub async fn org_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let org = db::lookup_asn_org(&reader, ip).unwrap_or_default(); + let org = db::lookup_asn_org(&state.db, ip).unwrap_or_default(); Ok(format!("{org}\n")) } @@ -699,9 +705,11 @@ pub async fn country_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let country = db::lookup_country_name(&reader, ip).unwrap_or_default(); + let country = db::lookup_country_name(&state.db, ip).unwrap_or_default(); Ok(format!("{country}\n")) } @@ -710,9 +718,11 @@ pub async fn city_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let city = db::lookup_city_name(&reader, ip).unwrap_or_default(); + let city = db::lookup_city_name(&state.db, ip).unwrap_or_default(); Ok(format!("{city}\n")) } @@ -721,9 +731,11 @@ pub async fn proxy_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?; + let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?; Ok(format!("{}\n", proxy.is_proxy)) } @@ -732,9 +744,11 @@ pub async fn vpn_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?; + let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?; Ok(format!("{}\n", proxy.is_vpn)) } @@ -743,9 +757,11 @@ pub async fn hosting_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?; + let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?; Ok(format!("{}\n", proxy.is_hosting)) } @@ -754,9 +770,11 @@ pub async fn tor_handler( headers: HeaderMap, Query(query): Query, ) -> Result { + if !db::is_ready(&state.db) { + return Err(AppError::DbNotReady); + } let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let reader = state.db.load(); - let proxy = db::lookup_proxy(&reader, ip).ok_or(AppError::DbLookupFailed)?; + let proxy = db::lookup_proxy(&state.db, ip).ok_or(AppError::DbLookupFailed)?; Ok(format!("{}\n", proxy.is_tor)) } diff --git a/src/main.rs b/src/main.rs index 9b19aa2..8783617 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,35 +19,30 @@ async fn main() { let config = config::Config::from_env(); let db_path = PathBuf::from(&config.db_path); - if let Err(e) = updater::ensure_db_exists(&db_path, &config.db_update_url).await { - error!("failed to initialize database: {e}"); - std::process::exit(1); - } - - let reader = match db::load_db(&db_path) { - Ok(r) => r, - Err(e) => { - error!("failed to load database: {e}"); - std::process::exit(1); - } + // Create shared DB in an empty (not-yet-loaded) state so the server can + // start immediately and respond to health checks while the database is + // still being downloaded on first deploy. + let shared_db = if db_path.exists() { + let reader = match db::load_db(&db_path) { + Ok(r) => r, + Err(e) => { + error!("failed to load database: {e}"); + std::process::exit(1); + } + }; + db::new_shared(reader) + } else { + info!("database file not found, server will start while download proceeds"); + db::new_empty() }; - let shared_db = db::new_shared(reader); let cancel = CancellationToken::new(); - let updater_handle = updater::spawn_updater( - shared_db.clone(), - db_path, - config.db_update_url, - config.db_update_interval_hours, - cancel.clone(), - ); - let site_domain: std::sync::Arc = config.site_domain.into(); let openapi_json = handlers::build_openapi_json(&site_domain); let state = handlers::AppState { - db: shared_db, + db: shared_db.clone(), site_domain, dev_mode: config.dev_mode, openapi_json, @@ -59,6 +54,33 @@ async fn main() { .await .expect("failed to bind listener"); + // Spawn background database initialization and periodic updater. + let init_db = shared_db.clone(); + let init_path = db_path.clone(); + let init_url = config.db_update_url.clone(); + let init_cancel = cancel.clone(); + let updater_handle = tokio::spawn(async move { + // Only download if the DB was not loaded synchronously above. + if !db::is_ready(&init_db) { + if let Err(e) = updater::init_db(&init_db, &init_path, &init_url).await { + error!("failed to initialize database: {e}"); + std::process::exit(1); + } + } + + // Start periodic updater after the initial load succeeds. + updater::run_updater( + init_db, + init_path, + init_url, + config.db_update_interval_hours, + init_cancel, + ) + .await; + }); + + info!("listening on {}", config.listen_addr); + axum::serve(listener, app) .with_graceful_shutdown(shutdown_signal()) .await diff --git a/src/updater.rs b/src/updater.rs index d07084f..e0a0729 100644 --- a/src/updater.rs +++ b/src/updater.rs @@ -10,6 +10,9 @@ use tracing::{error, info}; use crate::db::{DbReader, SharedDb, load_db}; +#[cfg(test)] +use crate::db::new_shared; + pub async fn ensure_db_exists( db_path: &Path, update_url: &str, @@ -26,6 +29,21 @@ pub async fn ensure_db_exists( Ok(()) } +/// Ensures the database file exists (downloading if needed), loads it, and +/// stores the reader into the shared DB. Called during startup to initialize +/// the database before the application can serve requests. +pub async fn init_db( + shared_db: &SharedDb, + db_path: &Path, + update_url: &str, +) -> Result<(), Box> { + ensure_db_exists(db_path, update_url).await?; + let reader = load_db(db_path)?; + shared_db.store(Arc::new(Some(reader))); + info!("database loaded successfully"); + Ok(()) +} + /// Downloads the database to a temp file, validates it, and atomically /// renames it into place. Returns the validated reader for immediate use. /// Cleans up the temp file on any failure. @@ -86,21 +104,37 @@ pub fn spawn_updater( interval_hours: u64, cancel: CancellationToken, ) -> JoinHandle<()> { - tokio::spawn(async move { - let interval = Duration::from_secs(interval_hours * 3600); - loop { - tokio::select! { - () = tokio::time::sleep(interval) => {} - () = cancel.cancelled() => { - info!("updater task cancelled, stopping"); - break; - } - } - if let Err(e) = update_db(&shared_db, &db_path, &update_url).await { - error!("database update failed: {e}"); + tokio::spawn(run_updater( + shared_db, + db_path, + update_url, + interval_hours, + cancel, + )) +} + +/// Runs the periodic database updater loop. This is the inner async function +/// that can be called directly when already inside a spawned task. +pub async fn run_updater( + shared_db: SharedDb, + db_path: PathBuf, + update_url: String, + interval_hours: u64, + cancel: CancellationToken, +) { + let interval = Duration::from_secs(interval_hours * 3600); + loop { + tokio::select! { + () = tokio::time::sleep(interval) => {} + () = cancel.cancelled() => { + info!("updater task cancelled, stopping"); + break; } } - }) + if let Err(e) = update_db(&shared_db, &db_path, &update_url).await { + error!("database update failed: {e}"); + } + } } async fn update_db( @@ -109,7 +143,7 @@ async fn update_db( update_url: &str, ) -> Result<(), Box> { let reader = download_db(update_url, db_path).await?; - shared_db.store(Arc::new(reader)); + shared_db.store(Arc::new(Some(reader))); info!("database updated successfully"); Ok(()) } @@ -210,7 +244,7 @@ mod tests { #[tokio::test] async fn spawn_updater_stops_on_cancellation() { let reader = load_db(&test_db_path()).unwrap(); - let shared_db = crate::db::new_shared(reader); + let shared_db = new_shared(reader); let cancel = CancellationToken::new(); let handle = spawn_updater( @@ -239,16 +273,17 @@ mod tests { std::fs::copy(test_db_path(), &dest).unwrap(); let original_reader = load_db(&dest).unwrap(); - let shared_db = crate::db::new_shared(original_reader); + let shared_db = new_shared(original_reader); // Load a new reader and manually swap it to verify the mechanism let new_reader = load_db(&dest).unwrap(); - shared_db.store(Arc::new(new_reader)); + shared_db.store(Arc::new(Some(new_reader))); // Verify the shared DB still works after the swap let guard = shared_db.load(); + let reader = Option::as_ref(&guard).unwrap(); let ip: IpAddr = "1.1.1.1".parse().unwrap(); - let result = guard.lookup(ip); + let result = reader.lookup(ip); assert!(result.is_ok()); } }