1 files changed, 16 insertions, 0 deletions
diff --git a/src/smtp_server.rs b/src/smtp_server.rs
index 6e48bce..42e5cf4 100644
--- a/src/smtp_server.rs
+++ b/src/smtp_server.rs
@@ -120,6 +120,9 @@ impl SessionHandler {
             };
             let res = self.apply(command).await?;
             writer.write_all(res.to_string().as_bytes()).await?;
+            if res == Reply::EndTransmission {
+                break;
+            }
         }
 
         log::info!("Connection closed by {}", self.addr);
@@ -132,6 +135,7 @@ impl SessionHandler {
             Command::MAIL(from) => self.mail(from).await,
             Command::RCPT(to) => self.rcpt(to).await,
             Command::DATA => self.start_data().await,
+            Command::QUIT => self.quit().await,
         }
     }
 
@@ -209,6 +213,10 @@ impl SessionHandler {
 
         Ok(Reply::Completed(String::from("Ok")))
     }
+
+    async fn quit(&self) -> Result<Reply> {
+        Ok(Reply::EndTransmission)
+    }
 }
 
 #[derive(PartialEq, Default)]
@@ -220,6 +228,7 @@ enum SessionState {
     Normal,
 }
 
+#[derive(PartialEq)]
 enum Reply {
     Ready(String),
     Completed(String),
@@ -227,6 +236,7 @@ enum Reply {
     InvalidCommand,
     InvalidParameter,
     BadSequence,
+    EndTransmission,
 }
 
 impl std::fmt::Display for Reply {
@@ -238,6 +248,7 @@ impl std::fmt::Display for Reply {
             Reply::InvalidParameter => write!(f, "501 Parameter error\r\n"),
             Reply::BadSequence => write!(f, "503 Bad sequence of commands\r\n"),
             Reply::StartMailInput => write!(f, "354 <CRLF>.<CRLF>\r\n"),
+            Reply::EndTransmission => write!(f, "221 OK\r\n"),
         }
     }
 }
@@ -247,6 +258,7 @@ enum Command {
     MAIL(String),
     RCPT(String),
     DATA,
+    QUIT,
 }
 
 impl TryFrom<&str> for Command {
@@ -275,6 +287,10 @@ impl TryFrom<&str> for Command {
             return Ok(Command::DATA);
         }
 
+        if value.to_lowercase().starts_with("quit") {
+            return Ok(Command::QUIT);
+        }
+
         Err(anyhow::format_err!("Invalid command"))
     }
 }