From 3298ba56eb11f92c9ed21c09cf87f97866397274 Mon Sep 17 00:00:00 2001 From: NetworkCats Date: Mon, 25 May 2026 18:55:23 +0800 Subject: [PATCH] Optimize MMDB response paths --- src/db.rs | 107 ++++++++++++++++++- src/handlers/api.rs | 19 ++-- src/handlers/html.rs | 78 +++++--------- src/handlers/mod.rs | 17 ++- src/models/api.rs | 12 +++ src/models/mod.rs | 244 ++++++++++++++++++++++++++++++++++++++++++- 6 files changed, 399 insertions(+), 78 deletions(-) diff --git a/src/db.rs b/src/db.rs index ace795d..9021db8 100644 --- a/src/db.rs +++ b/src/db.rs @@ -8,7 +8,10 @@ use arc_swap::ArcSwap; use maxminddb::{MaxMindDbError, PathElement, Reader}; use serde::Deserialize; -use crate::models::{IpInfo, MmdbProxy, MmdbRecord, ProxyInfo, from_mmdb_record}; +use crate::models::{ + IpInfo, MmdbProxy, MmdbRecord, PageInfo, ProxyInfo, from_mmdb_record, + serialize_mmdb_record_json, +}; pub type DbReader = Reader>; pub type SharedDb = Arc>>; @@ -39,6 +42,15 @@ pub fn lookup(db: &SharedDb, ip: IpAddr) -> LookupOutcome { }) } +pub fn lookup_json(db: &SharedDb, ip: IpAddr) -> LookupOutcome> { + with_reader(db, ip, |result| { + let Some(record) = result.decode::>()? else { + return Ok(None); + }; + Ok(Some(serialize_mmdb_record_json(ip, &record))) + }) +} + /// Loads the shared reader, performs a lookup, and passes the raw lookup result /// to `f` for decoding. This eliminates the repeated load-guard-lookup boilerplate /// while preserving lookup/decode errors for callers that need HTTP error mapping. @@ -71,6 +83,18 @@ struct ProxyOnly { proxy: MmdbProxy, } +#[derive(Deserialize, Default)] +struct PageRecord<'a> { + #[serde(borrow, default)] + city: crate::models::MmdbCity<'a>, + #[serde(default)] + country: crate::models::MmdbCountry<'a>, + #[serde(default)] + asn: crate::models::MmdbAsn<'a>, + #[serde(default)] + proxy: MmdbProxy, +} + pub fn lookup_proxy(db: &SharedDb, ip: IpAddr) -> LookupOutcome { with_reader(db, ip, |result| { let Some(record) = result.decode::()? else { @@ -80,6 +104,46 @@ pub fn lookup_proxy(db: &SharedDb, ip: IpAddr) -> LookupOutcome { }) } +pub fn lookup_page_info(db: &SharedDb, ip: IpAddr, mmdb_key: &str) -> LookupOutcome { + with_reader(db, ip, |result| { + let Some(record) = result.decode::>()? else { + return Ok(None); + }; + + let asn_number = record.asn.autonomous_system_number; + let asn = asn_number.map(|n| format!("AS{n}")).unwrap_or_default(); + + Ok(Some(PageInfo { + ip: ip.to_string(), + asn, + asn_number, + org: record + .asn + .autonomous_system_organization + .unwrap_or_default() + .to_owned(), + country: localized_mmdb_name(&record.country.names, mmdb_key).to_owned(), + country_code: record.country.iso_code.unwrap_or_default().to_owned(), + city: localized_mmdb_name(&record.city.names, mmdb_key).to_owned(), + proxy: record.proxy.into(), + })) + }) +} + +fn localized_mmdb_name<'a>(names: &'a crate::models::MmdbNames<'a>, mmdb_key: &str) -> &'a str { + match mmdb_key { + "de" => names.german.or(names.english), + "es" => names.spanish.or(names.english), + "fr" => names.french.or(names.english), + "ja" => names.japanese.or(names.english), + "pt-BR" => names.brazilian_portuguese.or(names.english), + "ru" => names.russian.or(names.english), + "zh-CN" => names.simplified_chinese.or(names.english), + _ => names.english, + } + .unwrap_or_default() +} + fn lookup_proxy_flag(db: &SharedDb, ip: IpAddr, field: &'static str) -> LookupOutcome { lookup_path( db, @@ -223,6 +287,47 @@ mod tests { assert!(lookup(&shared, ip).unwrap().is_none()); } + #[test] + fn lookup_json_matches_owned_ip_info_for_ipv4() { + assert_json_lookup_matches_ip_info("45.77.77.77"); + } + + #[test] + fn lookup_json_matches_owned_ip_info_for_ipv6() { + assert_json_lookup_matches_ip_info("2606:4700:4700::1111"); + } + + fn assert_json_lookup_matches_ip_info(ip: &str) { + let shared = test_shared_db(); + let ip: IpAddr = ip.parse().unwrap(); + let info = lookup(&shared, ip).unwrap().unwrap(); + let json = lookup_json(&shared, ip).unwrap().unwrap(); + let expected = serde_json::to_value(info).unwrap(); + let actual: serde_json::Value = serde_json::from_slice(&json).unwrap(); + assert_eq!(actual, expected); + } + + #[test] + fn lookup_json_returns_none_when_db_not_loaded() { + let shared = new_empty(); + let ip: IpAddr = "1.1.1.1".parse().unwrap(); + assert!(lookup_json(&shared, ip).unwrap().is_none()); + } + + #[test] + fn lookup_page_info_returns_rendered_fields() { + let shared = test_shared_db(); + let ip: IpAddr = "45.77.77.77".parse().unwrap(); + let info = lookup_page_info(&shared, ip, "en").unwrap().unwrap(); + assert_eq!(info.ip, "45.77.77.77"); + assert_eq!(info.asn, "AS20473"); + assert_eq!(info.asn_number, Some(20473)); + assert_eq!(info.org, "The Constant Company, LLC"); + assert_eq!(info.country, "United States"); + assert_eq!(info.country_code, "US"); + assert_eq!(info.city, "Piscataway"); + } + #[test] fn lookup_proxy_returns_none_for_unroutable_ip() { let shared = test_shared_db(); diff --git a/src/handlers/api.rs b/src/handlers/api.rs index 5980bd5..0388b16 100644 --- a/src/handlers/api.rs +++ b/src/handlers/api.rs @@ -6,9 +6,8 @@ use axum::response::{IntoResponse, Response}; use crate::db::{self, SharedDb}; use crate::error::AppError; -use crate::models::IpInfo; -use super::{AppState, IpQuery, extract_ip_from_header, lookup_ip, resolve_ip}; +use super::{AppState, IpQuery, extract_ip_from_header, lookup_json_response, resolve_ip}; pub async fn health(State(state): State) -> StatusCode { if db::is_ready(&state.db) { @@ -46,10 +45,9 @@ pub async fn json_handler( State(state): State, headers: HeaderMap, Query(query): Query, -) -> Result, AppError> { +) -> Result { let ip = resolve_ip(&headers, &query, state.dev_mode)?; - let info = lookup_ip(&state.db, ip)?; - Ok(axum::Json(info)) + lookup_json_response(&state.db, ip) } pub async fn ip_handler( @@ -68,12 +66,11 @@ fn field_handler( lookup: impl FnOnce(&SharedDb, IpAddr) -> db::LookupOutcome, ) -> Result { let ip = resolve_ip(headers, query, state.dev_mode)?; - if !db::is_ready(&state.db) { - return Err(AppError::DbNotReady); - } - let value = lookup(&state.db, ip) - .map_err(|_| AppError::DbLookupFailed)? - .unwrap_or_default(); + let value = match lookup(&state.db, ip).map_err(|_| AppError::DbLookupFailed)? { + Some(value) => value, + None if !db::is_ready(&state.db) => return Err(AppError::DbNotReady), + None => T::default(), + }; Ok(format!("{value}\n")) } diff --git a/src/handlers/html.rs b/src/handlers/html.rs index c5c2083..0f26d56 100644 --- a/src/handlers/html.rs +++ b/src/handlers/html.rs @@ -5,14 +5,13 @@ use axum::extract::{Path, Query, State}; use axum::http::{HeaderMap, StatusCode}; use axum::response::{Html, IntoResponse, Redirect, Response}; +use crate::db; use crate::error::AppError; use crate::i18n::{self, Locale}; -use crate::models::{IpInfo, get_localized_name}; +use crate::models::PageInfo; use crate::static_assets; -use super::{ - AppState, IpQuery, append_ip_query, handle_non_browser_request, lookup_ip, resolve_ip, -}; +use super::{AppState, IpQuery, append_ip_query, handle_non_browser_request, resolve_ip}; /// Askama custom filters. The `filters` module name is required by Askama's /// filter resolution. Usage in templates: `{{ "msgid"|t(locale) }}`. @@ -20,13 +19,12 @@ mod filters { use crate::i18n::{self, Locale}; #[askama::filter_fn] - pub fn t( - msgid: &T, + pub fn t<'a>( + msgid: &'a str, _env: &dyn askama::Values, locale: &Locale, - ) -> askama::Result { - let key = msgid.to_string(); - Ok(i18n::translate(*locale, &key).to_string()) + ) -> askama::Result<&'a str> { + Ok(i18n::translate(*locale, msgid)) } } @@ -83,42 +81,16 @@ struct ErrorPageTemplate { locale: Locale, } -fn asn_number(info: &IpInfo) -> Option { - info.asn.as_ref().and_then(|a| a.autonomous_system_number) -} - -fn format_asn(info: &IpInfo) -> String { - asn_number(info) - .map(|n| format!("AS{n}")) - .unwrap_or_default() -} - -fn format_org(info: &IpInfo) -> &str { - info.asn - .as_ref() - .and_then(|a| a.autonomous_system_organization.as_deref()) - .unwrap_or_default() -} - -fn format_country(info: &IpInfo, locale: Locale) -> &str { - info.country - .as_ref() - .map(|c| get_localized_name(&c.names, locale.mmdb_key())) - .unwrap_or_default() -} - -fn format_country_code(info: &IpInfo) -> &str { - info.country - .as_ref() - .and_then(|c| c.iso_code.as_deref()) - .unwrap_or_default() -} - -fn format_city(info: &IpInfo, locale: Locale) -> &str { - info.city - .as_ref() - .map(|c| get_localized_name(&c.names, locale.mmdb_key())) - .unwrap_or_default() +fn lookup_page_info( + db: &db::SharedDb, + ip: std::net::IpAddr, + locale: Locale, +) -> Result { + match db::lookup_page_info(db, ip, locale.mmdb_key()).map_err(|_| AppError::DbLookupFailed)? { + Some(info) => Ok(info), + None if !db::is_ready(db) => Err(AppError::DbNotReady), + None => Err(AppError::IpNotFound), + } } fn render_error_alert(state: &AppState, locale: Locale, query: &str, error: &AppError) -> Response { @@ -225,26 +197,26 @@ pub async fn root( Err(e) => return render_error_alert(&state, locale, &query_str, &e), }; - let info = match lookup_ip(&state.db, ip) { + let info = match lookup_page_info(&state.db, ip, locale) { Ok(info) => info, Err(e) => return render_error_alert(&state, locale, &query_str, &e), }; let is_query = query.ip.is_some(); let template = IndexTemplate { - ip: info.ip.clone(), + ip: info.ip, query: query.ip.unwrap_or_default(), is_query, site_domain: state.site_domain.clone(), ipv4_domain: state.ipv4_domain.clone(), css_version: static_assets::asset_version("style.css"), js_version: static_assets::asset_version("script.js"), - asn: format_asn(&info), - asn_number: asn_number(&info), - org: format_org(&info).to_owned(), - country: format_country(&info, locale).to_owned(), - country_code: format_country_code(&info).to_owned(), - city: format_city(&info, locale).to_owned(), + asn: info.asn, + asn_number: info.asn_number, + org: info.org, + country: info.country, + country_code: info.country_code, + city: info.city, is_proxy: info.proxy.is_proxy, is_vpn: info.proxy.is_vpn, is_hosting: info.proxy.is_hosting, diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index ca51d19..b1e14bd 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -13,7 +13,6 @@ use axum::response::{IntoResponse, Response}; use crate::db::{self, SharedDb}; use crate::error::AppError; use crate::ip_validate; -use crate::models::IpInfo; use crate::ua_detect::is_plain_text_agent; // Re-export all public handlers for use by routes.rs @@ -79,13 +78,12 @@ fn extract_client_ip(headers: &HeaderMap, dev_mode: bool) -> Result Result { - if !db::is_ready(db) { - return Err(AppError::DbNotReady); +fn lookup_json_response(db: &SharedDb, ip: IpAddr) -> Result { + match db::lookup_json(db, ip).map_err(|_| AppError::DbLookupFailed)? { + Some(body) => Ok(([(header::CONTENT_TYPE, "application/json")], body).into_response()), + None if !db::is_ready(db) => Err(AppError::DbNotReady), + None => Err(AppError::IpNotFound), } - db::lookup(db, ip) - .map_err(|_| AppError::DbLookupFailed)? - .ok_or(AppError::IpNotFound) } fn handle_non_browser_request( @@ -107,9 +105,10 @@ fn handle_non_browser_request( if wants_json(headers) && query.ip.is_none() { return Some( - match extract_client_ip(headers, state.dev_mode).and_then(|ip| lookup_ip(&state.db, ip)) + match extract_client_ip(headers, state.dev_mode) + .and_then(|ip| lookup_json_response(&state.db, ip)) { - Ok(info) => axum::Json(info).into_response(), + Ok(response) => response, Err(e) => e.into_response(), }, ); diff --git a/src/models/api.rs b/src/models/api.rs index f821fa2..733355d 100644 --- a/src/models/api.rs +++ b/src/models/api.rs @@ -123,6 +123,18 @@ pub struct ProxyInfo { pub is_anonymous: bool, } +#[derive(Debug)] +pub struct PageInfo { + pub ip: String, + pub asn: String, + pub asn_number: Option, + pub org: String, + pub country: String, + pub country_code: String, + pub city: String, + pub proxy: ProxyInfo, +} + impl From for ProxyInfo { fn from(p: MmdbProxy) -> Self { Self { diff --git a/src/models/mod.rs b/src/models/mod.rs index 10a057e..f40ea59 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -6,12 +6,251 @@ pub mod mmdb; use std::fmt::Write; use std::net::IpAddr; +use serde::Serialize; + pub use api::*; // Mmdb types are re-exported as pub(crate) since they are internal deserialization details. pub(crate) use mmdb::*; // --- Conversion from MMDB record to API output --- +#[derive(Serialize)] +struct JsonIpInfo<'a> { + ip: String, + #[serde(skip_serializing_if = "Option::is_none")] + city: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + continent: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + country: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + location: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + postal: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + registered_country: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + subdivisions: Option>>, + #[serde(skip_serializing_if = "Option::is_none")] + asn: Option>, + proxy: ProxyInfo, +} + +#[derive(Serialize)] +struct JsonNames<'a> { + #[serde(rename = "de", skip_serializing_if = "Option::is_none")] + german: Option<&'a str>, + #[serde(rename = "en", skip_serializing_if = "Option::is_none")] + english: Option<&'a str>, + #[serde(rename = "es", skip_serializing_if = "Option::is_none")] + spanish: Option<&'a str>, + #[serde(rename = "fr", skip_serializing_if = "Option::is_none")] + french: Option<&'a str>, + #[serde(rename = "ja", skip_serializing_if = "Option::is_none")] + japanese: Option<&'a str>, + #[serde(rename = "pt-BR", skip_serializing_if = "Option::is_none")] + brazilian_portuguese: Option<&'a str>, + #[serde(rename = "ru", skip_serializing_if = "Option::is_none")] + russian: Option<&'a str>, + #[serde(rename = "zh-CN", skip_serializing_if = "Option::is_none")] + simplified_chinese: Option<&'a str>, +} + +#[derive(Serialize)] +struct JsonCityInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + geoname_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + names: Option>, +} + +#[derive(Serialize)] +struct JsonContinentInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + code: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + geoname_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + names: Option>, +} + +#[derive(Serialize)] +struct JsonCountryInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + geoname_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + iso_code: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + names: Option>, +} + +#[derive(Serialize)] +struct JsonLocationInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + accuracy_radius: Option, + #[serde(skip_serializing_if = "Option::is_none")] + latitude: Option, + #[serde(skip_serializing_if = "Option::is_none")] + longitude: Option, + #[serde(skip_serializing_if = "Option::is_none")] + metro_code: Option, + #[serde(skip_serializing_if = "Option::is_none")] + time_zone: Option<&'a str>, +} + +#[derive(Serialize)] +struct JsonPostalInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + code: Option<&'a str>, +} + +#[derive(Serialize)] +struct JsonSubdivisionInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + geoname_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + iso_code: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + names: Option>, +} + +#[derive(Serialize)] +struct JsonAsnInfo<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + autonomous_system_number: Option, + #[serde(skip_serializing_if = "Option::is_none")] + autonomous_system_organization: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + as_domain: Option<&'a str>, +} + +fn json_names<'a>(names: &MmdbNames<'a>) -> Option> { + if names.is_empty() { + return None; + } + Some(JsonNames { + german: names.german, + english: names.english, + spanish: names.spanish, + french: names.french, + japanese: names.japanese, + brazilian_portuguese: names.brazilian_portuguese, + russian: names.russian, + simplified_chinese: names.simplified_chinese, + }) +} + +fn json_ip_string(ip: IpAddr) -> String { + let mut ip_str = String::with_capacity(45); + write!(ip_str, "{ip}").expect("IP address formatting cannot fail"); + ip_str +} + +fn json_record<'a>(ip: IpAddr, record: &MmdbRecord<'a>) -> JsonIpInfo<'a> { + let city = if has_city(&record.city) { + Some(JsonCityInfo { + geoname_id: record.city.geoname_id, + names: json_names(&record.city.names), + }) + } else { + None + }; + + let continent = if has_continent(&record.continent) { + Some(JsonContinentInfo { + code: record.continent.code, + geoname_id: record.continent.geoname_id, + names: json_names(&record.continent.names), + }) + } else { + None + }; + + let country = if has_country(&record.country) { + Some(JsonCountryInfo { + geoname_id: record.country.geoname_id, + iso_code: record.country.iso_code, + names: json_names(&record.country.names), + }) + } else { + None + }; + + let location = if has_location(&record.location) { + Some(JsonLocationInfo { + accuracy_radius: record.location.accuracy_radius, + latitude: record.location.latitude, + longitude: record.location.longitude, + metro_code: record.location.metro_code, + time_zone: record.location.time_zone, + }) + } else { + None + }; + + let postal = if has_postal(&record.postal) { + Some(JsonPostalInfo { + code: record.postal.code, + }) + } else { + None + }; + + let registered_country = if has_country(&record.registered_country) { + Some(JsonCountryInfo { + geoname_id: record.registered_country.geoname_id, + iso_code: record.registered_country.iso_code, + names: json_names(&record.registered_country.names), + }) + } else { + None + }; + + let subdivisions = if record.subdivisions.is_empty() { + None + } else { + Some( + record + .subdivisions + .iter() + .map(|s| JsonSubdivisionInfo { + geoname_id: s.geoname_id, + iso_code: s.iso_code, + names: json_names(&s.names), + }) + .collect(), + ) + }; + + let asn = if has_asn(&record.asn) { + Some(JsonAsnInfo { + autonomous_system_number: record.asn.autonomous_system_number, + autonomous_system_organization: record.asn.autonomous_system_organization, + as_domain: record.asn.as_domain, + }) + } else { + None + }; + + JsonIpInfo { + ip: json_ip_string(ip), + city, + continent, + country, + location, + postal, + registered_country, + subdivisions, + asn, + proxy: ProxyInfo::from(record.proxy), + } +} + +pub fn serialize_mmdb_record_json(ip: IpAddr, record: &MmdbRecord<'_>) -> Vec { + serde_json::to_vec(&json_record(ip, record)) + .expect("serializing MMDB lookup response cannot fail") +} + fn convert_names(names: &MmdbNames<'_>) -> Option { if names.is_empty() { return None; @@ -149,11 +388,8 @@ pub fn from_mmdb_record(ip: IpAddr, record: &MmdbRecord<'_>) -> IpInfo { let proxy = ProxyInfo::from(record.proxy); - let mut ip_str = String::with_capacity(45); - write!(ip_str, "{ip}").expect("IP address formatting cannot fail"); - IpInfo { - ip: ip_str, + ip: json_ip_string(ip), city, continent, country,