279 lines
10 KiB
Rust
279 lines
10 KiB
Rust
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<Body>, 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<SocketAddr>,
|
|
) -> 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<IpAddr> {
|
|
if let Some(value) = headers.get("x-real-ip").and_then(|v| v.to_str().ok()) {
|
|
if let Ok(ip) = value.trim().parse::<IpAddr>() {
|
|
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::<IpAddr>() {
|
|
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<u64> {
|
|
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::<IpAddr>().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::<IpAddr>().unwrap())
|
|
);
|
|
}
|
|
}
|