refactor get_body_parent_part_mut to get rid of unsafe code

src/lib.rs

@@ -1,4 +1,4 @@
-//#![forbid(unsafe_code)]
+#![forbid(unsafe_code)]
 
 pub mod header;
 pub mod headers;
@@ -18,7 +18,6 @@ use mediatype::{names::*, values::UTF_8, MediaTypeBuf, Name, ReadParams, WritePa
 use rand::{distr::Alphanumeric, rng, Rng as _};
 use std::borrow::Cow;
 use std::boxed::Box;
-use std::collections::VecDeque;
 use std::ops::Deref;
 
 #[derive(Debug)]
@@ -194,6 +193,16 @@ pub fn generate_boundary() -> String {
         .collect::<String>()
 }
 
+/// Result of checking whether a MIME part is a text body (plain/html).
+enum BodyStatus {
+    /// This part is a text/plain or text/html body (not an attachment)
+    IsBody,
+    /// This is a leaf node but not a body (attachment, other content type, etc.)
+    NotBody,
+    /// This is a multipart container - subparts need to be checked
+    NeedsTraversal,
+}
+
 #[derive(Debug, Clone)]
 pub struct CowMail<'a, H: Header<'a>> {
     pub bytes_orig: &'a [u8],
@@ -290,11 +299,27 @@ impl<'a, H: Header<'a> + 'a> CowMail<'a, H> {
         Ok(())
     }
 
-    fn contains_body(&self, mime: MediaType, bodytype: BodyType) -> Option<bool> {
+    /// Check if this MIME part is a text body (plain/html), not a body, or needs traversal.
+    fn body_status(&self, bodytype: Option<BodyType>) -> BodyStatus {
         match self.is_multipart() {
-            false => Some(!self.is_attachment() && mime.ty == TEXT && mime.subty == Name::from(bodytype)),
+            true if self.subparts.is_empty() => BodyStatus::NotBody,
-            _ if self.subparts.is_empty() => Some(false),
+            true => BodyStatus::NeedsTraversal,
-            _ => None,
+            false => {
+                let mime = self.get_mime_type();
+                if mime.ty() != TEXT || self.is_attachment() {
+                    return BodyStatus::NotBody;
+                }
+
+                let body_types = match bodytype {
+                    Some(bodytype) => vec![Name::from(bodytype)],
+                    None => vec![PLAIN, HTML],
+                };
+
+                match body_types.contains(&mime.subty()) {
+                    true => BodyStatus::IsBody,
+                    false => BodyStatus::NotBody,
+                }
+            }
         }
     }
 
@@ -310,15 +335,11 @@ impl<'a, H: Header<'a> + 'a> CowMail<'a, H> {
     }
 
     pub fn get_body_part(&self, bodytype: BodyType) -> Option<&CowMail<'a, H>> {
-        let mime = self.get_mime_type();
+        match self.body_status(Some(bodytype)) {
-
+            BodyStatus::IsBody => Some(self),
-        match self.contains_body(mime.essence(), bodytype) {
+            BodyStatus::NotBody => None,
-            Some(res) => match res {
+            BodyStatus::NeedsTraversal => {
-                true => Some(self),
+                if self.get_mime_type().subty() == RELATED {
-                false => None,
-            },
-            None => {
-                if mime.subty() == RELATED {
                     if let Some(candidate) = self.get_related_candidate() {
                         return self.subparts[candidate].get_body_part(bodytype);
                     }
@@ -329,14 +350,11 @@ impl<'a, H: Header<'a> + 'a> CowMail<'a, H> {
     }
 
     pub fn get_body_part_mut(&mut self, bodytype: BodyType) -> Option<&mut CowMail<'a, H>> {
-        let mime = self.get_mime_type();
+        match self.body_status(Some(bodytype)) {
-        match self.contains_body(mime.essence(), bodytype) {
+            BodyStatus::IsBody => Some(self),
-            Some(res) => match res {
+            BodyStatus::NotBody => None,
-                true => Some(self),
+            BodyStatus::NeedsTraversal => {
-                false => None,
+                if self.get_mime_type().subty() == RELATED {
-            },
-            None => {
-                if mime.subty() == RELATED {
                     if let Some(candidate) = self.get_related_candidate() {
                         return self.subparts[candidate].get_body_part_mut(bodytype);
                     }
@@ -439,51 +457,27 @@ impl<'a, H: Header<'a> + 'a> CowMail<'a, H> {
         self.add_multipart(subpart, MediaType::new(MULTIPART, MIXED))
     }
 
-    fn contains_plain_or_html_body(&self, mime: &MediaTypeBuf) -> Option<bool> {
+    fn find_body_parent_path(&self, path: &mut Vec<usize>, bodytype: Option<BodyType>) -> Option<Vec<usize>> {
-        match self.is_multipart() {
+        // For multipart/related, only check the start part (or first part if no start parameter)
-            false => {
+        let subpart_indices: Box<dyn Iterator<Item = usize>> = if self.get_mime_type().subty() == RELATED {
-                let subty = mime.subty();
+            Box::new(self.get_related_candidate().into_iter())
-                Some(!self.is_attachment() && mime.ty() == TEXT && (subty == PLAIN || subty == HTML))
+        } else {
-            }
+            Box::new(0..self.subparts.len())
-            true if self.subparts.is_empty() => Some(false),
+        };
-            true => None,
-        }
-    }
 
-    fn get_body_parent_part_mut(&mut self) -> Option<&mut CowMail<'a, H>> {
+        for i in subpart_indices {
-        let mime = self.get_mime_type();
+            let subpart = &self.subparts[i];
-
+            match subpart.body_status(bodytype) {
-        if let Some(res) = self.contains_plain_or_html_body(&mime) {
+                BodyStatus::IsBody => return Some(path.clone()),
-            return match res {
+                BodyStatus::NotBody => continue,
-                false => None,
+                BodyStatus::NeedsTraversal => {
-                true => Some(self),
+                    path.push(i);
-            };
+                    match subpart.find_body_parent_path(path, bodytype) {
-        }
-
-        let mut stack = VecDeque::new();
-        stack.push_back(self as *mut CowMail<H>);
-
-        while let Some(current_ptr) = stack.pop_front() {
-            let current = unsafe { &mut *current_ptr };
-
-            if current.get_mime_type().subty() == RELATED {
-                if let Some(candidate) = current.get_related_candidate() {
-                    match current.subparts[candidate].contains_plain_or_html_body(&mime) {
-                        Some(true) => return Some(current),
-                        Some(false) => return None,
                         None => {
-                            stack.clear();
+                            path.pop();
-                            stack.push_back(&mut current.subparts[candidate])
                         }
+                        found => return found,
                     }
-                    continue;
-                }
-            }
-            for p in current.subparts.iter_mut() {
-                match p.contains_plain_or_html_body(&mime) {
-                    Some(true) => return Some(current),
-                    Some(false) => continue,
-                    None => stack.push_back(p),
                 }
             }
         }
@@ -491,6 +485,23 @@ impl<'a, H: Header<'a> + 'a> CowMail<'a, H> {
         None
     }
 
+    fn get_body_parent_part_mut(&mut self) -> Option<&mut CowMail<'a, H>> {
+        match self.body_status(None) {
+            BodyStatus::IsBody => Some(self),
+            BodyStatus::NotBody => None,
+            BodyStatus::NeedsTraversal => match self.find_body_parent_path(&mut Vec::new(), None) {
+                Some(path) => {
+                    let mut current = self;
+                    for idx in path {
+                        current = &mut current.subparts[idx];
+                    }
+                    Some(current)
+                }
+                None => None,
+            },
+        }
+    }
+
     pub fn set_text_body<T: ToString>(&mut self, content: T) -> Result<(), CowMailError> {
         if let Some(parent) = self.get_body_parent_part_mut() {
             match parent.get_body_part_mut(BodyType::PLAIN) {
@@ -861,11 +872,9 @@ fn parse_mail_recursive<'a, H: Header<'a> + 'a>(
     let mut body_end_idx = raw_data.len();
     let mut end_bytes: &[u8] = &[];
 
-    if body_idx + 2 < raw_data.len() {
+    if body_idx + 2 < raw_data.len() && &raw_data[(raw_data.len() - 2)..] == b"\r\n" {
-        if &raw_data[(raw_data.len() - 2)..] == b"\r\n" {
+        body_end_idx = raw_data.len() - 2;
-            body_end_idx = raw_data.len() - 2;
+        end_bytes = b"\r\n";
-            end_bytes = b"\r\n";
-        }
     }
 
     let default_ctype = ContentType::default_conditional(in_multipart_digest);