From be5450f9cefb281051bbbbf025c48b6e75183bed Mon Sep 17 00:00:00 2001 From: NotAShelf Date: Thu, 1 May 2025 20:42:48 +0300 Subject: [PATCH] metrics: allow customizing address in full; add tests --- nix/module.nix | 10 +++--- src/main.rs | 27 +++++++++-------- src/metrics.rs | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 17 deletions(-) diff --git a/nix/module.nix b/nix/module.nix index 1aeb67d..0eea871 100644 --- a/nix/module.nix +++ b/nix/module.nix @@ -16,7 +16,7 @@ self: { # Generate the config.json content erisConfigFile = pkgs.writeText "eris-config.json" (toJSON { listen_addr = cfg.listenAddress; - metrics_port = cfg.metricsPort; + metrics_addr = cfg.metricsAddress; backend_addr = cfg.backendAddress; min_delay = cfg.minDelay; max_delay = cfg.maxDelay; @@ -50,11 +50,11 @@ in { example = "127.0.0.1:9999"; }; - metricsPort = mkOption { + metricsAddress = mkOption { type = port; - default = 9100; - example = 9110; - description = "The port for the Prometheus metrics endpoint."; + default = "0.0.0.0:9100"; + example = "127.0.0.1:9100"; + description = "The IP address and port for the Prometheus metrics endpoint."; }; backendAddress = mkOption { diff --git a/src/main.rs b/src/main.rs index 8bc3e98..2471ced 100644 --- a/src/main.rs +++ b/src/main.rs @@ -43,10 +43,10 @@ struct Args { #[clap( long, - default_value = "9100", - help = "Port to expose Prometheus metrics and status endpoint" + default_value = "0.0.0.0:9100", + help = "Address and port to expose Prometheus metrics and status endpoint (format: ip:port)" )] - metrics_port: u16, + metrics_addr: String, #[clap(long, help = "Disable Prometheus metrics server completely")] disable_metrics: bool, @@ -110,7 +110,7 @@ struct Args { #[derive(Clone, Debug, Deserialize, Serialize)] struct Config { listen_addr: String, - metrics_port: u16, + metrics_addr: String, disable_metrics: bool, backend_addr: String, min_delay: u64, @@ -130,7 +130,7 @@ impl Default for Config { fn default() -> Self { Self { listen_addr: "0.0.0.0:8888".to_string(), - metrics_port: 9100, + metrics_addr: "0.0.0.0:9100".to_string(), disable_metrics: false, backend_addr: "127.0.0.1:80".to_string(), min_delay: 1000, @@ -216,7 +216,7 @@ impl Config { Self { listen_addr: args.listen_addr.clone(), - metrics_port: args.metrics_port, + metrics_addr: args.metrics_addr.clone(), disable_metrics: args.disable_metrics, backend_addr: args.backend_addr.clone(), min_delay: args.min_delay, @@ -1192,8 +1192,7 @@ async fn main() -> std::io::Result<()> { log::info!("Metrics server disabled via configuration"); None } else { - let metrics_addr = format!("0.0.0.0:{}", metrics_config.metrics_port); - log::info!("Starting metrics server on {metrics_addr}"); + log::info!("Starting metrics server on {}", metrics_config.metrics_addr); let server = HttpServer::new(move || { App::new() @@ -1204,12 +1203,16 @@ async fn main() -> std::io::Result<()> { HttpResponse::Ok().body("Botpot Server is running. Visit /metrics for metrics or /status for status.") })) }) - .bind(&metrics_addr); + .bind(&metrics_config.metrics_addr); match server { Ok(server) => Some(server.run()), Err(e) => { - log::error!("Failed to bind metrics server to {metrics_addr}: {e}"); + log::error!( + "Failed to bind metrics server to {}: {}", + metrics_config.metrics_addr, + e + ); None } } @@ -1265,7 +1268,7 @@ mod tests { fn test_config_from_args() { let args = Args { listen_addr: "127.0.0.1:8080".to_string(), - metrics_port: 9000, + metrics_addr: "127.0.0.1:9000".to_string(), disable_metrics: true, backend_addr: "127.0.0.1:8081".to_string(), min_delay: 500, @@ -1279,7 +1282,7 @@ mod tests { let config = Config::from_args(&args); assert_eq!(config.listen_addr, "127.0.0.1:8080"); - assert_eq!(config.metrics_port, 9000); + assert_eq!(config.metrics_addr, "127.0.0.1:9000"); assert!(config.disable_metrics); assert_eq!(config.backend_addr, "127.0.0.1:8081"); assert_eq!(config.min_delay, 500); diff --git a/src/metrics.rs b/src/metrics.rs index c733752..48bceb9 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -59,3 +59,85 @@ pub async fn status_handler(state: web::Data>>) -> HttpResp .content_type("application/json") .body(serde_json::to_string_pretty(&info).unwrap()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::BotState; + use actix_web::{App, http, test}; + use std::net::{IpAddr, Ipv4Addr}; + + #[actix_web::test] + async fn test_metrics_handler() { + // For this test to be functional, we'll need to have some metrics + // with some values. + HITS_COUNTER.inc(); + PATH_HITS.with_label_values(&["/test/path"]).inc(); + UA_HITS.with_label_values(&["TestBot/1.0"]).inc(); + BLOCKED_IPS.set(5.0); + ACTIVE_CONNECTIONS.set(3.0); + + // Create test app + let app = + test::init_service(App::new().route("/metrics", web::get().to(metrics_handler))).await; + + // Send request + let req = test::TestRequest::get().uri("/metrics").to_request(); + let resp = test::call_service(&app, req).await; + + // Assert response + assert_eq!(resp.status(), http::StatusCode::OK); + let body = test::read_body(resp).await; + let body_str = String::from_utf8(body.to_vec()).unwrap(); + + // And now, lets verify metrics content + assert!(body_str.contains("eris_hits_total")); + assert!(body_str.contains("eris_path_hits_total")); + assert!(body_str.contains("eris_ua_hits_total")); + assert!(body_str.contains("eris_blocked_ips")); + assert!(body_str.contains("eris_active_connections")); + } + + #[actix_web::test] + async fn test_status_handler() { + // Test bot state + let mut bot_state = BotState::new("/tmp/eris_test", "/tmp/eris_test_cache"); + let ip1 = IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)); + let ip2 = IpAddr::V4(Ipv4Addr::new(5, 6, 7, 8)); + + // Test data + bot_state.blocked.insert(ip1); + bot_state.active_connections.insert(ip2); + bot_state.hits.insert(ip1, 3); + bot_state.hits.insert(ip2, 1); + + let bot_state = Arc::new(RwLock::new(bot_state)); + + let app = test::init_service( + App::new() + .app_data(web::Data::new(bot_state.clone())) + .route("/status", web::get().to(status_handler)), + ) + .await; + + // Send request + let req = test::TestRequest::get().uri("/status").to_request(); + let resp = test::call_service(&app, req).await; + + // Assert response + assert_eq!(resp.status(), http::StatusCode::OK); + assert_eq!( + resp.headers().get(http::header::CONTENT_TYPE).unwrap(), + "application/json" + ); + + let body = test::read_body(resp).await; + let status: serde_json::Value = serde_json::from_slice(&body).unwrap(); + + // Verify status JSON content + assert_eq!(status["status"], "running"); + assert_eq!(status["blocked_ips"], 1); + assert_eq!(status["active_connections"], 1); + assert_eq!(status["hit_count"], 2); + } +}