omni_dev/gmail/
attachments.rs1use std::collections::HashSet;
15
16use mail_parser::{ContentType, MessageParser, MimeHeaders};
17
18use crate::utils::path::attachment_filename;
19
20pub(crate) struct ExtractedAttachment {
25 pub(crate) filename: String,
26 pub(crate) contents: Vec<u8>,
27}
28
29pub(crate) fn extract_attachments(raw: &[u8]) -> Vec<ExtractedAttachment> {
42 let Some(message) = MessageParser::default().parse(raw) else {
43 return Vec::new();
44 };
45
46 let mut seen = HashSet::new();
47 message
48 .attachments()
49 .filter(|part| {
50 part.content_disposition()
51 .is_some_and(ContentType::is_attachment)
52 })
53 .enumerate()
54 .map(|(index, part)| {
55 let name =
56 attachment_filename(part.attachment_name().unwrap_or(""), &index.to_string());
57 ExtractedAttachment {
58 filename: dedupe_filename(&mut seen, name),
59 contents: part.contents().to_vec(),
60 }
61 })
62 .collect()
63}
64
65fn dedupe_filename(seen: &mut HashSet<String>, name: String) -> String {
76 if seen.insert(name.to_ascii_lowercase()) {
77 return name;
78 }
79 let (stem, ext) = match name.rsplit_once('.') {
80 Some((stem, ext)) if !stem.is_empty() => (stem.to_string(), Some(ext.to_string())),
81 _ => (name.clone(), None),
82 };
83 let mut n = 1u32;
84 loop {
85 let candidate = match &ext {
86 Some(ext) => format!("{stem}-{n}.{ext}"),
87 None => format!("{stem}-{n}"),
88 };
89 if seen.insert(candidate.to_ascii_lowercase()) {
90 return candidate;
91 }
92 n += 1;
93 }
94}
95
96#[cfg(test)]
97#[allow(clippy::unwrap_used, clippy::expect_used)]
98mod tests {
99 use super::*;
100 use base64::Engine as _;
101
102 #[test]
103 fn extract_attachments_decodes_base64_binary_content() {
104 let bytes: &[u8] = &[0x50, 0x44, 0x46, 0x00, 0x01, 0x02, 0xff];
105 let encoded = base64::engine::general_purpose::STANDARD.encode(bytes);
106 let raw = format!(
107 "Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
108--B\r\nContent-Type: text/plain\r\n\r\nHello\r\n\
109--B\r\nContent-Type: application/pdf\r\nContent-Transfer-Encoding: base64\r\n\
110Content-Disposition: attachment; filename=\"report.pdf\"\r\n\r\n{encoded}\r\n\
111--B--\r\n"
112 );
113 let extracted = extract_attachments(raw.as_bytes());
114 assert_eq!(extracted.len(), 1);
115 assert_eq!(extracted[0].filename, "report.pdf");
116 assert_eq!(extracted[0].contents, bytes);
117 }
118
119 #[test]
120 fn extract_attachments_decodes_quoted_printable_content() {
121 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
122--B\r\nContent-Type: text/plain\r\n\r\nHello\r\n\
123--B\r\nContent-Type: text/plain; charset=\"utf-8\"\r\nContent-Transfer-Encoding: quoted-printable\r\n\
124Content-Disposition: attachment; filename=\"notes.txt\"\r\n\r\nCaf=C3=A9\r\n\
125--B--\r\n";
126 let extracted = extract_attachments(raw);
127 assert_eq!(extracted.len(), 1);
128 assert_eq!(extracted[0].filename, "notes.txt");
129 assert_eq!(extracted[0].contents, "Café".as_bytes());
130 }
131
132 #[test]
133 fn extract_attachments_decodes_rfc2231_percent_encoded_filename() {
134 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
135--B\r\nContent-Type: application/pdf\r\nContent-Disposition: attachment; filename*=UTF-8''report%20Q3.pdf\r\n\r\ndata\r\n\
136--B--\r\n";
137 let extracted = extract_attachments(raw);
138 assert_eq!(extracted.len(), 1);
139 assert_eq!(extracted[0].filename, "report Q3.pdf");
140 }
141
142 #[test]
143 fn extract_attachments_resolves_rfc2231_continuation_filename() {
144 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
145--B\r\nContent-Type: application/pdf\r\nContent-Disposition: attachment; filename*0=\"report_\"; filename*1=\"Q3.pdf\"\r\n\r\ndata\r\n\
146--B--\r\n";
147 let extracted = extract_attachments(raw);
148 assert_eq!(extracted.len(), 1);
149 assert_eq!(extracted[0].filename, "report_Q3.pdf");
150 }
151
152 #[test]
153 fn extract_attachments_synthesizes_name_for_unnamed_attachment() {
154 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
155--B\r\nContent-Type: application/pdf\r\nContent-Disposition: attachment\r\n\r\ndata\r\n\
156--B--\r\n";
157 let extracted = extract_attachments(raw);
158 assert_eq!(extracted.len(), 1);
159 assert_eq!(extracted[0].filename, "attachment-0");
160 }
161
162 #[test]
163 fn extract_attachments_excludes_inline_disposition() {
164 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
165--B\r\nContent-Type: image/png\r\nContent-Disposition: inline; filename=\"logo.png\"\r\n\r\ndata\r\n\
166--B--\r\n";
167 let extracted = extract_attachments(raw);
168 assert!(extracted.is_empty());
169 }
170
171 #[test]
172 fn extract_attachments_dedupes_same_named_attachments_within_one_message() {
173 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
174--B\r\nContent-Type: image/png\r\nContent-Disposition: attachment; filename=\"image.png\"\r\n\r\nfirst\r\n\
175--B\r\nContent-Type: image/png\r\nContent-Disposition: attachment; filename=\"image.png\"\r\n\r\nsecond\r\n\
176--B--\r\n";
177 let extracted = extract_attachments(raw);
178 assert_eq!(extracted.len(), 2);
179 assert_eq!(extracted[0].filename, "image.png");
180 assert_eq!(extracted[1].filename, "image-1.png");
181 }
182
183 #[test]
184 fn extract_attachments_dedupes_names_differing_only_by_case() {
185 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
186--B\r\nContent-Type: application/pdf\r\nContent-Disposition: attachment; filename=\"Report.PDF\"\r\n\r\nfirst\r\n\
187--B\r\nContent-Type: application/pdf\r\nContent-Disposition: attachment; filename=\"report.pdf\"\r\n\r\nsecond\r\n\
188--B--\r\n";
189 let extracted = extract_attachments(raw);
190 assert_eq!(extracted.len(), 2);
191 assert_eq!(extracted[0].filename, "Report.PDF");
192 assert_eq!(extracted[1].filename, "report-1.pdf");
193 }
194
195 #[test]
196 fn extract_attachments_sanitizes_path_traversal_filename() {
197 let raw = b"Content-Type: multipart/mixed; boundary=\"B\"\r\n\r\n\
198--B\r\nContent-Type: application/octet-stream\r\nContent-Disposition: attachment; filename=\"../../etc/passwd\"\r\n\r\ndata\r\n\
199--B--\r\n";
200 let extracted = extract_attachments(raw);
201 assert_eq!(extracted.len(), 1);
202 assert_eq!(extracted[0].filename, "passwd");
203 }
204
205 #[test]
206 fn extract_attachments_returns_empty_vec_for_unparseable_message() {
207 assert!(extract_attachments(b"").is_empty());
208 }
209
210 #[test]
211 fn extract_attachments_returns_empty_vec_for_plain_text_message() {
212 let raw = b"Subject: Hi\r\nFrom: a@example.com\r\n\r\nJust text, no attachments.";
213 assert!(extract_attachments(raw).is_empty());
214 }
215}