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