Merge pull request #6 from furmur/feature/mmdb_country_lookup

add country geoip lookup for endpoint addresses

.github/workflows/ci.yml

@@ -19,7 +19,7 @@ jobs:
         rust:
           - stable
           - nightly
-          - 1.56.1 # MSRV
+          - 1.58.1 # MSRV
 
     steps:
       - uses: actions/checkout@v2

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"

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<Alias>,
+    /// 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<PathBuf>,
 }
 
 impl Args {

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(&registry));
+    let mut metrics = unwrap_or_exit!(Metrics::new(&registry, maxminddb_reader));
     let scrape_duration = unwrap_or_exit!(IntGauge::new(
         "wireguard_scrape_duration_milliseconds",
         "Duration in milliseconds of the scrape",

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};
@@ -12,10 +13,11 @@ pub struct Metrics {
     bytes_total: IntCounterVec,
     peer_bytes_total: IntCounterVec,
     duration_since_latest_handshake: IntGaugeVec,
+    maxminddb_reader: Option<maxminddb::Reader<Vec<u8>>>,
 }
 
 impl Metrics {
-    pub fn new(r: &Registry) -> Result<Self> {
+    pub fn new(r: &Registry, maxminddb_reader: Option<maxminddb::Reader<Vec<u8>>>) -> Result<Self> {
         trace!("Metrics::new");
 
         let interfaces_total = IntGaugeVec::new(
@@ -33,7 +35,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(
@@ -75,6 +86,7 @@ impl Metrics {
             peer_endpoint,
             peer_bytes_total,
             duration_since_latest_handshake,
+            maxminddb_reader,
         })
     }
 
@@ -117,17 +129,42 @@ impl Metrics {
             assert!(p.interface < state.interfaces.len());
 
             if let Some(endpoint) = p.endpoint {
-                self.peer_endpoint
-                    .with_label_values(&[
+                match &self.maxminddb_reader {
+                    Some(reader) => {
+                        let endpoint_country = match reader.lookup::<geoip2::Country>(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,
                         &p.alias
                             .as_ref()
                             .map(ToOwned::to_owned)
                             .unwrap_or_else(String::new),
-                    ])
-                    .set(1);
+                    ]),
+                }
+                .set(1);
             };
 
             let pbtt = self.peer_bytes_total.with_label_values(&[