@smbcloud / smbcloud-mail-imap
Code Commits Issues Pulls Cloud Agent CI/CD
main
copilot/add-pop3-protocol-examples copilot/add-pop3-protocol-implementation copilot/implement-mail-handling-separation development main
rs 137 lines 5.63 KB
Raw
1 use crate::send_commands::send_commands;
2 use crate::storage::MessageStorage;
3 use futures::StreamExt;
4 use tokio::io::{AsyncRead, AsyncWrite};
5 use tokio_util::codec::{Framed, LinesCodec};
6
7 enum SmtpState {
8 Command,
9 Data,
10 Quit,
11 }
12
13 pub async fn handle_session<S>(stream: S, storage: Option<MessageStorage>) -> anyhow::Result<()>
14 where
15 S: AsyncRead + AsyncWrite + Unpin,
16 {
17 let re_smtp_mail = regex::Regex::new(r"(?i)from: ?<(.+)>").unwrap();
18 let re_smtp_rcpt = regex::Regex::new(r"(?i)to: ?<(.+)>").unwrap();
19 let mut message = String::new();
20 let mut state = SmtpState::Command;
21 let mut mailfrom: Option<String> = None;
22 let mut rcpts: Vec<String> = Vec::new();
23 let mut framed = Framed::new(stream, LinesCodec::new());
24 while let Some(line_str) = framed.next().await {
25 let line = line_str?;
26 match state {
27 SmtpState::Command => {
28 let space_pos = line.find(" ").unwrap_or(line.len());
29 let (command, arg) = line.split_at(space_pos);
30 let arg = arg.trim();
31 match &*command.trim().to_uppercase() {
32 "HELO" | "EHLO" => {
33 send_commands(&mut framed, vec!["250 Hello".to_string()]).await?;
34 }
35 "MAIL" => {
36 // Handle MAIL FROM command
37 if let Some(address) = re_smtp_mail.captures(arg).and_then(|cap| cap.get(1))
38 {
39 mailfrom = Some(address.as_str().to_string());
40 send_commands(&mut framed, vec!["250 OK".to_string()]).await?;
41 } else {
42 send_commands(
43 &mut framed,
44 vec!["501 Syntax: MAIL From: <address>".to_string()],
45 )
46 .await?;
47 }
48 }
49 "RCPT" => {
50 // Handle RCPT TO command
51 if mailfrom.is_none() {
52 send_commands(
53 &mut framed,
54 vec!["503 Error: Send MAIL first".to_string()],
55 )
56 .await?;
57 } else if let Some(address) =
58 re_smtp_rcpt.captures(arg).and_then(|cap| cap.get(1))
59 {
60 rcpts.push(address.as_str().to_string());
61 send_commands(&mut framed, vec!["250 OK".to_string()]).await?;
62 } else {
63 send_commands(
64 &mut framed,
65 vec!["501 Syntax: RCPT TO: <address>".to_string()],
66 )
67 .await?;
68 }
69 }
70 "DATA" => {
71 if rcpts.is_empty() {
72 send_commands(&mut framed, vec!["503 Error: MAIL FROM and RCPT TO must be set before sending DATA".to_string()]).await?;
73 } else {
74 state = SmtpState::Data;
75 send_commands(
76 &mut framed,
77 vec!["354 End data with <CR><LF>.<CR><LF>".to_string()],
78 )
79 .await?;
80 }
81 }
82 "NOOP" => {
83 send_commands(&mut framed, vec!["250 OK".to_string()]).await?;
84 }
85 "RSET" => {
86 mailfrom = None;
87 rcpts = Vec::new();
88 message = String::new();
89 send_commands(&mut framed, vec!["250 OK".to_string()]).await?;
90 }
91 "QUIT" => {
92 send_commands(&mut framed, vec!["221 Bye".to_string()]).await?;
93 state = SmtpState::Quit;
94 }
95 _ => {
96 send_commands(&mut framed, vec!["500 Unknown command".to_string()]).await?;
97 }
98 }
99 }
100 SmtpState::Data => {
101 if line.trim() == "." {
102 // The end of the email content has been received
103 send_commands(&mut framed, vec!["250 OK".to_string()]).await?;
104
105 // Log the received email
106 tracing::info!(
107 "Email received - From: {:?}, To: {:?}, Size: {} bytes",
108 mailfrom,
109 rcpts,
110 message.len()
111 );
112
113 // Store the message if storage is available
114 if let Some(ref storage) = storage {
115 if let Some(ref from) = mailfrom {
116 storage.store_message(from.clone(), rcpts.clone(), message.clone());
117 }
118 }
119
120 // reset the state and variables for the next email
121 mailfrom = None;
122 rcpts = Vec::new();
123 message = String::new();
124 state = SmtpState::Command;
125 } else {
126 // Add the received line to the email content
127 message.push_str(&line);
128 message.push('\n');
129 }
130 }
131 SmtpState::Quit => {
132 break;
133 }
134 }
135 }
136 Ok(())
137 }