feat(su): limits with whitelist

Author: VinceJuliano

servers/su/src/domain/config.rs

@@ -54,6 +54,12 @@ pub struct AoConfig {
     pub assignment: String,
 
     pub cache_url: String,
+
+    pub enable_message_max_size: bool,
+    pub max_message_size: usize,
+    pub max_size_owner_whitelist: Vec<String>,
+    pub max_size_from_owner_whitelist: Vec<String>,
+    pub max_size_from_whitelist: Vec<String>,
 }
 
 fn get_db_dirs() -> (String, String, String, String) {
@@ -218,6 +224,31 @@ impl AoConfig {
             Err(_e) => "https://cache.forward.computer".to_string(),
         };
 
+        let enable_message_max_size = match env::var("ENABLE_MESSAGE_MAX_SIZE") {
+            Ok(val) => val == "true",
+            Err(_e) => false,
+        };
+
+        let max_message_size = match env::var("MAX_MESSAGE_SIZE_BYTES") {
+            Ok(val) => val.parse().unwrap(),
+            Err(_e) => 102400,
+        };
+
+        let max_size_owner_whitelist: Vec<String> = match env::var("MAX_SIZE_OWNER_WHITELIST") {
+            Ok(val) => val.split(',').map(|s| s.trim().to_string()).collect(),
+            Err(_e) => vec![],
+        };
+
+        let max_size_from_whitelist: Vec<String> = match env::var("MAX_SIZE_FROM_WHITELIST") {
+            Ok(val) => val.split(',').map(|s| s.trim().to_string()).collect(),
+            Err(_e) => vec![],
+        };
+
+        let max_size_from_owner_whitelist: Vec<String> = match env::var("MAX_SIZE_FROM_OWNER_WHITELIST") {
+            Ok(val) => val.split(',').map(|s| s.trim().to_string()).collect(),
+            Err(_e) => vec![],
+        };
+
         Ok(AoConfig {
             database_url: env::var("DATABASE_URL")?,
             database_read_url,
@@ -249,7 +280,12 @@ impl AoConfig {
             enable_router_check,
             router_url,
             assignment,
-            cache_url
+            cache_url,
+            enable_message_max_size,
+            max_message_size,
+            max_size_owner_whitelist,
+            max_size_from_owner_whitelist,
+            max_size_from_whitelist
         })
     }
 }
@@ -291,4 +327,19 @@ impl Config for AoConfig {
     fn assignment(&self) -> String {
         self.assignment.clone()
     }
+    fn enable_message_max_size(&self) -> bool {
+        self.enable_message_max_size.clone()
+    }
+    fn max_message_size(&self) -> usize {
+        self.max_message_size.clone()
+    }
+    fn max_size_owner_whitelist(&self) -> Vec<String> {
+        self.max_size_owner_whitelist.clone()
+    }
+    fn max_size_from_owner_whitelist(&self) -> Vec<String> {
+        self.max_size_from_owner_whitelist.clone()
+    }
+    fn max_size_from_whitelist(&self) -> Vec<String> {
+        self.max_size_from_whitelist.clone()
+    }
 }

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

@@ -77,6 +77,11 @@ pub trait Config: Send + Sync {
     fn enable_router_check(&self) -> bool;
     fn router_url(&self) -> String;
     fn assignment(&self) -> String;
+    fn enable_message_max_size(&self) -> bool;
+    fn max_message_size(&self) -> usize;
+    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>;
 }
 
 #[derive(Debug)]

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

@@ -7,6 +7,7 @@ use dotenv::dotenv;
 use serde_json::json;
 use simd_json::to_string as simd_to_string;
 use tokio::sync::Mutex;
+use sha2::{Digest, Sha256};
 
 use super::builder::Builder;
 use super::bytes::{DataBundle, DataItem};
@@ -213,6 +214,75 @@ async fn maybe_recalc_deephashes(deps: Arc<Deps>, process_id: &String) -> Result
     Ok(())
 }
 
+fn limit_message_size(
+  deps: &Arc<Deps>, 
+  input: &Vec<u8>, 
+  data_item: &Option<DataItem>
+) -> Result<(), String> {
+    let enable_message_max_size = deps.config.enable_message_max_size();
+    let max_size_owner_whitelist = deps.config.max_size_owner_whitelist();
+    let max_size_from_owner_whitelist = deps.config.max_size_from_owner_whitelist();
+    let max_size_from_whitelist = deps.config.max_size_from_whitelist();
+    let max_message_size = deps.config.max_message_size();
+
+    if !enable_message_max_size {
+        return Ok(());
+    }
+
+    if let Some(item) = data_item {
+        let tags = item.tags();
+        let from_process = tags.iter().find(
+          |tag| tag.name == "From-Process" || tag.name == "from-process"
+        );
+
+        let owner = item.owner().to_string();
+        let owner_bytes = match base64_url::decode(&owner) {
+            Ok(b) => b,
+            Err(_) => return Err("Unable to decode owner".to_string()),
+        };
+        let mut hasher = Sha256::new();
+        hasher.update(owner_bytes);
+        let result = hasher.finalize();
+        let address_hash = result.to_vec();
+        let address = base64_url::encode(&address_hash);
+
+        println!("Owner: {}", address);
+        println!("From-process: {:?}", from_process);
+
+        match tags
+            .iter()
+            .find(|tag| tag.name == "Type" || tag.name == "type")
+        {
+            Some(type_tag) => match type_tag.value.as_str() {
+                "Process" => return Ok(()),
+                "Message" => (),
+                _ => return Err("Unsupported Type tag value".to_string()),
+            },
+            None => return Err("Type tag not present".to_string()),
+        }
+
+        if max_size_owner_whitelist.contains(&address) {
+            return Ok(());
+        }
+
+        if let Some(fp) = from_process {
+            if max_size_from_owner_whitelist.contains(&address) 
+              && max_size_from_whitelist.contains(&fp.value) {
+                return Ok(());
+            }
+        }
+    }
+
+    if input.len() > max_message_size {
+        return Err(format!(
+            "Message size exceeds maximum of {} bytes",
+            max_message_size
+        ));
+    }
+
+    Ok(())
+}
+
 /*
     This writes a message or process data item,
     it detects which it is creating by the tags.
@@ -250,6 +320,8 @@ pub async fn write_item(
         }
     };
 
+    limit_message_size(&deps, &input, &data_item)?;
+
     deps.logger.log(format!(
         "builder initialized item parsed target - {}",
         &target_id