Commit
+46 -21 +/-2 browse
| 1 | diff --git a/maitred/src/rewrite.rs b/maitred/src/rewrite.rs |
| 2 | index a455288..8f02ee0 100644 |
| 3 | --- a/maitred/src/rewrite.rs |
| 4 | +++ b/maitred/src/rewrite.rs |
| 5 | @@ -1,19 +1,39 @@ |
| 6 | - use mail_parser::{Message, MessageParser}; |
| 7 | + use mail_parser::{HeaderName, Message, MessageParser}; |
| 8 | |
| 9 | - /// Wrapper around raw bytes that allows for easy modification of an email. |
| 10 | - pub struct Rewrite<'a>(pub &'a mut Vec<u8>); |
| 11 | + /// Basically a hack that can modify messages expensively re-parsing them on |
| 12 | + /// each modificaiton. The mail_parser project has mentioned adding this |
| 13 | + /// functionality and perhaps this could be upstreamed. |
| 14 | + pub struct Rewrite<'a> { |
| 15 | + parser: MessageParser, |
| 16 | + raw_message: &'a mut Vec<u8>, |
| 17 | + } |
| 18 | |
| 19 | impl<'a> Rewrite<'a> { |
| 20 | - /// Prepend a header to the message body |
| 21 | - pub fn with_header(&mut self, key: &str, value: &str) { |
| 22 | - let header: Vec<u8> = format!("{}: {}\n", key, value.trim_end()).bytes().collect(); |
| 23 | - self.0.splice(0..0, header); |
| 24 | + pub fn new(parser: Option<MessageParser>, raw_message: &'a mut Vec<u8>) -> Rewrite<'a> { |
| 25 | + Self { |
| 26 | + parser: parser.unwrap_or_default(), |
| 27 | + raw_message, |
| 28 | + } |
| 29 | } |
| 30 | |
| 31 | - /// Return a parsed Message |
| 32 | - pub fn message(&'a self) -> Option<Message<'a>> { |
| 33 | - let message = MessageParser::default().parse(self.0.as_slice()); |
| 34 | - message |
| 35 | + pub fn message(&'a self) -> Message<'a> { |
| 36 | + self.parser |
| 37 | + .parse(self.raw_message.as_slice()) |
| 38 | + .expect("Cannot parse message") |
| 39 | + } |
| 40 | + |
| 41 | + /// Set the header to a string value replacing it if it already exists |
| 42 | + pub fn set_header(&mut self, key: HeaderName, value: &str) { |
| 43 | + let message = self |
| 44 | + .parser |
| 45 | + .parse_headers(self.raw_message.as_slice()) |
| 46 | + .expect("Cannot parse message"); |
| 47 | + if let Some(header) = message.headers().iter().find(|header| header.name() == key.as_str()) { |
| 48 | + let (start, end) = (header.offset_field(), header.offset_end()); |
| 49 | + self.raw_message.drain(start..end); |
| 50 | + } |
| 51 | + let header: Vec<u8> = format!("{}: {}\n", key, value.trim_end()).bytes().collect(); |
| 52 | + self.raw_message.splice(0..0, header); |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | @@ -34,15 +54,20 @@ Hello World |
| 57 | "#; |
| 58 | |
| 59 | #[test] |
| 60 | - fn test_apppend_header() { |
| 61 | - let mut email_bytes = TEST_EMAIL.as_bytes().to_vec(); |
| 62 | - let mut rewrite = Rewrite(&mut email_bytes); |
| 63 | - rewrite.with_header("a", "b"); |
| 64 | - let message = rewrite.message().unwrap(); |
| 65 | + fn rewrite() { |
| 66 | + let email_bytes = &mut TEST_EMAIL.as_bytes().to_vec(); |
| 67 | + let mut rewrite = Rewrite::new(None, email_bytes); |
| 68 | + rewrite.set_header(HeaderName::Other("a".into()), "b"); |
| 69 | + rewrite.set_header(HeaderName::Subject, "Bar"); |
| 70 | + let message = rewrite.message(); |
| 71 | + println!("{}", String::from_utf8_lossy(message.raw_message())); |
| 72 | let value = message.header("a").unwrap(); |
| 73 | assert!(value.as_text().unwrap() == "b"); |
| 74 | + let value = message.header("Subject").unwrap(); |
| 75 | + assert!(value.as_text().unwrap() == "Bar"); |
| 76 | let message_str = String::from_utf8(message.raw_message().to_vec()).unwrap(); |
| 77 | - assert!(message_str.split("\n").next().unwrap() == "a: b"); |
| 78 | - assert!(message_str.split("\n").nth(1).unwrap() == "Date: Mon, 2 Sep 2024 00:17:18 +0200"); |
| 79 | + assert!(message_str.split("\n").next().unwrap() == "Subject: Bar"); |
| 80 | + assert!(message_str.split("\n").nth(1).unwrap() == "a: b"); |
| 81 | + assert!(message_str.split("\n").nth(2).unwrap() == "Date: Mon, 2 Sep 2024 00:17:18 +0200"); |
| 82 | } |
| 83 | } |
| 84 | diff --git a/maitred/src/worker.rs b/maitred/src/worker.rs |
| 85 | index b015d2c..6f224f8 100644 |
| 86 | --- a/maitred/src/worker.rs |
| 87 | +++ b/maitred/src/worker.rs |
| 88 | @@ -85,12 +85,12 @@ impl Worker { |
| 89 | } |
| 90 | |
| 91 | let mut modified_message_bytes = message.raw_message().to_vec(); |
| 92 | - let mut modified = Rewrite(&mut modified_message_bytes); |
| 93 | + let mut modified = Rewrite::new(None, &mut modified_message_bytes); |
| 94 | if let Some(dkim_passed) = dkim_passed { |
| 95 | - modified.with_header(HEADER_DKIM_RESULT, &dkim_passed.to_string()) |
| 96 | + modified.set_header(HEADER_DKIM_RESULT.into(), &dkim_passed.to_string()) |
| 97 | } |
| 98 | |
| 99 | - let to_deliver = modified.message().unwrap(); |
| 100 | + let to_deliver = modified.message(); |
| 101 | |
| 102 | if let Some(delivery) = &self.delivery { |
| 103 | tracing::info!("Delivering message {}", message_id); |