[Management] Implement assigning regions to servers on registation.

Author: davidv1992

Cargo.lock

@@ -43,6 +43,9 @@ COPY --from=builder /build/artifacts/nts-pool-management /usr/local/bin/nts-pool
 COPY --from=builder /build/artifacts/nts-pool-monitor /usr/local/bin/nts-pool-monitor
 COPY --from=builder /build/nts-pool-management/assets /opt/nts-pool-management/assets
 
+# Temporarily copy test geodb, until we actually implement proper loading of the geolocation database.
+COPY --from=builder /build/nts-pool-ke/testdata/GeoLite2-Country-Test.mmdb /opt/nts-pool-management/GeoLite2-Country-Test.mmdb
+
 # Set a default assets directory
 ENV NTSPOOL_ASSETS_DIR=/opt/nts-pool-management/assets
 

Dockerfile

@@ -79,6 +79,7 @@ services:
       DATABASE_URL: "postgres://nts-pool@psql:5432/nts-pool"
       NTSPOOL_DATABASE_RUN_MIGRATIONS: "false"
       NTSPOOL_ASSETS_DIR: "/app/nts-pool-management/assets"
+      NTSPOOL_GEOLOCATION_DB: "/app/nts-pool-ke/testdata/GeoLite2-Country-Test.mmdb"
       NTSPOOL_JWT_SECRET: "a-string-secret-at-least-256-bits-long"
       NTSPOOL_COOKIE_SECRET: "another-string-secret-at-least-256-bits-long"
       NTSPOOL_SMTP_URL: "smtp://mailcrab:1025"

compose.yml

@@ -101,3 +101,5 @@ spec:
             value: "1"
           - name: RESTART_BUMP
             value: "{{ default 1 .Values.management.restartBump}}"
+          - name: NTSPOOL_GEOLOCATION_DB
+            value: "/opt/nts-pool-management/GeoLite2-Country-Test.mmdb" # Temporary, until proper loading of geolocation database is implemented

helm-charts/nts-pool/templates/management.yaml

@@ -15,7 +15,10 @@ askama = { workspace = true }
 axum = { workspace = true, features = ["macros"] }
 axum-extra = { workspace = true, features = ["cookie", "cookie-private", "cookie-key-expansion", "typed-header"] }
 headers.workspace = true
+maxminddb.workspace = true
+notify.workspace = true
 nts-pool-shared = { workspace = true, features = ["sqlx"] }
+phf.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
 serde_qs = { workspace = true }

…ede2cc53928bdf0b2521091dc38f0c75afe.json …a5eae6fd24d3a9715970fc7923f6b4cae00.jsonnts-pool-management/.sqlx/query-8f1990990f9c7bc98787dfdc41df5ede2cc53928bdf0b2521091dc38f0c75afe.json renamed to nts-pool-management/.sqlx/query-41819cf933c55bdc06b40379f6499a5eae6fd24d3a9715970fc7923f6b4cae00.json nts-pool-management/.sqlx/query-8f1990990f9c7bc98787dfdc41df5ede2cc53928bdf0b2521091dc38f0c75afe.json renamed to nts-pool-management/.sqlx/query-41819cf933c55bdc06b40379f6499a5eae6fd24d3a9715970fc7923f6b4cae00.json

@@ -8,6 +8,7 @@ pub struct AppConfig {
     // Base website configuration
     pub base_url: BaseUrl,
     pub assets_path: String,
+    pub geolocation_db: std::path::PathBuf,
     // Database config
     pub database_url: String,
     pub database_run_migrations: RunDatabaseMigrations,
@@ -42,6 +43,9 @@ impl AppConfig {
                 .wrap_err("NTSPOOL_BASE_URL not set in environment")?,
         );
         let assets_path = std::env::var("NTSPOOL_ASSETS_DIR").unwrap_or("./assets".into());
+        let geolocation_db = std::env::var("NTSPOOL_GEOLOCATION_DB")
+            .wrap_err("NTSPOOL_GEOLOCATION_DB not set in environment")?
+            .into();
 
         // Database configuration
         let database_url = std::env::var("NTSPOOL_DATABASE_URL")
@@ -108,6 +112,7 @@ impl AppConfig {
         Ok(Self {
             base_url,
             assets_path,
+            geolocation_db,
             database_url,
             database_run_migrations,
             jwt_secret,
@@ -136,6 +141,7 @@ impl Default for AppConfig {
             cookie_secret: "UNSAFE_SECRET".into(),
             mail_from_address: "noreply@example.com".into(),
             mail_smtp_url: "smtp://localhost:25".into(),
+            geolocation_db: "../nts-pool-ke/testdata/GeoLite2-Country-Test.mmdb".into(),
             assets_path: "./assets".into(),
             monitor_result_batchsize: 4,
             monitor_result_batchtime: Duration::from_secs(60),

nts-pool-management/Cargo.toml

@@ -1,12 +1,18 @@
-use std::str::FromStr;
+use std::{
+    str::FromStr,
+    sync::{Arc, RwLock},
+};
 
 use axum::{extract::FromRef, middleware};
 use axum_extra::extract::cookie;
+use eyre::Context;
+use notify::Watcher;
 use sqlx::PgPool;
 use tracing::info;
 
 use crate::{
     common::config::RunDatabaseMigrations,
+    common::error::AppError,
     config::AppConfig,
     email::{MailTransport, Mailer},
 };
@@ -29,6 +35,7 @@ pub struct AppState {
     private_cookie_key: cookie::Key,
     mailer: Mailer,
     config: AppConfig,
+    geodb: Arc<RwLock<Arc<maxminddb::Reader<Vec<u8>>>>>,
 }
 
 pub trait DbConnLike<'a>:
@@ -77,6 +84,61 @@ pub async fn pool_conn(
     }
 }
 
+async fn load_geodb(
+    geodb_path: &std::path::Path,
+) -> Result<Arc<maxminddb::Reader<Vec<u8>>>, AppError> {
+    let geodb_raw = tokio::fs::read(geodb_path)
+        .await
+        .wrap_err("Could not load geolocation database from disk")?;
+
+    Ok(Arc::new(
+        maxminddb::Reader::from_source(geodb_raw).wrap_err("Invalid geolocation database")?,
+    ))
+}
+
+async fn manage_geodb(
+    geodb_path: impl AsRef<std::path::Path> + Send + 'static,
+) -> Result<Arc<RwLock<Arc<maxminddb::Reader<Vec<u8>>>>>, AppError> {
+    let (change_sender, mut change_receiver) = tokio::sync::mpsc::unbounded_channel::<()>();
+    // Use a poll watcher here as INotify can be unreliable in many ways and I don't want to deal with that.
+    let mut watcher = notify::poll::PollWatcher::new(
+        move |event: notify::Result<notify::Event>| {
+            if event.is_ok() {
+                let _ = change_sender.send(());
+            }
+        },
+        notify::Config::default()
+            .with_poll_interval(std::time::Duration::from_secs(60))
+            .with_compare_contents(true),
+    )
+    .wrap_err("Could not setup watcher for changes in geolocation database")?;
+
+    watcher
+        .watch(geodb_path.as_ref(), notify::RecursiveMode::NonRecursive)
+        .wrap_err("Could not watch geolocation database for changes")?;
+
+    let geodb = Arc::new(RwLock::new(load_geodb(geodb_path.as_ref()).await?));
+    let geodb_cloned = geodb.clone();
+
+    tokio::spawn(async move {
+        // keep the watcher alive
+        let _w = watcher;
+        loop {
+            change_receiver.recv().await;
+            match load_geodb(geodb_path.as_ref()).await {
+                Ok(new_geodb) => {
+                    *geodb.write().unwrap() = new_geodb;
+                }
+                Err(e) => {
+                    tracing::error!("Could not refresh geolocation database: {e}");
+                }
+            }
+        }
+    });
+
+    Ok(geodb_cloned)
+}
+
 #[tokio::main]
 async fn main() {
     let config = AppConfig::from_env().expect("Failed to load configuration");
@@ -122,6 +184,10 @@ async fn main() {
 
     let serve_dir_service = tower_http::services::ServeDir::new(&config.assets_path);
 
+    let geodb = manage_geodb(config.geolocation_db.clone())
+        .await
+        .expect("Unable to initialize geolocation database");
+
     // construct the application state
     let state = AppState {
         db,
@@ -130,6 +196,7 @@ async fn main() {
         private_cookie_key,
         mailer,
         config,
+        geodb,
     };
 
     // setup routes

nts-pool-management/src/common/config.rs

@@ -1,4 +1,7 @@
+use std::collections::HashSet;
+
 use eyre::Context;
+use phf::phf_map;
 use serde::Deserialize;
 
 use crate::{
@@ -10,6 +13,16 @@ use crate::{
     },
 };
 
+static CONTINENTS: phf::Map<&'static str, &'static str> = phf_map! {
+    "AF" => "AFRICA",
+    "AN" => "ANTARCTICA",
+    "AS" => "ASIA",
+    "EU" => "EUROPE",
+    "NA" => "NORTH-AMERICA",
+    "OC" => "OCEANIA",
+    "SA" => "SOUTH_AMERICA",
+};
+
 uuid!(TimeSourceId);
 
 #[derive(Debug, Clone, Deserialize, sqlx::FromRow)]
@@ -74,19 +87,66 @@ impl TryFrom<NewTimeSourceForm> for NewTimeSource {
     }
 }
 
+pub async fn infer_regions(
+    hostname: impl AsRef<str>,
+    geodb: &maxminddb::Reader<impl AsRef<[u8]>>,
+) -> Vec<String> {
+    // Note: port doesn't matter but is needed for the lookup_host interface
+    let addresses = match tokio::net::lookup_host((hostname.as_ref(), 4460)).await {
+        Ok(addresses) => addresses,
+        Err(e) => {
+            if e.raw_os_error().is_some() {
+                // Definitely an issue
+                tracing::error!("Could not resolve hostname of time source: {e}");
+            }
+            return vec![];
+        }
+    };
+
+    let mut result = HashSet::new();
+    for addr in addresses {
+        let Some(lookup) = (match geodb.lookup::<maxminddb::geoip2::Country>(addr.ip()) {
+            Ok(lookup) => lookup,
+            Err(e) => {
+                tracing::error!("Failure during geoip lookup: {e}");
+                None
+            }
+        }) else {
+            continue;
+        };
+
+        if let Some(continent) = lookup
+            .continent
+            .and_then(|v| v.code)
+            .and_then(|c| CONTINENTS.get(c))
+        {
+            result.insert((*continent).to_owned());
+        }
+        if let Some(country) = lookup.country.and_then(|v| v.iso_code) {
+            result.insert(country.to_owned());
+        }
+    }
+
+    result.into_iter().collect()
+}
+
 pub async fn create(
     conn: impl DbConnLike<'_>,
     owner: UserId,
     new_time_source: NewTimeSource,
+    geodb: &maxminddb::Reader<impl AsRef<[u8]>>,
 ) -> Result<(), sqlx::Error> {
+    let regions = infer_regions(&new_time_source.hostname, geodb).await;
+
     sqlx::query!(
         r#"
-            INSERT INTO time_sources (owner, hostname, port)
-            VALUES ($1, $2, $3)
+            INSERT INTO time_sources (owner, hostname, port, countries)
+            VALUES ($1, $2, $3, $4)
         "#,
         owner as _,
         new_time_source.hostname,
         new_time_source.port as _,
+        regions.as_slice(),
     )
     .execute(conn)
     .await?;

nts-pool-management/src/main.rs

@@ -51,10 +51,12 @@ pub async fn create_time_source(
     flash: FlashMessageService,
     Form(new_time_source): Form<NewTimeSourceForm>,
 ) -> Result<impl IntoResponse, AppError> {
-    let flash = match time_source::create(&state.db, user.id, new_time_source.try_into()?).await {
-        Ok(_) => flash.success("Time source added successfully".to_string()),
-        Err(_) => flash.error("Could not add time source".to_string()),
-    };
+    let geodb = state.geodb.read().unwrap().clone();
+    let flash =
+        match time_source::create(&state.db, user.id, new_time_source.try_into()?, &geodb).await {
+            Ok(_) => flash.success("Time source added successfully".to_string()),
+            Err(_) => flash.error("Could not add time source".to_string()),
+        };
 
     Ok((flash, Redirect::to(TIME_SOURCES_ENDPOINT)))
 }