Fix last touchups on SmtpServer and add test

Fix some mistakes in the SmtpServer implementation and add test.

This test is only for the ideal interaction, it doesn't cover all checks
for the malformed or out of order commands. This is left for later.

Signed-off-by: Mroik <mroik@delayed.space>

1 files changed, 158 insertions, 1 deletions

diff --git a/src/smtp_server.rs b/src/smtp_server.rs
index 83a1093..bee00e9 100644
--- a/src/smtp_server.rs
+++ b/src/smtp_server.rs
@@ -1,4 +1,7 @@
-use std::net::{IpAddr, SocketAddr};
+use std::{
+    fmt::Display,
+    net::{IpAddr, SocketAddr},
+};
 
 use anyhow::Result;
 use tokio::{
@@ -44,6 +47,7 @@ impl SmtpServer {
                     }
                 },
             };
+            log::info!("Connected: {}", addr);
             let session = SessionHandler {
                 addr,
                 server_name: self.server_name.clone(),
@@ -101,6 +105,9 @@ impl SessionHandler {
             if self.state == SessionState::AwaitingMailInput {
                 if buffer.starts_with(".\r\n") {
                     self.process_mail().await?;
+                    writer
+                        .write_all(Reply::Completed("Ok").to_string().as_bytes())
+                        .await?;
                     continue;
                 }
 
@@ -117,13 +124,19 @@ impl SessionHandler {
                 }
                 Ok(v) => v,
             };
+            log::debug!("COMMAND: {}", command);
             let res = self.apply(command).await?;
+            log::debug!("RESPONSE: {}", res);
             writer.write_all(res.to_string().as_bytes()).await?;
             if res == Reply::EndTransmission {
                 break;
             }
         }
 
+        // Should wait for the client before closing but meh. Any transaction has already gone
+        // through anyway so... waiting for the client to close the connection as part of the
+        // protocol was a dumb choice, the client can't issue new commands anyway.
+        stream.shutdown().await?;
         Ok(())
     }
 
@@ -292,3 +305,147 @@ impl<'a> TryFrom<&'a str> for Command<'a> {
         Err(anyhow::format_err!("Invalid command"))
     }
 }
+
+impl Display for Command<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Command::HELO(host) => write!(f, "HELO {}\r\n", host),
+            Command::MAIL(from) => write!(f, "MAIL FROM:<{}>\r\n", from),
+            Command::RCPT(to) => write!(f, "RCPT TO:<{}>\r\n", to),
+            Command::DATA => write!(f, "DATA\r\n"),
+            Command::QUIT => write!(f, "QUIT\r\n"),
+        }
+    }
+}
+
+mod tests {
+    use std::net::SocketAddr;
+
+    use tokio::{
+        io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
+        net::TcpSocket,
+        spawn,
+        sync::mpsc::channel,
+    };
+
+    use crate::smtp_server::{Command, SmtpServer};
+
+    const DOMAIN: &str = "mail.example.local";
+    const CLIENT_HOST: &str = "exploit";
+
+    #[tokio::test]
+    async fn receive_email() {
+        env_logger::init();
+        let mut server = SmtpServer::new([127, 0, 0, 1], 9090, DOMAIN).await.unwrap();
+        let (tx, mut rx) = channel(5);
+        let j = spawn(async move {
+            server.run(tx).await.unwrap();
+            log::info!("IM DONE");
+        });
+
+        let sock = TcpSocket::new_v4().unwrap();
+        let mut s = sock
+            .connect(SocketAddr::new([127, 0, 0, 1].into(), 9090))
+            .await
+            .unwrap();
+        let (reader, mut writer) = s.split();
+        let mut reader = BufReader::new(reader);
+        let mut buf = String::new();
+
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), format!("220 {}", DOMAIN));
+
+        writer
+            .write_all(Command::HELO(CLIENT_HOST).to_string().as_bytes())
+            .await
+            .unwrap();
+        buf.clear();
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), format!("250 {}", DOMAIN));
+
+        writer
+            .write_all(Command::MAIL("mroik@poul.org").to_string().as_bytes())
+            .await
+            .unwrap();
+        buf.clear();
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), "250 Ok");
+
+        writer
+            .write_all(Command::RCPT("mroik@delayed.space").to_string().as_bytes())
+            .await
+            .unwrap();
+        buf.clear();
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), "250 Ok");
+
+        writer
+            .write_all(Command::DATA.to_string().as_bytes())
+            .await
+            .unwrap();
+        buf.clear();
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), "354 <CRLF>.<CRLF>");
+
+        writer.write_all(b"From: mroik@poul.org\r\n").await.unwrap();
+        writer
+            .write_all(b"To: mroik@delayed.space\r\n")
+            .await
+            .unwrap();
+        writer
+            .write_all(b"Date: Tue, 14 Apr 2026 23:00:00 +0200\r\n")
+            .await
+            .unwrap();
+        writer
+            .write_all(b"Subject: This is a test\r\n")
+            .await
+            .unwrap();
+        writer.write_all(b"\r\n").await.unwrap();
+        writer
+            .write_all(b"Hey, this is a test. You can ignore this!\r\n")
+            .await
+            .unwrap();
+        writer.write_all(b"\r\n").await.unwrap();
+        writer.write_all(b"Mroik\r\n").await.unwrap();
+        writer.write_all(b".\r\n").await.unwrap();
+        buf.clear();
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), "250 Ok");
+
+        let lines = [
+            "From: mroik@poul.org",
+            "To: mroik@delayed.space",
+            "Date: Tue, 14 Apr 2026 23:00:00 +0200",
+            "Subject: This is a test",
+            "",
+            "Hey, this is a test. You can ignore this!",
+            "",
+            "Mroik",
+        ];
+
+        let email = rx.recv().await.unwrap();
+        assert_eq!(email.from, "mroik@poul.org");
+        assert_eq!(email.recipient.len(), 1);
+        assert_eq!(email.recipient[0], "mroik@delayed.space");
+        assert_eq!(email.data.trim(), lines.join("\r\n").trim());
+
+        // Close connection
+        writer
+            .write_all(Command::QUIT.to_string().as_bytes())
+            .await
+            .unwrap();
+        buf.clear();
+        assert!(reader.read_line(&mut buf).await.unwrap() > 0);
+        assert_eq!(buf.trim(), "221 OK");
+
+        buf.clear();
+        assert_eq!(reader.read_line(&mut buf).await.unwrap(), 0);
+
+        s.shutdown().await.unwrap();
+        // This is necessary because the server continues to run even after closing a connection as
+        // it continues to listen for new connections.
+        //
+        // In order to ensure that this doesn't affect other tests we have to abort it.
+        j.abort();
+    }
+}