From 121669de3883d37801a92814cd5113d7ebd15018 Mon Sep 17 00:00:00 2001 From: Michael Furmur Date: Tue, 26 Mar 2024 02:35:24 +0200 Subject: [PATCH 1/4] add country geoip lookup for endpoint addresses * add CLI option 'geoip_path' to specify Country MMDB path * add attribute 'endpoint_country' to the 'wireguard_peer_endpoint' metric if 'geoip_path' was specified --- Cargo.toml | 1 + src/cli.rs | 5 ++++- src/main.rs | 9 ++++++-- src/metrics.rs | 56 +++++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 61 insertions(+), 10 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 1702ea8..2bfff46 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ tracing = "0.1.25" tracing-subscriber = "0.3.1" base64 = "0.13.0" time = { version = "0.3.5", features = ["local-offset"] } +maxminddb = "0.24.0" [build-dependencies] clap = "3.0.0-beta.5" diff --git a/src/cli.rs b/src/cli.rs index 17a09e9..c7ba4c6 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, env, net::IpAddr, str::FromStr}; +use std::{collections::HashMap, env, net::IpAddr, path::PathBuf, str::FromStr}; use clap::{crate_authors, Parser}; @@ -27,6 +27,9 @@ pub struct Args { /// Add an alias for a given public key in the form of 'pubkey:alias' (separate multiple with commas) #[clap(long, short, value_delimiter = ',', multiple_occurrences = true)] pub alias: Vec, + /// Do geoip lookup using Country MMDB from the PATH for 'endpoint_ip' attribute in the 'wireguard_peer_endpoint' metric and add attribute 'endpoint_country' + #[clap(short, long, value_name = "PATH")] + pub geoip_path: Option, } impl Args { diff --git a/src/main.rs b/src/main.rs index 4be1ae6..9d7b965 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,6 +12,7 @@ use color_eyre::{ config::{HookBuilder, Theme}, eyre::Result, }; +use maxminddb; use prometheus::{IntGauge, Registry}; use prometheus_hyper::Server; use tokio::time::{Duration, Instant}; @@ -57,11 +58,15 @@ async fn try_main(args: Args) -> Result<()> { let aliases = args.aliases(); + let maxminddb_reader = args.geoip_path.as_ref().map_or(None, |path| { + Some(unwrap_or_exit!(maxminddb::Reader::open_readfile(path))) + }); + let running = Arc::new(AtomicBool::new(true)); info!("Registering metrics..."); let registry = Arc::new(Registry::new()); - let mut metrics = unwrap_or_exit!(Metrics::new(®istry)); + let mut metrics = unwrap_or_exit!(Metrics::new(®istry, &maxminddb_reader)); let scrape_duration = unwrap_or_exit!(IntGauge::new( "wireguard_scrape_duration_milliseconds", "Duration in milliseconds of the scrape", @@ -89,7 +94,7 @@ async fn try_main(args: Args) -> Result<()> { debug!("Updating metrics..."); metrics - .update(&WireguardState::scrape(&aliases).await?) + .update(&WireguardState::scrape(&aliases).await?, &maxminddb_reader) .await; let after = Instant::now(); diff --git a/src/metrics.rs b/src/metrics.rs index 24207c6..4eadcdb 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -1,4 +1,5 @@ use color_eyre::eyre::Result; +use maxminddb::geoip2; use prometheus::{IntCounterVec, IntGaugeVec, Opts, Registry}; use time::OffsetDateTime; use tracing::{debug, trace}; @@ -15,7 +16,10 @@ pub struct Metrics { } impl Metrics { - pub fn new(r: &Registry) -> Result { + pub fn new( + r: &Registry, + maxminddb_reader: &Option>>, + ) -> Result { trace!("Metrics::new"); let interfaces_total = IntGaugeVec::new( @@ -33,7 +37,16 @@ impl Metrics { let peer_endpoint = IntGaugeVec::new( Opts::new("wireguard_peer_endpoint", "Peers info. static value"), - &["interface", "endpoint_ip", "peer", "alias"], + match maxminddb_reader { + Some(_) => &[ + "interface", + "endpoint_ip", + "endpoint_country", + "peer", + "alias", + ], + None => &["interface", "endpoint_ip", "peer", "alias"], + }, )?; let bytes_total = IntCounterVec::new( @@ -78,7 +91,11 @@ impl Metrics { }) } - pub async fn update(&mut self, state: &WireguardState) { + pub async fn update( + &mut self, + state: &WireguardState, + maxminddb_reader: &Option>>, + ) { let it = self.interfaces_total.with_label_values(&[]); it.set(state.interfaces.len() as i64); @@ -117,8 +134,32 @@ impl Metrics { assert!(p.interface < state.interfaces.len()); if let Some(endpoint) = p.endpoint { - self.peer_endpoint - .with_label_values(&[ + match maxminddb_reader { + Some(reader) => { + let endpoint_country = match reader.lookup::(endpoint.ip()) + { + Ok(reader_response) => { + reader_response.country.map_or_else(String::new, |c| { + c.iso_code + .map(|code| code.to_string()) + .unwrap_or_else(String::new) + }) + } + _ => String::new(), + }; + + self.peer_endpoint.with_label_values(&[ + &state.interfaces[p.interface], + &endpoint.ip().to_string(), + &endpoint_country.to_string(), + &p.pubkey, + &p.alias + .as_ref() + .map(ToOwned::to_owned) + .unwrap_or_else(String::new), + ]) + } + None => self.peer_endpoint.with_label_values(&[ &state.interfaces[p.interface], &endpoint.ip().to_string(), &p.pubkey, @@ -126,8 +167,9 @@ impl Metrics { .as_ref() .map(ToOwned::to_owned) .unwrap_or_else(String::new), - ]) - .set(1); + ]), + } + .set(1); }; let pbtt = self.peer_bytes_total.with_label_values(&[ From 736f1ae8ff0575650141087da16f9ce859ba3d07 Mon Sep 17 00:00:00 2001 From: Michael Furmur Date: Tue, 9 Apr 2024 02:08:20 +0300 Subject: [PATCH 2/4] store maxminddb_reader in the Metrics struct --- src/main.rs | 4 ++-- src/metrics.rs | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/main.rs b/src/main.rs index 9d7b965..0b989f5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -66,7 +66,7 @@ async fn try_main(args: Args) -> Result<()> { info!("Registering metrics..."); let registry = Arc::new(Registry::new()); - let mut metrics = unwrap_or_exit!(Metrics::new(®istry, &maxminddb_reader)); + let mut metrics = unwrap_or_exit!(Metrics::new(®istry, maxminddb_reader)); let scrape_duration = unwrap_or_exit!(IntGauge::new( "wireguard_scrape_duration_milliseconds", "Duration in milliseconds of the scrape", @@ -94,7 +94,7 @@ async fn try_main(args: Args) -> Result<()> { debug!("Updating metrics..."); metrics - .update(&WireguardState::scrape(&aliases).await?, &maxminddb_reader) + .update(&WireguardState::scrape(&aliases).await?) .await; let after = Instant::now(); diff --git a/src/metrics.rs b/src/metrics.rs index 4eadcdb..95514d3 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -13,12 +13,13 @@ pub struct Metrics { bytes_total: IntCounterVec, peer_bytes_total: IntCounterVec, duration_since_latest_handshake: IntGaugeVec, + maxminddb_reader: Option>> } impl Metrics { pub fn new( r: &Registry, - maxminddb_reader: &Option>>, + maxminddb_reader: Option>>, ) -> Result { trace!("Metrics::new"); @@ -88,13 +89,13 @@ impl Metrics { peer_endpoint, peer_bytes_total, duration_since_latest_handshake, + maxminddb_reader }) } pub async fn update( &mut self, - state: &WireguardState, - maxminddb_reader: &Option>>, + state: &WireguardState ) { let it = self.interfaces_total.with_label_values(&[]); it.set(state.interfaces.len() as i64); @@ -134,7 +135,7 @@ impl Metrics { assert!(p.interface < state.interfaces.len()); if let Some(endpoint) = p.endpoint { - match maxminddb_reader { + match &self.maxminddb_reader { Some(reader) => { let endpoint_country = match reader.lookup::(endpoint.ip()) { From 3713673adc0b1391076086eb41c4b2d14403b43e Mon Sep 17 00:00:00 2001 From: Michael Furmur Date: Tue, 9 Apr 2024 02:10:28 +0300 Subject: [PATCH 3/4] bump MSRV to 1.58.1 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1f7e6fe..4ff90e3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,7 +11,7 @@ jobs: rust: - stable - nightly - - 1.56.1 # MSRV + - 1.58.1 # MSRV steps: - uses: actions/checkout@v2 From 6814784e13ad24f77dc4ffb8c69681c71a466a8e Mon Sep 17 00:00:00 2001 From: Michael Furmur Date: Tue, 9 Apr 2024 02:26:50 +0300 Subject: [PATCH 4/4] fix formatting --- src/metrics.rs | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/src/metrics.rs b/src/metrics.rs index 95514d3..ff8185f 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -13,14 +13,11 @@ pub struct Metrics { bytes_total: IntCounterVec, peer_bytes_total: IntCounterVec, duration_since_latest_handshake: IntGaugeVec, - maxminddb_reader: Option>> + maxminddb_reader: Option>>, } impl Metrics { - pub fn new( - r: &Registry, - maxminddb_reader: Option>>, - ) -> Result { + pub fn new(r: &Registry, maxminddb_reader: Option>>) -> Result { trace!("Metrics::new"); let interfaces_total = IntGaugeVec::new( @@ -89,14 +86,11 @@ impl Metrics { peer_endpoint, peer_bytes_total, duration_since_latest_handshake, - maxminddb_reader + maxminddb_reader, }) } - pub async fn update( - &mut self, - state: &WireguardState - ) { + pub async fn update(&mut self, state: &WireguardState) { let it = self.interfaces_total.with_label_values(&[]); it.set(state.interfaces.len() as i64);