1use core::fmt;
42
43use alloc::{
44 string::{String, ToString},
45 vec::Vec,
46};
47
48use log::trace;
49use thiserror::Error;
50
51use crate::{coroutine::*, rfc5321::types::reply_code::ReplyCode, send::*, smtp_try};
52
53pub struct SmtpDataCommand;
55
56impl From<SmtpDataCommand> for Vec<u8> {
57 fn from(_: SmtpDataCommand) -> Vec<u8> {
58 b"DATA\r\n".to_vec()
59 }
60}
61
62pub struct SmtpDataBody(pub Vec<u8>);
65
66impl From<SmtpDataBody> for Vec<u8> {
67 fn from(body: SmtpDataBody) -> Vec<u8> {
68 body.0
69 }
70}
71
72#[derive(Clone, Debug, Error)]
74pub enum SmtpDataError {
75 #[error("SMTP DATA command failed: rejected {code} {message}")]
76 CommandRejected { code: u16, message: String },
77 #[error("SMTP DATA body failed: rejected {code} {message}")]
78 BodyRejected { code: u16, message: String },
79 #[error("SMTP DATA failed: {0}")]
80 Send(#[from] SendSmtpCommandError),
81}
82
83pub struct SmtpData {
85 state: State,
86 body: Option<Vec<u8>>,
87}
88
89impl SmtpData {
90 pub fn new(message: Vec<u8>) -> Self {
93 Self {
94 state: State::SendCommand(SendSmtpCommand::new(SmtpDataCommand)),
95 body: Some(message),
96 }
97 }
98
99 fn prepare_body(message: Vec<u8>) -> Vec<u8> {
102 let mut result = Vec::with_capacity(message.len() + 5);
103
104 let mut at_line_start = true;
105 for &byte in &message {
106 if at_line_start && byte == b'.' {
107 result.push(b'.');
108 }
109 result.push(byte);
110 at_line_start = byte == b'\n';
111 }
112
113 if !result.ends_with(b"\r\n") {
114 if result.ends_with(b"\n") {
115 result.pop();
116 result.extend_from_slice(b"\r\n");
117 } else if result.ends_with(b"\r") {
118 result.push(b'\n');
119 } else {
120 result.extend_from_slice(b"\r\n");
121 }
122 }
123
124 result.extend_from_slice(b".\r\n");
125
126 result
127 }
128}
129
130impl SmtpCoroutine for SmtpData {
131 type Yield = SmtpYield;
132 type Return = Result<(), SmtpDataError>;
133
134 fn resume(&mut self, arg: Option<&[u8]>) -> SmtpCoroutineState<Self::Yield, Self::Return> {
135 loop {
136 trace!("data: {}", self.state);
137
138 match &mut self.state {
139 State::SendCommand(send) => {
140 let out = smtp_try!(send, arg);
141
142 if out.response.code != ReplyCode::START_MAIL_INPUT {
143 let code = out.response.code.code();
144 let message = out.response.text().to_string();
145 return SmtpCoroutineState::Complete(Err(SmtpDataError::CommandRejected {
146 code,
147 message,
148 }));
149 }
150
151 let body = self.body.take().expect("body taken twice");
152 let prepared = Self::prepare_body(body);
153 trace!("message body prepared: {} bytes", prepared.len());
154
155 self.state = State::SendBody(SendSmtpCommand::new(SmtpDataBody(prepared)));
156 }
157 State::SendBody(send) => {
158 let out = smtp_try!(send, arg);
159
160 if out.response.code == ReplyCode::OK {
161 return SmtpCoroutineState::Complete(Ok(()));
162 }
163
164 let code = out.response.code.code();
165 let message = out.response.text().to_string();
166 return SmtpCoroutineState::Complete(Err(SmtpDataError::BodyRejected {
167 code,
168 message,
169 }));
170 }
171 }
172 }
173 }
174}
175
176enum State {
177 SendCommand(SendSmtpCommand<SmtpDataCommand>),
178 SendBody(SendSmtpCommand<SmtpDataBody>),
179}
180
181impl fmt::Display for State {
182 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183 match self {
184 Self::SendCommand(_) => f.write_str("send data"),
185 Self::SendBody(_) => f.write_str("send body"),
186 }
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193
194 #[test]
195 fn success_returns_ok() {
196 let mut data = SmtpData::new(b"Subject: hi\r\n\r\nbody\r\n".to_vec());
197
198 let bytes = expect_wants_write(&mut data, None);
199 assert_eq!(bytes, b"DATA\r\n");
200
201 expect_wants_read(&mut data);
202 let body_bytes = expect_wants_write(&mut data, Some(b"354 send body\r\n"));
203 assert!(body_bytes.ends_with(b"\r\n.\r\n"));
204
205 expect_wants_read(&mut data);
206 expect_complete_ok(&mut data, b"250 message accepted\r\n");
207 }
208
209 #[test]
210 fn command_rejected_returns_command_error() {
211 let mut data = SmtpData::new(b"hi\r\n".to_vec());
212 let _ = expect_wants_write(&mut data, None);
213 expect_wants_read(&mut data);
214
215 let err = expect_complete_err(&mut data, b"503 bad sequence\r\n");
216 let SmtpDataError::CommandRejected { code, message } = err else {
217 panic!("expected SmtpDataError::CommandRejected, got {err:?}");
218 };
219 assert_eq!(code, 503);
220 assert_eq!(message, "bad sequence");
221 }
222
223 #[test]
224 fn body_rejected_returns_body_error() {
225 let mut data = SmtpData::new(b"hi\r\n".to_vec());
226 let _ = expect_wants_write(&mut data, None);
227 expect_wants_read(&mut data);
228 let _ = expect_wants_write(&mut data, Some(b"354 send body\r\n"));
229 expect_wants_read(&mut data);
230
231 let err = expect_complete_err(&mut data, b"552 too large\r\n");
232 let SmtpDataError::BodyRejected { code, message } = err else {
233 panic!("expected SmtpDataError::BodyRejected, got {err:?}");
234 };
235 assert_eq!(code, 552);
236 assert_eq!(message, "too large");
237 }
238
239 #[test]
240 fn dot_stuffs_leading_dot_lines() {
241 let body = SmtpData::prepare_body(b".hello\r\n".to_vec());
242 assert_eq!(body, b"..hello\r\n.\r\n");
243 }
244
245 #[test]
246 fn eof_returns_eof_error() {
247 let mut data = SmtpData::new(b"hi\r\n".to_vec());
248 let _ = expect_wants_write(&mut data, None);
249 expect_wants_read(&mut data);
250
251 let err = expect_complete_err(&mut data, b"");
252 assert!(matches!(
253 err,
254 SmtpDataError::Send(SendSmtpCommandError::Eof)
255 ));
256 }
257
258 fn expect_wants_write(cor: &mut SmtpData, arg: Option<&[u8]>) -> Vec<u8> {
261 match cor.resume(arg) {
262 SmtpCoroutineState::Yielded(SmtpYield::WantsWrite(bytes)) => bytes,
263 state => panic!("expected WantsWrite, got {state:?}"),
264 }
265 }
266
267 fn expect_wants_read(cor: &mut SmtpData) {
268 match cor.resume(None) {
269 SmtpCoroutineState::Yielded(SmtpYield::WantsRead) => {}
270 state => panic!("expected WantsRead, got {state:?}"),
271 }
272 }
273
274 fn expect_complete_ok(cor: &mut SmtpData, reply: &[u8]) {
275 match cor.resume(Some(reply)) {
276 SmtpCoroutineState::Complete(Ok(())) => {}
277 state => panic!("expected Complete(Ok), got {state:?}"),
278 }
279 }
280
281 fn expect_complete_err(cor: &mut SmtpData, reply: &[u8]) -> SmtpDataError {
282 match cor.resume(Some(reply)) {
283 SmtpCoroutineState::Complete(Err(err)) => err,
284 state => panic!("expected Complete(Err), got {state:?}"),
285 }
286 }
287}