Commit

Author:

Hash:

Timestamp:

+46 -21 +/-2 browse

Kevin Schoon [me@kevinschoon.com]

faa59826ea0333091cb28006c21c564a4c44e26d

Sat, 11 Jan 2025 22:36:46 +0000 (9 months ago)

improve rewrite module
1diff --git a/maitred/src/rewrite.rs b/maitred/src/rewrite.rs
2index 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
85index 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);