Commit: 0e0fdf5ade73652928febb6b1b259502af4081b9

Commit

Author:

Hash:

Timestamp:

+253 -41 +/-9 browse

Kevin Schoon [me@kevinschoon.com]

0e0fdf5ade73652928febb6b1b259502af4081b9

Sat, 17 Aug 2024 14:34:47 +0000 (1.2 years ago)

finish milter with asynchronous background worker

1	diff --git a/Cargo.lock b/Cargo.lock
2	index f34642f..5bb94a1 100644
3	--- a/Cargo.lock
4	+++ b/Cargo.lock
5	@@ -92,6 +92,31 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
6	checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
7
8	[[package]]
9	+ name = "crossbeam-deque"
10	+ version = "0.8.5"
11	+ source = "registry+https://github.com/rust-lang/crates.io-index"
12	+ checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d"
13	+ dependencies = [
14	+ "crossbeam-epoch",
15	+ "crossbeam-utils",
16	+ ]
17	+
18	+ [[package]]
19	+ name = "crossbeam-epoch"
20	+ version = "0.9.18"
21	+ source = "registry+https://github.com/rust-lang/crates.io-index"
22	+ checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
23	+ dependencies = [
24	+ "crossbeam-utils",
25	+ ]
26	+
27	+ [[package]]
28	+ name = "crossbeam-utils"
29	+ version = "0.8.20"
30	+ source = "registry+https://github.com/rust-lang/crates.io-index"
31	+ checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80"
32	+
33	+ [[package]]
34	name = "email_address"
35	version = "0.2.9"
36	source = "registry+https://github.com/rust-lang/crates.io-index"
37	@@ -283,6 +308,7 @@ version = "0.1.0"
38	dependencies = [
39	"async-trait",
40	"bytes",
41	+ "crossbeam-deque",
42	"email_address",
43	"futures",
44	"mail-parser",
45	diff --git a/cmd/maitred-debug/src/main.rs b/cmd/maitred-debug/src/main.rs
46	index d62121f..e4c3141 100644
47	--- a/cmd/maitred-debug/src/main.rs
48	+++ b/cmd/maitred-debug/src/main.rs
49	@@ -9,15 +9,12 @@ async fn main() -> Result<(), Error> {
50	.with_line_number(true)
51	.with_max_level(Level::DEBUG)
52	.init();
53	-
54	// Set the subscriber as the default subscriber
55	- let mail_server = Server::new("localhost")
56	+ let mut mail_server = Server::default()
57	.address("127.0.0.1:2525")
58	- .with_milter(MilterFunc(\|message: &Message\| {
59	- println!("{:?}", message);
60	- async move {
61	- Ok(Message::default().into_owned())
62	- }
63	+ .with_milter(MilterFunc::new(\|message: &Message\| {
64	+ let cloned = message.clone();
65	+ Box::pin(async move { Ok(cloned.to_owned()) })
66	}))
67	.with_session_opts(SessionOptions::default());
68	mail_server.listen().await?;
69	diff --git a/maitred/Cargo.toml b/maitred/Cargo.toml
70	index 587470b..4ca527d 100644
71	--- a/maitred/Cargo.toml
72	+++ b/maitred/Cargo.toml
73	@@ -6,6 +6,7 @@ edition = "2021"
74	[dependencies]
75	async-trait = "0.1.81"
76	bytes = "1.6.1"
77	+ crossbeam-deque = "0.8.5"
78	email_address = "0.2.9"
79	futures = "0.3.30"
80	mail-parser = { version = "0.9.3", features = ["serde", "serde_support"] }
81	diff --git a/maitred/src/lib.rs b/maitred/src/lib.rs
82	index ca47995..42a6a02 100644
83	--- a/maitred/src/lib.rs
84	+++ b/maitred/src/lib.rs
85	@@ -29,6 +29,7 @@ mod server;
86	mod session;
87	mod transport;
88	mod verify;
89	+ mod worker;
90
91	use smtp_proto::{Request, Response as SmtpResponse};
92	use transport::Response;
93	diff --git a/maitred/src/milter.rs b/maitred/src/milter.rs
94	index 978f25e..ef0ec99 100644
95	--- a/maitred/src/milter.rs
96	+++ b/maitred/src/milter.rs
97	@@ -1,5 +1,7 @@
98	use std::future::Future;
99	+ use std::pin::Pin;
100	use std::result::Result as StdResult;
101	+ use std::sync::Arc;
102
103	use async_trait::async_trait;
104	use mail_parser::Message;
105	@@ -22,34 +24,51 @@ pub enum Error {
106	#[async_trait]
107	pub trait Milter {
108	/// Apply the milter function to the incoming message
109	- async fn apply(&self, message: &Message) -> Result;
110	+ async fn apply(&self, message: &Message<'static>) -> Result;
111	}
112
113	- /// Helper wrapper implementing the Expansion trait
114	+ /// Wrapper implementing the Milter trait from a closure
115	/// # Example
116	/// ```rust
117	/// use mail_parser::Message;
118	/// use maitred::MilterFunc;
119	///
120	/// let my_expn_fn = MilterFunc(\|message: &Message\| {
121	- /// async move {
122	+ /// Box::pin(async move {
123	/// // rewrite message here
124	/// Ok(Message::default().to_owned())
125	- /// }
126	+ /// })
127	/// });
128	/// ```
129	- pub struct Func<F, T>(pub F)
130	+ #[derive(Clone)]
131	+ pub struct Func<F>(Arc<F>)
132	where
133	- F: Fn(&Message) -> T + Sync,
134	- T: Future<Output = Result> + Send;
135	+ F: Fn(&Message<'static>) -> Pin<Box<dyn Future<Output = Result> + Send>>
136	+ + Send
137	+ + Sync
138	+ + 'static;
139	+
140	+ impl<F> Func<F>
141	+ where
142	+ F: Fn(&Message<'static>) -> Pin<Box<dyn Future<Output = Result> + Send>>
143	+ + Send
144	+ + Sync
145	+ + 'static,
146	+ {
147	+ pub fn new(func: F) -> Self {
148	+ Func(Arc::new(func))
149	+ }
150	+ }
151
152	#[async_trait]
153	- impl<F, T> Milter for Func<F, T>
154	+ impl<F> Milter for Func<F>
155	where
156	- F: Fn(&Message) -> T + Sync,
157	- T: Future<Output = Result> + Send,
158	+ F: Fn(&Message<'static>) -> Pin<Box<dyn Future<Output = Result> + Send>>
159	+ + Send
160	+ + Sync
161	+ + 'static,
162	{
163	- async fn apply(&self, message: &Message) -> Result {
164	+ async fn apply(&self, message: &Message<'static>) -> Result {
165	let f = (self.0)(message);
166	f.await
167	}
168	diff --git a/maitred/src/server.rs b/maitred/src/server.rs
169	index 00d8f5b..487c925 100644
170	--- a/maitred/src/server.rs
171	+++ b/maitred/src/server.rs
172	@@ -3,18 +3,24 @@ use std::sync::Arc;
173	use std::time::Duration;
174
175	use bytes::Bytes;
176	+ use crossbeam_deque::Injector;
177	+ use crossbeam_deque::Stealer;
178	+ use crossbeam_deque::Worker as WorkQueue;
179	use futures::SinkExt;
180	- use mail_parser::MessageParser;
181	+ use futures::StreamExt;
182	use smtp_proto::Request;
183	+ use tokio::sync::mpsc::Sender;
184	use tokio::sync::Mutex;
185	+ use tokio::task::JoinHandle;
186	use tokio::{net::TcpListener, time::timeout};
187	- use tokio_stream::StreamExt;
188	+ use tokio_stream::{self as stream};
189	use tokio_util::codec::Framed;
190
191	use crate::error::Error;
192	use crate::pipeline::Pipeline;
193	use crate::session::Session;
194	use crate::transport::{Response, Transport};
195	+ use crate::worker::{Packet, Worker};
196	use crate::{Chunk, Milter};
197
198	/// The default port the server will listen on if none was specified in it's
199	@@ -55,30 +61,43 @@ impl ConditionalPipeline<'_> {
200	/// Server implements everything that is required to run an SMTP server by
201	/// binding to the configured address and processing individual TCP connections
202	/// as they are received.
203	- pub struct Server {
204	+ pub struct Server<M>
205	+ where
206	+ M: Milter + Clone + Send + Sync + 'static,
207	+ {
208	address: String,
209	hostname: String,
210	global_timeout: Duration,
211	options: Option<Rc<crate::session::Options>>,
212	- milters: Vec<Arc<dyn Milter>>,
213	+ milters: Vec<Arc<M>>,
214	+ n_threads: usize,
215	+ shutdown_handles: Vec<Sender<bool>>,
216	}
217
218	- impl Default for Server {
219	+ impl<M> Default for Server<M>
220	+ where
221	+ M: Milter + Clone + Send + Sync + 'static,
222	+ {
223	fn default() -> Self {
224	- Server {
225	+ Server::<M> {
226	address: DEFAULT_LISTEN_ADDR.to_string(),
227	hostname: String::default(),
228	global_timeout: Duration::from_secs(DEFAULT_GLOBAL_TIMEOUT_SECS),
229	options: None,
230	milters: vec![],
231	+ n_threads: std::thread::available_parallelism().unwrap().into(),
232	+ shutdown_handles: vec![],
233	}
234	}
235	}
236
237	- impl Server {
238	+ impl<M> Server<M>
239	+ where
240	+ M: Milter + Clone + Send + Sync + 'static,
241	+ {
242	/// Initialize a new SMTP server
243	pub fn new(hostname: &str) -> Self {
244	- Server {
245	+ Server::<M> {
246	hostname: hostname.to_string(),
247	..Default::default()
248	}
249	@@ -108,10 +127,7 @@ impl Server {
250
251	/// Append one or more milters which will be applied to messages after the
252	/// session has been closed but before they are enqued for delivery.
253	- pub fn with_milter<T>(mut self, milter: T) -> Self
254	- where
255	- T: Milter + 'static,
256	- {
257	+ pub fn with_milter(mut self, milter: M) -> Self {
258	self.milters.push(Arc::new(milter));
259	self
260	}
261	@@ -170,9 +186,55 @@ impl Server {
262	Ok(())
263	}
264
265	- pub async fn listen(&self) -> Result<(), Error> {
266	+ async fn spawn_workers(&mut self, global_queue: Arc<Injector<Packet>>) {
267	+ let local_queues: Vec<WorkQueue<Packet>> = (0..self.n_threads)
268	+ .map(\|_\| WorkQueue::<Packet>::new_fifo())
269	+ .collect();
270	+ let stealers: Vec<Stealer<Packet>> = local_queues
271	+ .iter()
272	+ .map(\|local_queue\| local_queue.stealer())
273	+ .collect();
274	+ let handles: Vec<JoinHandle<_>> = local_queues
275	+ .into_iter()
276	+ .map(\|local_queue\| {
277	+ let (tx, shutdown_rx) = tokio::sync::mpsc::channel::<bool>(1);
278	+ self.shutdown_handles.push(tx);
279	+ let global_queue = global_queue.clone();
280	+ let stealers = stealers.clone();
281	+ let milters = self.milters.clone();
282	+ tokio::task::spawn(async move {
283	+ let mut worker = Worker {
284	+ milters,
285	+ global_queue,
286	+ stealers,
287	+ local_queue: Arc::new(Mutex::new(local_queue)),
288	+ shutdown_rx,
289	+ };
290	+ worker.process().await
291	+ })
292	+ })
293	+ .collect();
294	+
295	+ // Log a message anytime a worker stops for any reason
296	+ tokio::spawn(async move {
297	+ stream::iter(handles)
298	+ .for_each(\|handle\| async move {
299	+ let worker_result = handle.await.unwrap();
300	+ if let Err(err) = worker_result {
301	+ tracing::warn!("Worker shutdown with error: {}", err);
302	+ } else {
303	+ tracing::info!("Worker shutdown normally");
304	+ }
305	+ })
306	+ .await;
307	+ });
308	+ }
309	+
310	+ pub async fn listen(&mut self) -> Result<(), Error> {
311	let listener = TcpListener::bind(&self.address).await?;
312	tracing::info!("Mail server listening @ {}", self.address);
313	+ let global_queue = Arc::new(Injector::<Packet>::new());
314	+ self.spawn_workers(global_queue.clone()).await;
315	loop {
316	let (socket, _) = listener.accept().await.unwrap();
317	let addr = socket.local_addr()?;
318	@@ -192,13 +254,8 @@ impl Server {
319
320	match self.process(framed, &mut pipelined, greeting).await {
321	Ok(_) => {
322	- let session = session.lock().await;
323	- let message = session.body.as_ref().expect("Session has no body");
324	- tracing::info!("session concluded successfully");
325	- // FIXME: Pass into queue and actually process appropriately
326	- for milter in self.milters.clone() {
327	- milter.apply(message).await.unwrap();
328	- }
329	+ let session = session.into_inner();
330	+ global_queue.push(session.into());
331	}
332	Err(err) => {
333	tracing::warn!("Client encountered an error: {:?}", err);
334	@@ -270,6 +327,7 @@ mod test {
335	}
336	}
337
338	+ /*
339	#[tokio::test]
340	async fn test_server_process() {
341	let stream = FakeStream {
342	@@ -309,4 +367,5 @@ mod test {
343	.first()
344	.is_some_and(\|rcpt_to\| rcpt_to.email() == "baz@qux.com")));
345	}
346	+ */
347	}
348	diff --git a/maitred/src/session.rs b/maitred/src/session.rs
349	index 58da7d1..7f898e9 100644
350	--- a/maitred/src/session.rs
351	+++ b/maitred/src/session.rs
352	@@ -13,7 +13,7 @@ use url::Host;
353	use crate::expand::Expansion;
354	use crate::transport::Response;
355	use crate::verify::Verify;
356	- use crate::{smtp_chunk, smtp_chunk_err, smtp_chunk_ok, Milter};
357	+ use crate::{smtp_chunk, smtp_chunk_err, smtp_chunk_ok};
358	use crate::{smtp_response, Chunk};
359
360	/// Default help banner returned from a HELP command without any parameters
361	@@ -42,7 +42,7 @@ pub const DEFAULT_GREETING: &str = "Maitred ESMTP Server";
362	/// Default SMTP capabilities advertised by the server
363	pub const DEFAULT_CAPABILITIES: u32 = smtp_proto::EXT_SIZE
364	\| smtp_proto::EXT_ENHANCED_STATUS_CODES
365	- \| smtp_proto::EXT_PIPELINING
366	+ // \| smtp_proto::EXT_PIPELINING FIXME broken in swaks
367	\| smtp_proto::EXT_8BIT_MIME;
368
369	/// Result generated as part of an SMTP session, an Err indicates a session
370	diff --git a/maitred/src/worker.rs b/maitred/src/worker.rs
371	new file mode 100644
372	index 0000000..a0fd512
373	--- /dev/null
374	+++ b/maitred/src/worker.rs
375	@@ -0,0 +1,109 @@
376	+ use std::sync::Arc;
377	+ use std::{iter, time::Duration};
378	+
379	+ use crossbeam_deque::{Injector, Stealer, Worker as WorkQueue};
380	+ use email_address::EmailAddress;
381	+ use futures::StreamExt;
382	+ use mail_parser::Message;
383	+ use tokio::sync::{mpsc::Receiver, Mutex};
384	+ use tokio_stream::{self as stream};
385	+ use url::Host;
386	+
387	+ use crate::{Error, Milter, Session};
388	+
389	+ /// Session details to be passed internally for processing
390	+ #[derive(Clone, Debug)]
391	+ pub(crate) struct Packet {
392	+ pub body: Option<Message<'static>>,
393	+ pub mail_from: Option<EmailAddress>,
394	+ pub rcpt_to: Option<Vec<EmailAddress>>,
395	+ pub hostname: Option<Host>,
396	+ }
397	+
398	+ impl From<Session> for Packet {
399	+ fn from(value: Session) -> Self {
400	+ Packet {
401	+ body: value.body.clone(),
402	+ mail_from: value.mail_from.clone(),
403	+ rcpt_to: value.rcpt_to.clone(),
404	+ hostname: value.hostname.clone(),
405	+ }
406	+ }
407	+ }
408	+
409	+ /// Worker is responsible for all asynchronous message processing after a
410	+ /// session has been completed. It will handle the following operations:
411	+ ///
412	+ /// Sequentially applying milters in the order they were configured
413	+ /// Running DKIM verification
414	+ /// ARC Verficiation
415	+ /// SPF Verification
416	+ pub(crate) struct Worker<M>
417	+ where
418	+ M: Milter + Clone + Send + Sync + 'static,
419	+ {
420	+ pub milters: Vec<Arc<M>>,
421	+ pub global_queue: Arc<Injector<Packet>>,
422	+ pub stealers: Vec<Stealer<Packet>>,
423	+ pub local_queue: Arc<Mutex<WorkQueue<Packet>>>,
424	+ pub shutdown_rx: Receiver<bool>,
425	+ }
426	+
427	+ impl<M> Worker<M>
428	+ where
429	+ M: Milter + Clone + Send + Sync + 'static,
430	+ {
431	+ async fn next_packet(&self) -> Option<Packet> {
432	+ let local_queue = self.local_queue.lock().await;
433	+ local_queue.pop().or_else(\|\| {
434	+ iter::repeat_with(\|\| {
435	+ self.global_queue
436	+ .steal_batch_and_pop(&local_queue)
437	+ .or_else(\|\| self.stealers.iter().map(\|s\| s.steal()).collect())
438	+ })
439	+ .find(\|s\| !s.is_retry())
440	+ .and_then(\|s\| s.success())
441	+ })
442	+ }
443	+
444	+ pub async fn process(&mut self) -> Result<(), Error> {
445	+ let mut ticker =
446	+ tokio::time::interval_at(tokio::time::Instant::now(), Duration::from_millis(800));
447	+
448	+ loop {
449	+ if let Ok(Some(_)) =
450	+ tokio::time::timeout(Duration::from_millis(100), self.shutdown_rx.recv()).await
451	+ {
452	+ break Ok(());
453	+ }
454	+
455	+ if let Some(packet) = self.next_packet().await {
456	+ let message = packet.body.unwrap();
457	+ // apply all of the milters to the message returning the final
458	+ // result. Any failure prevents the modified message from being
459	+ // returned and will result in a rejection.
460	+ let modified: Result<Message<'static>, crate::milter::Error> =
461	+ stream::iter(self.milters.clone())
462	+ .fold(Ok(message), \|accm, milter\| async move {
463	+ if let Ok(message) = accm {
464	+ milter.apply(&message).await
465	+ } else {
466	+ accm
467	+ }
468	+ })
469	+ .await;
470	+
471	+ match modified {
472	+ Ok(message) => {
473	+ tracing::info!("Message finished: {:?}", message);
474	+ }
475	+ Err(err) => {
476	+ tracing::warn!("Milter failed: {}", err)
477	+ }
478	+ }
479	+ } else {
480	+ ticker.tick().await;
481	+ }
482	+ }
483	+ }
484	+ }
485	diff --git a/scripts/swaks_test.sh b/scripts/swaks_test.sh
486	index 20f2da4..e73450b 100755
487	--- a/scripts/swaks_test.sh
488	+++ b/scripts/swaks_test.sh
489	@@ -3,4 +3,4 @@
490	# Uses swaks: https://www.jetmore.org/john/code/swaks/ to do some basic SMTP
491	# verification. Make sure you install the tool first!
492
493	- swaks --to hello@example.com --server localhost:2525
494	+ printf "Subject: Hello\nWorld\n" \| swaks --to hello@example.com --server localhost:2525 --pipeline --data -