FIX - Snapshot Memory Pressure: Both save_to_cold and

     start_snapshot_task clone the entire database state to
     serialize it. This effectively doubles memory consumption
     during the operation and can lead to OOM on large datasets.

Author: Istrate Andrei-Eduard

src/network/broadcaster.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::env;
 use std::time::Duration;
 use reqwest;
-use actix_web::{web, HttpResponse, Responder};
+use actix_web::{HttpRequest, HttpResponse, Responder, web};
 use serde_json::json;
 use crate::storage::engine::{
     current_timestamp, ClusterData, NodeInfo, NodeStatus,
@@ -181,18 +181,50 @@ pub async fn heartbeat() -> impl Responder {
     }))
 }
 
+// Define a safe upper bound for your cluster size
+const MAX_CLUSTER_SIZE: usize = 100; 
+
 /// Process membership updates received from other nodes
 pub async fn update_membership(
+    req: HttpRequest, // Added to inspect headers
     cluster_data: web::Data<ClusterData>,
     payload: web::Json<HashMap<String, NodeInfo>>,
 ) -> impl Responder {
+    // 1. Authentication Check
+    // In production, load this once at startup rather than on every request
+    let expected_token = env::var("CLUSTER_SECRET")
+        .unwrap_or_else(|_| "default_insecure_secret".to_string());
+    
+    let is_authenticated = req.headers()
+        .get("X-Cluster-Token")
+        .and_then(|h| h.to_str().ok())
+        .map_or(false, |token| token == expected_token);
+
+    if !is_authenticated {
+        return HttpResponse::Unauthorized().json(json!({
+            "error": "Unauthorized: Invalid or missing cluster token"
+        }));
+    }
+
     let mut local_guard = cluster_data.nodes.write().await;
     let incoming = payload.into_inner();
     let mut updated = false;
 
     for (node, info) in incoming.into_iter() {
-        // For nodes not in our membership, add them
+        // 2. Input Validation (Basic check to ensure it looks like a URL)
+        if !node.starts_with("http://") && !node.starts_with("https://") {
+            continue; // Ignore malformed node addresses
+        }
+
+        // For nodes not in our membership, add them safely
         if !local_guard.contains_key(&node) {
+            // 3. Hard Limit Check
+            if local_guard.len() >= MAX_CLUSTER_SIZE {
+                // Silently ignore or log a warning. Avoid heavy logging to prevent log-flooding DoS.
+                eprintln!("Warning: Cluster at max capacity ({}). Ignoring new node: {}", MAX_CLUSTER_SIZE, node);
+                continue; 
+            }
+
             local_guard.insert(node, info);
             updated = true;
             continue;
@@ -212,7 +244,6 @@ pub async fn update_membership(
 
     HttpResponse::Ok().finish()
 }
-
 /// Get the current cluster membership state
 pub async fn get_membership(
     cluster_data: web::Data<ClusterData>,

src/storage/engine.rs

@@ -208,7 +208,7 @@ pub async fn get_value(
     // 2️⃣ Authenticate external user via JWT
     let user = match extract_user_from_token(&req) {
         Ok(u) => u,
-        Err(resp) => return HttpResponse::Unauthorized().body("You are not owning this record"),
+        Err(_) => return HttpResponse::Unauthorized().body("You are not owning this record"),
     };
 
     // 3️⃣ Request replicas (including local node) to find the latest version

src/storage/persistance.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 use std::fs::{create_dir_all, File, read_dir};
-use std::io::{self, BufRead, BufReader, Read, Write};
+use std::io::{self, BufRead, BufReader,  Write};
 use std::path::Path;
 use std::process;
 use std::time::Duration;
@@ -108,59 +108,85 @@ pub struct WalEntry {
 // }
 
 /// Saves the current in‑memory state (snapshot) for each table to disk.
-/// Also clears the WAL after snapshot.
+/// Processes one table at a time, and encrypts data in 1MB chunks to prevent OOM.
 pub async fn save_to_cold(state: web::Data<AppState>) -> io::Result<()> {
-    // Clone the store so that we can release the lock.
-    let store_snapshot = {
+    let table_names: Vec<String> = {
         let store = state.store.read().await;
-        store.clone()
+        store.keys().cloned().collect()
     };
-    let state_clone = state.clone();
-    task::spawn_blocking(move || {
-        for (table_name, table_data) in store_snapshot.into_iter() {
-            let folder_path = Path::new(state_clone.base_dir).join(&table_name);
+
+    for table_name in table_names {
+        // Grab the read lock just long enough to clone this specific table
+        let table_data = {
+            let store = state.store.read().await;
+            match store.get(&table_name) {
+                Some(data) => data.clone(),
+                None => continue,
+            }
+        };
+
+        let state_clone = state.clone();
+        let t_name = table_name.clone();
+
+        task::spawn_blocking(move || -> io::Result<()> {
+            let folder_path = Path::new(state_clone.base_dir).join(&t_name);
             create_dir_all(&folder_path)?;
             let file_path = folder_path.join("storage.db");
             let mut file = File::create(&file_path)?;
 
-            // build your plaintext
-            let mut plain_data = String::new();
-            for (key, versioned_value) in table_data.iter() {
-                plain_data.push_str(&format!(
-                    "{} = {}\n",
-                    key,
-                    serde_json::to_string(versioned_value)?
-                ));
-            }
+            let mut plain_buffer = String::new();
+            // Define a chunk size limit (e.g., 1 MB)
+            let chunk_size_limit = 1024 * 1024; 
+
+            // Use into_iter() to consume the cloned table and free memory as we go
+            for (key, versioned_value) in table_data.into_iter() {
+                let json_str = serde_json::to_string(&versioned_value)
+                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+                
+                plain_buffer.push_str(&format!("{} = {}\n", key, json_str));
 
-            // ✂️ unwrap the Result<String,_> here and map to io::Error
-            let encrypted_data = encrypt(&plain_data)
-                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+                // If our buffer hits the limit, encrypt, write, and clear
+                if plain_buffer.len() >= chunk_size_limit {
+                    let encrypted_chunk = encrypt(&plain_buffer)
+                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+                    
+                    // Write the base64 string followed by a newline
+                    writeln!(file, "{}", encrypted_chunk)?;
+                    plain_buffer.clear();
+                }
+            }
 
-            // write the actual String, not a Result<_,_>
-            file.write_all(encrypted_data.as_bytes())?;
+            // Flush any remaining records in the buffer
+            if !plain_buffer.is_empty() {
+                let encrypted_chunk = encrypt(&plain_buffer)
+                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+                writeln!(file, "{}", encrypted_chunk)?;
+            }
 
-            // Clear the WAL
+            // Clear the WAL for this table
             let wal_path = folder_path.join("wal.log");
             File::create(&wal_path)?;
-        }
-        Ok::<(), io::Error>(())
-    })
+
+            Ok(())
+        })
         .await??;
+    }
 
     Ok(())
 }
-
 /// Loads all tables from disk by replaying the snapshot and WAL.
 pub async fn load_all_tables(state: &web::Data<AppState>) -> io::Result<()> {
     let state_cloned = state.clone(); // clone to ensure 'static lifetime in blocking task
+    
     let new_store = task::spawn_blocking(move || {
-        let base_path = Path::new(state_cloned.base_dir);
+        let base_path = Path::new(&state_cloned.base_dir);
         let mut store: HashMap<String, HashMap<String, VersionedValue>> = HashMap::new();
+        
         if base_path.exists() && base_path.is_dir() {
             for entry in read_dir(base_path)? {
                 let entry = entry?;
                 let table_folder = entry.path();
+                
                 if table_folder.is_dir() {
                     let table_name = match table_folder.file_name() {
                         Some(name) => name.to_string_lossy().to_string(),
@@ -169,44 +195,56 @@ pub async fn load_all_tables(state: &web::Data<AppState>) -> io::Result<()> {
 
                     let mut table_data: HashMap<String, VersionedValue> = HashMap::new();
                     let snapshot_path = table_folder.join("storage.db");
+                    
+                    // --- 1. Load Snapshot via Chunked Decryption ---
                     if snapshot_path.exists() {
-                        let mut file = File::open(&snapshot_path)?;
-                        let mut encrypted_content = String::new();
-                        file.read_to_string(&mut encrypted_content)?;
-                        match decrypt(&encrypted_content) {
-                            Ok(decrypted_content) => {
-                                for line in decrypted_content.lines() {
-                                    if let Some((key, json_str)) = line.split_once('=') {
-                                        let key = key.trim().to_string();
-                                        let json_str = json_str.trim();
-                                        if let Ok(value) =
-                                            serde_json::from_str::<VersionedValue>(json_str)
-                                        {
-                                            table_data.insert(key, value);
+                        let file = File::open(&snapshot_path)?;
+                        let reader = BufReader::new(file);
+                        
+                        // Read line by line (each line is a Base64 encrypted chunk)
+                        for encrypted_line in reader.lines() {
+                            let encrypted_line = encrypted_line?;
+                            let trimmed_line = encrypted_line.trim();
+                            
+                            if trimmed_line.is_empty() {
+                                continue;
+                            }
+
+                            match decrypt(trimmed_line) {
+                                Ok(decrypted_content) => {
+                                    // Process the decrypted block of plaintext
+                                    for line in decrypted_content.lines() {
+                                        if let Some((key, json_str)) = line.split_once('=') {
+                                            let key = key.trim().to_string();
+                                            let json_str = json_str.trim();
+                                            if let Ok(value) =
+                                                serde_json::from_str::<VersionedValue>(json_str)
+                                            {
+                                                table_data.insert(key, value);
+                                            }
                                         }
                                     }
                                 }
-                            }
-                            Err(e) => {
-                                eprintln!(
-                                    "Decryption failed for table {}: {}",
-                                    table_name, e
-                                );
-                                process::exit(1);
+                                Err(e) => {
+                                    eprintln!(
+                                        "Decryption failed for table {} chunk: {}",
+                                        table_name, e
+                                    );
+                                    process::exit(1);
+                                }
                             }
                         }
                     }
 
-                    // Replay WAL entries.
+                    // --- 2. Replay WAL entries ---
                     let wal_path = table_folder.join("wal.log");
                     if wal_path.exists() {
                         let file = File::open(&wal_path)?;
                         let reader = BufReader::new(file);
+                        
                         for line in reader.lines() {
                             let line = line?;
-                            if let Ok(entry) =
-                                serde_json::from_str::<WalEntry>(&line)
-                            {
+                            if let Ok(entry) = serde_json::from_str::<WalEntry>(&line) {
                                 match entry.op.as_str() {
                                     "put" => {
                                         if let Some(val) = entry.value {
@@ -237,12 +275,13 @@ pub async fn load_all_tables(state: &web::Data<AppState>) -> io::Result<()> {
         }
         Ok::<HashMap<String, HashMap<String, VersionedValue>>, io::Error>(store)
     })
-        .await??;
+    .await??;
+    
+    // Apply the newly built store to the live application state
     let mut store_write = state.store.write().await;
     *store_write = new_store;
     Ok(())
 }
-
 /// Spawns a background Tokio task that periodically saves the current state to disk.
 pub async fn cold_save(state: web::Data<AppState>, interval: usize) {
     tokio::spawn(async move {