use axum::body::Body; use axum::extract::{ConnectInfo, Request}; use axum::http::{HeaderMap, StatusCode}; use axum::middleware::Next; use axum::response::{IntoResponse, Response}; use metrics::{counter, gauge, histogram}; use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle}; use std::net::{IpAddr, SocketAddr}; use std::time::Instant; /// Initialize the Prometheus metrics exporter and return a handle for rendering metrics. /// /// Configures histogram bucket boundaries so the exporter renders Prometheus histograms /// (with `_bucket` suffix) instead of summaries. Without this, `histogram_quantile()` /// queries in Grafana find no `_bucket` metrics and return empty. pub fn init_metrics() -> PrometheusHandle { // Standard Prometheus buckets for HTTP latencies (seconds) const LATENCY_BUCKETS: &[f64] = &[ 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, ]; // Wider buckets for screenshot generation (can take 30s+) const SCREENSHOT_BUCKETS: &[f64] = &[0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0]; // Count-based buckets for response sizes (number of hexagons/postcodes returned) const RESPONSE_COUNT_BUCKETS: &[f64] = &[ 10.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 2500.0, 5000.0, 10000.0, ]; PrometheusBuilder::new() .set_buckets_for_metric( Matcher::Full("http_request_duration_seconds".to_string()), LATENCY_BUCKETS, ) .expect("Failed to set HTTP latency buckets") .set_buckets_for_metric( Matcher::Full("screenshot_duration_seconds".to_string()), SCREENSHOT_BUCKETS, ) .expect("Failed to set screenshot duration buckets") .set_buckets_for_metric( Matcher::Full("hexagons_response_count".to_string()), RESPONSE_COUNT_BUCKETS, ) .expect("Failed to set hexagons response count buckets") .set_buckets_for_metric( Matcher::Full("postcodes_response_count".to_string()), RESPONSE_COUNT_BUCKETS, ) .expect("Failed to set postcodes response count buckets") .install_recorder() .expect("Failed to install Prometheus recorder") } /// Middleware to track HTTP request metrics. pub async fn track_metrics(request: Request, next: Next) -> Response { let path = request.uri().path().to_string(); let method = request.method().to_string(); // Skip metrics endpoint itself to avoid recursion if path == "/metrics" { return next.run(request).await; } let start = Instant::now(); let response = next.run(request).await; let duration = start.elapsed(); let status = response.status().as_u16().to_string(); // Normalize path for metrics (avoid high cardinality from dynamic segments) let normalized_path = normalize_path(&path); // Record metrics histogram!("http_request_duration_seconds", "method" => method.clone(), "path" => normalized_path.clone(), "status" => status.clone()) .record(duration.as_secs_f64()); counter!("http_requests_total", "method" => method, "path" => normalized_path, "status" => status) .increment(1); response } /// Normalize paths to avoid high cardinality from dynamic segments. /// /// Groups dynamic segments into parameterized placeholders and collapses /// static assets, PocketBase proxy paths, and unknown paths to prevent /// Prometheus label cardinality explosion from bot scans and unique URLs. fn normalize_path(path: &str) -> String { // Tiles: /api/tiles/5/16/10 → /api/tiles/:z/:x/:y if path.starts_with("/api/tiles/") && !path.ends_with("style.json") { return "/api/tiles/:z/:x/:y".to_string(); } // Invite API: /api/invite/abc123 → /api/invite/:code if path.starts_with("/api/invite/") { return "/api/invite/:code".to_string(); } // PocketBase proxy: /pb/api/files/... → /pb/api/files/:path if path.starts_with("/pb/api/files/") { return "/pb/api/files/:path".to_string(); } // PocketBase proxy: /pb/api/... → keep collection-level granularity if path.starts_with("/pb/api/collections/") { // /pb/api/collections/users/auth-with-password → keep as-is (bounded set) // /pb/api/collections/saved_searches/records/abc → /pb/api/collections/saved_searches/records/:id let parts: Vec<&str> = path.splitn(6, '/').collect(); if parts.len() >= 6 && parts[4] == "records" { return format!("/pb/api/collections/{}/records/:id", parts[3]); } return path.to_string(); } // Short URLs: /s/abc → /s/:code if path.starts_with("/s/") { return "/s/:code".to_string(); } // Invite pages: /invite/abc → /invite/:code if path.starts_with("/invite/") { return "/invite/:code".to_string(); } // Static assets: /assets/* → /assets/:file if path.starts_with("/assets/") { return "/assets/:file".to_string(); } // Known application routes and API endpoints: keep as-is if path.starts_with("/api/") || matches!( path, "/" | "/health" | "/metrics" | "/dashboard" | "/pricing" | "/account" | "/saved" | "/invites" | "/learn" | "/bundle.js" | "/main.css" | "/favicon.ico" | "/house.png" | "/robots.txt" | "/sitemap.xml" ) { return path.to_string(); } // Everything else (bot scans, probes, etc.) → /other "/other".to_string() } /// Handler for the /metrics endpoint. Only accepts requests from peers on the /// same private network (loopback, RFC1918, or IPv6 unique/link-local). /// /// The TCP peer is always the reverse proxy (a private container IP), so the /// real client IP is taken from the proxy-set `X-Real-IP` / `X-Forwarded-For` /// headers when present; a public client reaching this endpoint through the /// public domain is therefore rejected. Direct in-network connections (e.g. the /// Prometheus scraper hitting the container directly) carry no such header and /// fall back to the trusted TCP peer address. pub async fn metrics_handler( handle: PrometheusHandle, headers: HeaderMap, ConnectInfo(peer): ConnectInfo, ) -> Response { let client_ip = client_ip_from_headers(&headers).unwrap_or_else(|| peer.ip()); if !is_same_network(client_ip) { return StatusCode::FORBIDDEN.into_response(); } update_process_metrics(); match handle.render() { output if !output.is_empty() => (StatusCode::OK, output).into_response(), _ => ( StatusCode::INTERNAL_SERVER_ERROR, "Failed to render metrics".to_string(), ) .into_response(), } } /// Resolve the real client IP from trusted reverse-proxy headers. nginx sets /// `X-Real-IP` to the post-`real_ip` client address and overwrites any /// client-supplied value, so when present it is authoritative; the left-most /// `X-Forwarded-For` hop is used as a fallback. Returns `None` when neither /// header is set, so a direct in-network connection uses its TCP peer address. fn client_ip_from_headers(headers: &HeaderMap) -> Option { if let Some(value) = headers.get("x-real-ip").and_then(|v| v.to_str().ok()) { if let Ok(ip) = value.trim().parse::() { return Some(ip); } } if let Some(value) = headers.get("x-forwarded-for").and_then(|v| v.to_str().ok()) { if let Some(first) = value.split(',').next() { if let Ok(ip) = first.trim().parse::() { return Some(ip); } } } None } fn is_same_network(ip: IpAddr) -> bool { match ip { IpAddr::V4(v4) => v4.is_loopback() || v4.is_private() || v4.is_link_local(), IpAddr::V6(v6) => { v6.is_loopback() || (v6.segments()[0] & 0xfe00) == 0xfc00 || (v6.segments()[0] & 0xffc0) == 0xfe80 || v6 .to_ipv4_mapped() .is_some_and(|v4| v4.is_loopback() || v4.is_private() || v4.is_link_local()) } } } /// Update process-level metrics (memory, etc.) fn update_process_metrics() { #[cfg(target_os = "linux")] { if let Ok(status) = std::fs::read_to_string("/proc/self/status") { for line in status.lines() { if let Some(value) = line.strip_prefix("VmRSS:") { if let Some(kb) = parse_kb(value) { gauge!("process_resident_memory_bytes").set((kb * 1024) as f64); } } else if let Some(value) = line.strip_prefix("VmSize:") { if let Some(kb) = parse_kb(value) { gauge!("process_virtual_memory_bytes").set((kb * 1024) as f64); } } } } } } #[cfg(target_os = "linux")] fn parse_kb(value: &str) -> Option { value .trim() .strip_suffix(" kB") .or_else(|| value.trim().strip_suffix("kB")) .and_then(|num| num.trim().parse().ok()) } /// Record a custom gauge for data loading (call at startup). pub fn record_data_stats(property_count: usize, poi_count: usize, postcode_count: usize) { gauge!("data_properties_loaded").set(property_count as f64); gauge!("data_pois_loaded").set(poi_count as f64); gauge!("data_postcodes_loaded").set(postcode_count as f64); } #[cfg(test)] mod tests { use super::*; #[test] fn x_real_ip_is_authoritative_and_public_ip_is_rejected() { let mut headers = HeaderMap::new(); headers.insert("x-real-ip", "203.0.113.7".parse().unwrap()); let ip = client_ip_from_headers(&headers).expect("resolves from x-real-ip"); assert_eq!(ip, "203.0.113.7".parse::().unwrap()); assert!(!is_same_network(ip), "a public client must be rejected"); } #[test] fn missing_proxy_headers_fall_back_to_peer() { // No X-Real-IP / X-Forwarded-For → caller uses the TCP peer (e.g. an // in-network Prometheus scraper hitting the container directly). assert_eq!(client_ip_from_headers(&HeaderMap::new()), None); } #[test] fn forwarded_for_uses_first_hop() { let mut headers = HeaderMap::new(); headers.insert("x-forwarded-for", "198.51.100.4, 10.0.0.1".parse().unwrap()); assert_eq!( client_ip_from_headers(&headers), Some("198.51.100.4".parse::().unwrap()) ); } }