Enable prometheus metrics in binaries

ewpratten · Aug 4, 2023 · 41708f1 · 41708f1 · github-actions · Aug 4, 2023
1 parent a0e4ead
commit 41708f1
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 10 deletions.
diff --git a/libs/protomask-metrics/Cargo.toml b/libs/protomask-metrics/Cargo.toml
@@ -13,6 +13,7 @@ keywords = []
 categories = []
 
 [dependencies]
+hyper = { version = "0.14.27", features = ["server", "http1", "tcp"] }
 log = "^0.4"
 prometheus = "0.13.3"
 lazy_static = "1.4.0"
diff --git a/libs/protomask-metrics/src/http.rs b/libs/protomask-metrics/src/http.rs
@@ -0,0 +1,44 @@
+use hyper::{
+    service::{make_service_fn, service_fn},
+    Body, Method, Request, Response, Server,
+};
+use prometheus::{Encoder, TextEncoder};
+use std::{convert::Infallible, net::SocketAddr};
+
+/// Handle an HTTP request
+#[allow(clippy::unused_async)]
+async fn handle_request(request: Request<Body>) -> Result<Response<Body>, Infallible> {
+    // If the request is targeting the metrics endpoint
+    if request.method() == Method::GET && request.uri().path() == "/metrics" {
+        // Gather metrics
+        let metric_families = prometheus::gather();
+        let body = {
+            let mut buffer = Vec::new();
+            let encoder = TextEncoder::new();
+            encoder.encode(&metric_families, &mut buffer).unwrap();
+            String::from_utf8(buffer).unwrap()
+        };
+
+        // Return the response
+        return Ok(Response::new(Body::from(body)));
+    }
+
+    // Otherwise, just return a 404
+    Ok(Response::builder()
+        .status(404)
+        .body(Body::from("Not found"))
+        .unwrap())
+}
+
+/// Bring up an HTTP server that listens for metrics requests
+pub async fn serve_metrics(bind_addr: SocketAddr) {
+    // Set up the server
+    let make_service =
+        make_service_fn(|_| async { Ok::<_, Infallible>(service_fn(handle_request)) });
+    let server = Server::bind(&bind_addr).serve(make_service);
+
+    // Run the server
+    if let Err(e) = server.await {
+        eprintln!("Metrics server error: {e}");
+    }
+}
diff --git a/libs/protomask-metrics/src/lib.rs b/libs/protomask-metrics/src/lib.rs
@@ -5,6 +5,7 @@
 #![allow(clippy::missing_panics_doc)]
 #![allow(clippy::doc_markdown)]
 
+pub mod http;
 pub mod metrics;
 
 #[macro_use]

diff --git a/src/protomask-clat.rs b/src/protomask-clat.rs
@@ -10,7 +10,10 @@ use interproto::protocols::ip::{translate_ipv4_to_ipv6, translate_ipv6_to_ipv4};
 use ipnet::{IpNet, Ipv4Net, Ipv6Net};
 use nix::unistd::Uid;
 use rfc6052::{embed_ipv4_addr_unchecked, extract_ipv4_addr_unchecked};
-use std::io::{Read, Write};
+use std::{
+    io::{Read, Write},
+    net::SocketAddr,
+};
 
 use crate::common::packet_handler::handle_packet;
 
@@ -19,14 +22,18 @@ mod common;
 #[derive(Debug, Parser)]
 #[clap(author, version, about="IPv4 to IPv6 Customer-side transLATor (CLAT)", long_about = None)]
 struct Args {
-    /// RFC6052 IPv6 prefix to encapsulate IPv4 packets within
-    #[clap(long="via", default_value_t = ("64:ff9b::/96").parse().unwrap(), value_parser = parse_network_specific_prefix)]
-    embed_prefix: Ipv6Net,
-
     /// One or more customer-side IPv4 prefixes to allow through CLAT
     #[clap(short = 'c', long = "customer-prefix", required = true)]
     customer_pool: Vec<Ipv4Net>,
 
+    /// Enable prometheus metrics on a given address
+    #[clap(long = "prometheus")]
+    prom_bind_addr: Option<SocketAddr>,
+
+    /// RFC6052 IPv6 prefix to encapsulate IPv4 packets within
+    #[clap(long="via", default_value_t = ("64:ff9b::/96").parse().unwrap(), value_parser = parse_network_specific_prefix)]
+    embed_prefix: Ipv6Net,
+
     /// Explicitly set the interface name to use
     #[clap(short, long, default_value_t = ("clat%d").to_string())]
     interface: String,
@@ -93,6 +100,12 @@ pub async fn main() {
         .unwrap();
     }
 
+    // If we are configured to serve prometheus metrics, start the server
+    if let Some(bind_addr) = args.prom_bind_addr {
+        log::info!("Starting prometheus server on {}", bind_addr);
+        tokio::spawn(protomask_metrics::http::serve_metrics(bind_addr));
+    }
+
     // Translate all incoming packets
     log::info!("Translating packets on {}", tun.name());
     let mut buffer = vec![0u8; 1500];

diff --git a/src/protomask.rs b/src/protomask.rs
@@ -9,7 +9,7 @@ use rfc6052::{embed_ipv4_addr_unchecked, extract_ipv4_addr_unchecked};
 use std::{
     cell::RefCell,
     io::{BufRead, Read, Write},
-    net::{Ipv4Addr, Ipv6Addr},
+    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
     path::PathBuf,
     time::Duration,
 };
@@ -21,17 +21,21 @@ mod common;
 #[derive(Parser)]
 #[clap(author, version, about="Fast and simple NAT64", long_about = None)]
 struct Args {
-    /// RFC6052 IPv6 translation prefix
-    #[clap(long, default_value_t = ("64:ff9b::/96").parse().unwrap(), value_parser = parse_network_specific_prefix)]
-    translation_prefix: Ipv6Net,
-
     #[command(flatten)]
     pool: PoolArgs,
 
     /// A CSV file containing static address mappings from IPv6 to IPv4
     #[clap(long = "static-file")]
     static_file: Option<PathBuf>,
 
+    /// Enable prometheus metrics on a given address
+    #[clap(long = "prometheus")]
+    prom_bind_addr: Option<SocketAddr>,
+
+    /// RFC6052 IPv6 translation prefix
+    #[clap(long, default_value_t = ("64:ff9b::/96").parse().unwrap(), value_parser = parse_network_specific_prefix)]
+    translation_prefix: Ipv6Net,
+
     /// NAT reservation timeout in seconds
     #[clap(long, default_value = "7200")]
     reservation_timeout: u64,
@@ -150,6 +154,12 @@ pub async fn main() {
             .unwrap();
     }
 
+    // If we are configured to serve prometheus metrics, start the server
+    if let Some(bind_addr) = args.prom_bind_addr {
+        log::info!("Starting prometheus server on {}", bind_addr);
+        tokio::spawn(protomask_metrics::http::serve_metrics(bind_addr));
+    }
+
     // Translate all incoming packets
     log::info!("Translating packets on {}", tun.name());
     let mut buffer = vec![0u8; 1500];