From 5e327c22f46c1b8a1eb251087e4d9d6ce762124f Mon Sep 17 00:00:00 2001 From: nodebreaker-carlsagan Date: Fri, 31 Jan 2020 18:19:14 +0900 Subject: [PATCH] /metrics url Input check --- utils/prometheus/src/lib.rs | 34 +++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/utils/prometheus/src/lib.rs b/utils/prometheus/src/lib.rs index f1027ce7ef87d..a59f17a14a738 100644 --- a/utils/prometheus/src/lib.rs +++ b/utils/prometheus/src/lib.rs @@ -16,7 +16,7 @@ use futures_util::{FutureExt, future::Future}; use hyper::http::StatusCode; -use hyper::{Server, Body, Response, service::{service_fn, make_service_fn}}; +use hyper::{Server, Body, Request, Response, service::{service_fn, make_service_fn}}; use prometheus::{Encoder, TextEncoder, core::Collector}; use std::net::SocketAddr; #[cfg(not(target_os = "unknown"))] @@ -59,17 +59,25 @@ impl std::error::Error for Error { } } -async fn request_metrics(registry: Registry) -> Result, Error> { - let metric_families = registry.gather(); - let mut buffer = vec![]; - let encoder = TextEncoder::new(); - encoder.encode(&metric_families, &mut buffer).unwrap(); - - Response::builder() - .status(StatusCode::OK) - .header("Content-Type", encoder.format_type()) - .body(Body::from(buffer)) +async fn request_metrics(req: Request, registry: Registry) -> Result, Error> { + if req.uri().path() == "/metrics" { + let metric_families = registry.gather(); + let mut buffer = vec![]; + let encoder = TextEncoder::new(); + encoder.encode(&metric_families, &mut buffer).unwrap(); + + Response::builder() + .status(StatusCode::OK) + .header("Content-Type", encoder.format_type()) + .body(Body::from(buffer)) + .map_err(Error::Http) + } else { + Response::builder() + .status(StatusCode::NOT_FOUND) + .body(Body::from("Not found.")) .map_err(Error::Http) + } + } #[derive(Clone)] @@ -101,8 +109,8 @@ pub async fn init_prometheus(prometheus_addr: SocketAddr, registry: Registry) -> let registry = registry.clone(); async move { - Ok::<_, hyper::Error>(service_fn(move |_| { - request_metrics(registry.clone()) + Ok::<_, hyper::Error>(service_fn(move |req: Request| { + request_metrics(req, registry.clone()) })) } });