feat(su): su improvements

Author: VinceJuliano

servers/su/src/domain/config.rs

@@ -60,6 +60,8 @@ pub struct AoConfig {
     pub max_size_owner_whitelist: Vec<String>,
     pub max_size_from_owner_whitelist: Vec<String>,
     pub max_size_from_whitelist: Vec<String>,
+
+    pub ip_whitelist_url: String,
 }
 
 fn get_db_dirs() -> (String, String, String, String) {
@@ -249,6 +251,11 @@ impl AoConfig {
             Err(_e) => vec![],
         };
 
+        let ip_whitelist_url: String = match env::var("IP_WHITELIST_URL") {
+            Ok(val) => val,
+            Err(_e) => "".to_string(),
+        };
+
         Ok(AoConfig {
             database_url: env::var("DATABASE_URL")?,
             database_read_url,
@@ -285,7 +292,8 @@ impl AoConfig {
             max_message_size,
             max_size_owner_whitelist,
             max_size_from_owner_whitelist,
-            max_size_from_whitelist
+            max_size_from_whitelist,
+            ip_whitelist_url
         })
     }
 }
@@ -342,4 +350,7 @@ impl Config for AoConfig {
     fn max_size_from_whitelist(&self) -> Vec<String> {
         self.max_size_from_whitelist.clone()
     }
+    fn ip_whitelist_url(&self) -> String {
+        self.ip_whitelist_url.clone()
+    }
 }

servers/su/src/domain/core/dal.rs

@@ -82,6 +82,7 @@ pub trait Config: Send + Sync {
     fn max_size_owner_whitelist(&self) -> Vec<String>;
     fn max_size_from_owner_whitelist(&self) -> Vec<String>;
     fn max_size_from_whitelist(&self) -> Vec<String>;
+    fn ip_whitelist_url(&self) -> String;
 }
 
 #[derive(Debug)]

servers/su/src/main.rs

@@ -1,7 +1,9 @@
 use std::env;
 use std::io::{self, Error, ErrorKind};
 use std::sync::Arc;
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::time::{SystemTime, UNIX_EPOCH, Duration};
+use std::collections::HashSet;
+use std::sync::RwLock;
 
 use actix_cors::Cors;
 use actix_web::{
@@ -11,9 +13,13 @@ use actix_web::{
 
 use serde::Deserialize;
 use serde_json::json;
+use tokio::time::interval;
+use reqwest;
 
 use su::domain::{flows, init_deps, router, Deps, PromMetrics};
 
+type IpWhitelist = Arc<RwLock<HashSet<String>>>;
+
 #[derive(Deserialize)]
 struct FromTo {
     from: Option<String>,
@@ -61,6 +67,40 @@ fn err_response(err: String) -> HttpResponse {
         .body(error_json.to_string())
 }
 
+fn get_client_ip(req: &HttpRequest) -> Option<String> {
+    req.connection_info()
+        .realip_remote_addr()
+        .map(|ip| ip.to_string())
+}
+
+fn is_ip_allowed(ip: &str, whitelist: &IpWhitelist) -> bool {
+    println!("Checking if IP {} is allowed", ip);
+    match whitelist.read() {
+        Ok(ips) => ips.contains(ip),
+        Err(_) => false,
+    }
+}
+
+async fn fetch_ip_whitelist(url: String) -> Result<HashSet<String>, reqwest::Error> {
+    let response = reqwest::get(url).await?;
+    let ips: Vec<String> = response.json().await?;
+    Ok(ips.into_iter().collect())
+}
+
+async fn update_ip_whitelist_loop(ip_whitelist: IpWhitelist, whitelist_url: String) {
+    let mut interval = interval(Duration::from_secs(30)); 
+
+    loop {
+        interval.tick().await;
+
+        if let Ok(new_ips) = fetch_ip_whitelist(whitelist_url.clone()).await {
+            if let Ok(mut ips) = ip_whitelist.write() {
+                *ips = new_ips;
+            }
+        }
+    }
+}
+
 async fn base(
     data: web::Data<AppState>,
     query_params: web::Query<ProcessId>,
@@ -128,6 +168,16 @@ async fn main_post_route(
         return HttpResponse::ServiceUnavailable()
             .json(json!({"error": "Server is warming up. Please try again later."}));
     }
+
+    if let Some(client_ip) = get_client_ip(&req) {
+        if !is_ip_allowed(&client_ip, &data.ip_whitelist) {
+            return HttpResponse::Forbidden()
+                .json(json!({"error": "Access denied"}));
+        }
+    } else {
+        return HttpResponse::BadRequest()
+            .json(json!({"error": "Access denied"}));
+    }
     match router::redirect_data_item(
         data.deps.clone(),
         req_body.to_vec(),
@@ -281,6 +331,7 @@ struct AppState {
     deps: Arc<Deps>,
     metrics: Arc<PromMetrics>,
     startup_time: u64,
+    ip_whitelist: IpWhitelist,
 }
 
 #[actix_web::main]
@@ -311,13 +362,32 @@ async fn main() -> io::Result<()> {
         .as_secs();
 
     let (deps, metrics) = init_deps(mode).await;
+
+    let run_deps = deps.clone();
+
+    let ip_whitelist: IpWhitelist = Arc::new(RwLock::new(HashSet::new()));
+
+    match fetch_ip_whitelist(run_deps.config.ip_whitelist_url()).await {
+        Ok(initial_ips) => {
+            if let Ok(mut ips) = ip_whitelist.write() {
+                *ips = initial_ips;
+            }
+        }
+        Err(_) => {}
+    }
+
     let app_state = web::Data::new(AppState {
         deps,
         metrics,
         startup_time,
+        ip_whitelist: ip_whitelist.clone(),
     });
 
-    let run_deps = app_state.deps.clone();
+    let whitelist_clone = ip_whitelist.clone();
+    let whitelist_url = run_deps.config.ip_whitelist_url();
+    tokio::spawn(async move {
+        update_ip_whitelist_loop(whitelist_clone, whitelist_url).await;
+    });
 
     if run_deps.config.mode() == "router" {
         match router::init_schedulers(run_deps.clone()).await {