diff --git a/.env.example b/.env.example index 015f7ae..13d3f8b 100644 --- a/.env.example +++ b/.env.example @@ -13,6 +13,9 @@ DB_UPDATE_INTERVAL_HOURS=24 # Site domain name (used for display/metadata) SITE_DOMAIN=localhost +# Enable development mode (uses 1.1.1.1 as fallback when CF-Connecting-IP is absent) +DEV_MODE=false + # Path to the TLS certificate directory for HAproxy # The directory must contain origin.pem (combined cert + key) CERT_PATH=./haproxy/certs diff --git a/src/config.rs b/src/config.rs index fc1b882..0001a7c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -7,6 +7,7 @@ pub struct Config { pub db_update_url: String, pub db_update_interval_hours: u64, pub site_domain: String, + pub dev_mode: bool, } impl Config { @@ -30,12 +31,17 @@ impl Config { let site_domain = env::var("SITE_DOMAIN").unwrap_or_else(|_| "localhost".to_string()); + let dev_mode = env::var("DEV_MODE") + .unwrap_or_default() + .eq_ignore_ascii_case("true"); + Self { listen_addr, db_path, db_update_url, db_update_interval_hours, site_domain, + dev_mode, } } } @@ -61,6 +67,7 @@ mod tests { "DB_UPDATE_URL", "DB_UPDATE_INTERVAL_HOURS", "SITE_DOMAIN", + "DEV_MODE", ]; // SAFETY: These env var helpers are only called while holding ENV_LOCK, @@ -91,6 +98,7 @@ mod tests { assert!(config.db_update_url.contains("Merged-IP.mmdb")); assert_eq!(config.db_update_interval_hours, 24); assert_eq!(config.site_domain, "localhost"); + assert!(!config.dev_mode); } #[test] @@ -173,6 +181,48 @@ mod tests { unsafe { remove_var("SITE_DOMAIN") }; } + #[test] + fn dev_mode_enabled() { + let _guard = lock_env(); + // SAFETY: ENV_LOCK is held. + unsafe { + clear_config_vars(); + set_var("DEV_MODE", "true"); + } + + let config = Config::from_env(); + assert!(config.dev_mode); + + // SAFETY: ENV_LOCK is held. + unsafe { remove_var("DEV_MODE") }; + } + + #[test] + fn dev_mode_case_insensitive() { + let _guard = lock_env(); + // SAFETY: ENV_LOCK is held. + unsafe { + clear_config_vars(); + set_var("DEV_MODE", "True"); + } + + let config = Config::from_env(); + assert!(config.dev_mode); + + // SAFETY: ENV_LOCK is held. + unsafe { remove_var("DEV_MODE") }; + } + + #[test] + fn dev_mode_disabled_by_default() { + let _guard = lock_env(); + // SAFETY: ENV_LOCK is held. + unsafe { clear_config_vars() }; + + let config = Config::from_env(); + assert!(!config.dev_mode); + } + #[test] #[should_panic(expected = "LISTEN_ADDR must be a valid socket address")] fn invalid_listen_addr_panics() { diff --git a/src/handlers.rs b/src/handlers.rs index 2301c9d..c44025c 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -17,10 +17,13 @@ pub struct IpQuery { pub ip: Option, } +const DEV_FALLBACK_IP: IpAddr = IpAddr::V4(std::net::Ipv4Addr::new(1, 1, 1, 1)); + #[derive(Clone)] pub struct AppState { pub db: SharedDb, pub site_domain: Arc, + pub dev_mode: bool, } #[derive(Template)] @@ -40,22 +43,30 @@ struct IndexTemplate { is_tor: bool, } -fn resolve_ip(headers: &HeaderMap, query: &IpQuery) -> Result { +fn resolve_ip(headers: &HeaderMap, query: &IpQuery, dev_mode: bool) -> Result { if let Some(ref ip_str) = query.ip { let trimmed = ip_str.trim(); if !trimmed.is_empty() { return trimmed.parse().map_err(|_| AppError::InvalidIp); } } - extract_client_ip(headers) + extract_client_ip(headers, dev_mode) } -fn extract_client_ip(headers: &HeaderMap) -> Result { - headers +fn extract_client_ip(headers: &HeaderMap, dev_mode: bool) -> Result { + if let Some(ip) = headers .get("CF-Connecting-IP") .and_then(|v| v.to_str().ok()) .and_then(|s| s.trim().parse().ok()) - .ok_or(AppError::MissingClientIp) + { + return Ok(ip); + } + + if dev_mode { + return Ok(DEV_FALLBACK_IP); + } + + Err(AppError::MissingClientIp) } fn lookup_ip(db: &SharedDb, ip: IpAddr) -> Result { @@ -110,17 +121,17 @@ pub async fn root( .unwrap_or(""); if is_cli_user_agent(ua) && query.ip.is_none() { - let ip = extract_client_ip(&headers)?; + let ip = extract_client_ip(&headers, state.dev_mode)?; return Ok(format!("{ip}\n").into_response()); } if wants_json(&headers) && query.ip.is_none() { - let ip = extract_client_ip(&headers)?; + let ip = extract_client_ip(&headers, state.dev_mode)?; let info = lookup_ip(&state.db, ip)?; return Ok(axum::Json(info).into_response()); } - let ip = resolve_ip(&headers, &query)?; + let ip = resolve_ip(&headers, &query, state.dev_mode)?; let info = lookup_ip(&state.db, ip)?; let template = IndexTemplate { @@ -150,16 +161,17 @@ pub async fn json_handler( headers: HeaderMap, Query(query): Query, ) -> Result, AppError> { - let ip = resolve_ip(&headers, &query)?; + let ip = resolve_ip(&headers, &query, state.dev_mode)?; let info = lookup_ip(&state.db, ip)?; Ok(axum::Json(info)) } pub async fn ip_handler( + State(state): State, headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + let ip = resolve_ip(&headers, &query, state.dev_mode)?; Ok(format!("{ip}\n")) } @@ -168,7 +180,7 @@ pub async fn asn_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + let ip = resolve_ip(&headers, &query, state.dev_mode)?; let reader = state.db.load(); let text = db::lookup_asn_number(&reader, ip) .map(|n| format!("AS{n}")) @@ -181,7 +193,7 @@ pub async fn org_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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(); Ok(format!("{org}\n")) @@ -192,7 +204,7 @@ pub async fn country_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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(); Ok(format!("{country}\n")) @@ -203,7 +215,7 @@ pub async fn city_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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(); Ok(format!("{city}\n")) @@ -214,7 +226,7 @@ pub async fn proxy_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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)?; Ok(format!("{}\n", proxy.is_proxy)) @@ -225,7 +237,7 @@ pub async fn vpn_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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)?; Ok(format!("{}\n", proxy.is_vpn)) @@ -236,7 +248,7 @@ pub async fn hosting_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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)?; Ok(format!("{}\n", proxy.is_hosting)) @@ -247,7 +259,7 @@ pub async fn tor_handler( headers: HeaderMap, Query(query): Query, ) -> Result { - let ip = resolve_ip(&headers, &query)?; + 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)?; Ok(format!("{}\n", proxy.is_tor)) diff --git a/src/main.rs b/src/main.rs index 085e748..ff43e0a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -41,6 +41,7 @@ async fn main() { let state = handlers::AppState { db: shared_db, site_domain: config.site_domain.into(), + dev_mode: config.dev_mode, }; let app = routes::build_router(state); diff --git a/tests/integration.rs b/tests/integration.rs index bdf5cf4..88891bb 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -18,11 +18,16 @@ fn test_db_path() -> String { } fn build_test_app() -> axum::Router { + build_test_app_with_dev_mode(false) +} + +fn build_test_app_with_dev_mode(dev_mode: bool) -> axum::Router { let reader = db::load_db(Path::new(&test_db_path())).expect("failed to load test DB"); let shared_db = db::new_shared(reader); let state = AppState { db: shared_db, site_domain: "test.example.com".into(), + dev_mode, }; build_router(state) } @@ -474,3 +479,81 @@ async fn html_uses_th_for_row_headers() { assert!(body.contains("Country")); assert!(body.contains("autocapitalize=\"none\"")); } + +// --- Dev mode tests --- + +#[tokio::test] +async fn dev_mode_root_returns_html_without_cf_header() { + let app = build_test_app_with_dev_mode(true); + let (status, body) = get_with_headers( + &app, + "/", + vec![( + "User-Agent", + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", + )], + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!(body.contains("")); + assert!(body.contains("1.1.1.1")); +} + +#[tokio::test] +async fn dev_mode_cli_returns_fallback_ip() { + let app = build_test_app_with_dev_mode(true); + let (status, body) = get_with_headers(&app, "/", vec![("User-Agent", "curl/8.7.1")]).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body.trim(), "1.1.1.1"); +} + +#[tokio::test] +async fn dev_mode_json_returns_fallback_ip() { + let app = build_test_app_with_dev_mode(true); + let (status, body) = get(&app, "/json").await; + assert_eq!(status, StatusCode::OK); + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(json["ip"], "1.1.1.1"); +} + +#[tokio::test] +async fn dev_mode_cf_header_takes_precedence() { + let app = build_test_app_with_dev_mode(true); + let (status, body) = get_with_headers( + &app, + "/", + vec![ + ("User-Agent", "curl/8.7.1"), + ("CF-Connecting-IP", "8.8.8.8"), + ], + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body.trim(), "8.8.8.8"); +} + +#[tokio::test] +async fn dev_mode_query_param_overrides_fallback() { + let app = build_test_app_with_dev_mode(true); + let (status, body) = get(&app, "/json?ip=45.77.77.77").await; + assert_eq!(status, StatusCode::OK); + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(json["ip"], "45.77.77.77"); +} + +#[tokio::test] +async fn dev_mode_field_endpoints_return_fallback_data() { + let app = build_test_app_with_dev_mode(true); + + let (status, body) = get(&app, "/ip").await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body.trim(), "1.1.1.1"); + + let (status, body) = get(&app, "/asn").await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body.trim(), "AS13335"); + + let (status, body) = get(&app, "/org").await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body.trim(), "Cloudflare, Inc."); +}