From 66fd50bc23043e6c07a30cb2970a37a47b3cdc81 Mon Sep 17 00:00:00 2001 From: rustmailer Date: Mon, 18 May 2026 20:25:13 +0800 Subject: [PATCH] fix: bichon-cli OOMs on import #233 --- crates/cli/src/mbox/mod.rs | 223 +++++++++++++++++++++++++++++++++- crates/core/src/import/mod.rs | 27 +++- 2 files changed, 246 insertions(+), 4 deletions(-) diff --git a/crates/cli/src/mbox/mod.rs b/crates/cli/src/mbox/mod.rs index 2a4fb59..f3065a2 100644 --- a/crates/cli/src/mbox/mod.rs +++ b/crates/cli/src/mbox/mod.rs @@ -32,6 +32,11 @@ use mail_parser::parsers::MessageStream; use mail_parser::MessageParser; use reqwest::Client; +/// Skip emails larger than this with a warning (100 MB). +const MAX_EMAIL_BYTES: usize = 100 * 1024 * 1024; +/// Flush a folder buffer when accumulated base64 bytes exceed this (200 MB). +const MAX_BUFFER_BYTES: usize = 200 * 1024 * 1024; + pub mod gmail; pub mod reader; @@ -133,13 +138,29 @@ pub async fn run_import( }; let mut folder_buffers: HashMap> = HashMap::new(); + let mut total_buffered_bytes: usize = 0; let batch_limit = 50; + let mut skipped_count: u64 = 0; println!("Starting import process..."); for (index, e) in mbox.iter().enumerate() { let msg_num = index + 1; let body = e.data; + + if body.len() > MAX_EMAIL_BYTES { + let size_mb = body.len() as f64 / 1024.0 / 1024.0; + eprintln!( + "{} {}: email #{} is {:.1} MB (limit 100 MB). Skipping...", + style("Warning").yellow().bold(), + style(format!("oversized")).dim(), + msg_num, + size_mb, + ); + skipped_count += 1; + continue; + } + let message = match MessageParser::new().parse(body) { Some(msg) => msg, None => { @@ -149,6 +170,7 @@ pub async fn run_import( style(format!("at message #{}", msg_num)).dim(), "Failed to parse email structure. Skipping..." ); + skipped_count += 1; continue; } }; @@ -178,14 +200,22 @@ pub async fn run_import( get_default_folder() }; + // Drop message before base64-encoding to free MIME parse memory. + drop(message); + let b64_eml = base64_encode_url_safe!(&body); + let encoded_len = b64_eml.len(); + let buffer = folder_buffers .entry(folder_name.clone()) - .or_insert_with(|| Vec::new()); + .or_insert_with(Vec::new); buffer.push(b64_eml); + total_buffered_bytes += encoded_len; - if buffer.len() >= batch_limit { + if buffer.len() >= batch_limit || total_buffered_bytes >= MAX_BUFFER_BYTES { let emls_to_send = folder_buffers.remove(&folder_name).unwrap(); + let freed: usize = emls_to_send.iter().map(|s| s.len()).sum(); + total_buffered_bytes = total_buffered_bytes.saturating_sub(freed); send_batch_request(&client, config, account_id, &folder_name, emls_to_send).await; } } @@ -196,5 +226,194 @@ pub async fn run_import( } } + if skipped_count > 0 { + println!( + "{}", + style(format!( + "Skipped {} email(s) (oversized or unparseable).", + skipped_count + )) + .yellow() + .bold() + ); + } + println!("{}", style("Import completed successfully!").green().bold()); } + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + /// Fake sender: records every flushed batch as (folder_name, email_count, total_bytes). + struct FakeSender { + batches: Vec<(String, usize, usize)>, + } + + impl FakeSender { + fn new() -> Self { + Self { batches: vec![] } + } + fn send(&mut self, folder: &str, emls: Vec) { + let count = emls.len(); + let bytes: usize = emls.iter().map(|s| s.len()).sum(); + self.batches.push((folder.to_string(), count, bytes)); + // emls is dropped here, simulating real send + } + } + + fn fake_encode(size: usize) -> String { + // base64 expands ~1.33x, so the encoded string is roughly this long. + // We just need a predictable byte size, so use a repeated character. + "x".repeat(size) + } + + #[test] + fn flush_on_global_byte_threshold() { + let mut buffers: HashMap> = HashMap::new(); + let mut total_bytes: usize = 0; + let batch_limit = 50; + let mut sender = FakeSender::new(); + + // Simulate 3 emails, each 80 MB encoded, spread across 3 folders. + // After each email, global total goes up by 80 MB. + // After the 3rd email: 240 MB > 200 MB → flush the folder that got the 3rd email. + let emails = vec![ + ("Inbox", 80_000_000), + ("Sent", 80_000_000), + ("Archive", 80_000_000), + ]; + + for (folder, eml_size) in emails { + let encoded = fake_encode(eml_size); + let len = encoded.len(); + let buffer = buffers.entry(folder.to_string()).or_insert_with(Vec::new); + buffer.push(encoded); + total_bytes += len; + + if buffer.len() >= batch_limit || total_bytes >= MAX_BUFFER_BYTES { + let sent = buffers.remove(folder).unwrap(); + let freed: usize = sent.iter().map(|s| s.len()).sum(); + total_bytes = total_bytes.saturating_sub(freed); + sender.send(folder, sent); + } + } + + // The 3rd email should trigger a global flush of "Archive". + assert_eq!(sender.batches.len(), 1); + assert_eq!(sender.batches[0].0, "Archive"); + assert_eq!(sender.batches[0].1, 1); + // "Inbox" and "Sent" are still buffered (160 MB total). + assert_eq!(buffers.len(), 2); + assert!(buffers.contains_key("Inbox")); + assert!(buffers.contains_key("Sent")); + assert_eq!(total_bytes, 160_000_000); + } + + #[test] + fn flush_on_count_threshold() { + let mut buffers: HashMap> = HashMap::new(); + let mut total_bytes: usize = 0; + let batch_limit = 3; + let mut sender = FakeSender::new(); + + // 4 small emails all to Inbox, well under byte threshold. + for _ in 0..4 { + let encoded = fake_encode(100); // tiny + let len = encoded.len(); + let buffer = buffers + .entry("Inbox".to_string()) + .or_insert_with(Vec::new); + buffer.push(encoded); + total_bytes += len; + + if buffer.len() >= batch_limit || total_bytes >= MAX_BUFFER_BYTES { + let sent = buffers.remove("Inbox").unwrap(); + let freed: usize = sent.iter().map(|s| s.len()).sum(); + total_bytes = total_bytes.saturating_sub(freed); + sender.send("Inbox", sent); + } + } + + // Count=3 should trigger flush once; the 4th email stays buffered. + assert_eq!(sender.batches.len(), 1); + assert_eq!(sender.batches[0].1, 3); // 3 emails flushed + let remaining = buffers.get("Inbox").unwrap(); + assert_eq!(remaining.len(), 1); // 1 still buffered + } + + #[test] + fn global_bytes_exact_boundary() { + let mut buffers: HashMap> = HashMap::new(); + let mut total_bytes: usize = 0; + let mut sender = FakeSender::new(); + + // Push one email that puts us right at 200 MB. + let encoded = fake_encode(MAX_BUFFER_BYTES); + let len = encoded.len(); + buffers + .entry("Inbox".to_string()) + .or_insert_with(Vec::new) + .push(encoded); + total_bytes += len; + + if total_bytes >= MAX_BUFFER_BYTES { + let sent = buffers.remove("Inbox").unwrap(); + let freed: usize = sent.iter().map(|s| s.len()).sum(); + total_bytes = total_bytes.saturating_sub(freed); + sender.send("Inbox", sent); + } + + // Should have flushed on the boundary. + assert_eq!(sender.batches.len(), 1); + assert_eq!(total_bytes, 0); + } + + #[test] + fn flush_one_folder_does_not_lose_others() { + let mut buffers: HashMap> = HashMap::new(); + let mut total_bytes: usize = 0; + let batch_limit = 50; + let mut sender = FakeSender::new(); + + // Build up A to 150 MB, B to 100 MB (total 250 MB > 200 MB). + // A should trigger flush; B should stay buffered. + let folder_a = "A".to_string(); + let folder_b = "B".to_string(); + + // Folder A: 150 MB + let encoded = fake_encode(150_000_000); + let len = encoded.len(); + buffers.entry(folder_a.clone()).or_insert_with(Vec::new).push(encoded); + total_bytes += len; + + // Folder B: 100 MB → total 250 MB → trigger flush on B + let encoded = fake_encode(100_000_000); + let len = encoded.len(); + buffers.entry(folder_b.clone()).or_insert_with(Vec::new).push(encoded); + total_bytes += len; + + // Check trigger on B + let b_buffer = buffers.get(&folder_b).unwrap(); + if b_buffer.len() >= batch_limit || total_bytes >= MAX_BUFFER_BYTES { + let sent = buffers.remove(&folder_b).unwrap(); + let freed: usize = sent.iter().map(|s| s.len()).sum(); + total_bytes = total_bytes.saturating_sub(freed); + sender.send(&folder_b, sent); + } + + assert_eq!(sender.batches.len(), 1); + assert_eq!(sender.batches[0].0, "B"); // B flushed + assert!(buffers.contains_key("A")); // A still there + assert_eq!(total_bytes, 150_000_000); + } + + #[test] + fn skip_oversized_email() { + assert!(100 <= MAX_EMAIL_BYTES); + // Use vec! so the 100 MB array lives on the heap, not the stack. + let huge = vec![0u8; MAX_EMAIL_BYTES + 1]; + assert!(huge.len() > MAX_EMAIL_BYTES); + } +} diff --git a/crates/core/src/import/mod.rs b/crates/core/src/import/mod.rs index 0a6e897..8cd533e 100644 --- a/crates/core/src/import/mod.rs +++ b/crates/core/src/import/mod.rs @@ -32,6 +32,9 @@ use crate::{ raise_error, }; +/// Skip individual emails larger than this after decoding (100 MB). +const MAX_SINGLE_EML_BYTES: usize = 100 * 1024 * 1024; + #[derive(Debug, Clone, Default, Eq, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "web-api", derive(poem_openapi::Object))] pub struct BatchEmlRequest { @@ -66,7 +69,7 @@ pub struct BatchEmlResult { pub struct ImportEmls; impl ImportEmls { - pub async fn do_import(request: BatchEmlRequest) -> BichonResult { + pub async fn do_import(mut request: BatchEmlRequest) -> BichonResult { let account = AccountModel::check_account_exists(request.account_id)?; if !account.enabled { @@ -115,7 +118,8 @@ impl ImportEmls { let mut failed_details: Vec = Vec::new(); // Store failure details let total = request.emls.len(); - for (index, eml_base64) in request.emls.into_iter().enumerate() { + let mut index: usize = 0; + while let Some(eml_base64) = request.emls.pop() { let decoded = match base64_decode_url_safe!(eml_base64.as_bytes()) { Ok(bytes) => bytes, Err(e) => { @@ -126,9 +130,26 @@ impl ImportEmls { index, error_message: error_msg, }); + index += 1; continue; } }; + // eml_base64 string dropped here — frees base64 memory before parsing + + if decoded.len() > MAX_SINGLE_EML_BYTES { + let size_mb = decoded.len() as f64 / 1024.0 / 1024.0; + let error_msg = format!( + "Email at index {} is {:.1} MB (limit 50 MB). Skipping.", + index, size_mb, + ); + tracing::warn!("{}", error_msg); + failed_details.push(FailedEmlDetail { + index, + error_message: error_msg, + }); + index += 1; + continue; + } match extract_envelope_from_eml(&decoded, account_id, mailbox_id).await { Ok(_) => { @@ -144,9 +165,11 @@ impl ImportEmls { index, error_message: error_msg, }); + index += 1; continue; } }; + index += 1; } let failed_count = failed_details.len();