path: root/src/smtp_server.rs

use std::{
    io::{BufRead, BufReader, Write},
    net::{IpAddr, SocketAddr, TcpListener, TcpStream},
};

use anyhow::Result;
use tokio::{spawn, sync::mpsc::Sender};

use crate::model::Mail;

const SERVER_NAME: &str = "";

pub struct SmtpServer {
    listener: TcpListener,
    running: bool,
}

impl SmtpServer {
    pub fn new(ip: [u8; 4], port: u16) -> Result<Self> {
        Ok(Self {
            listener: TcpListener::bind((IpAddr::from(ip), port))?,
            running: false,
        })
    }

    // TODO: trap SIGINT to stop server?
    pub async fn run(&mut self, tx_processor: Sender<Mail>) -> Result<()> {
        self.running = true;
        while self.running {
            let (stream, addr) = self.listener.accept()?;
            let session = SessionHandler {
                addr,
                tx_processor: tx_processor.clone(),
                state: SessionState::default(),
                client_host: String::new(),
                from: None,
                to: Vec::new(),
                data: String::new(),
            };
            spawn(session.run(stream));
        }
        Ok(())
    }
}

struct SessionHandler {
    addr: SocketAddr,
    tx_processor: Sender<Mail>,
    state: SessionState,
    client_host: String,
    from: Option<String>,
    to: Vec<String>,
    data: String,
}

impl SessionHandler {
    async fn run(mut self, stream: TcpStream) -> Result<()> {
        let mut writer = stream.try_clone()?;
        let mut r = BufReader::new(&stream);
        let mut buffer = String::new();

        writer.write_all(
            Reply::Ready(String::from(SERVER_NAME))
                .to_string()
                .as_bytes(),
        )?;

        loop {
            buffer.clear();
            if r.read_line(&mut buffer)? == 0 {
                break;
            }
            log::debug!("Received '{}' from '{}'", buffer.trim(), self.addr);

            // In this state the server is not expecting commands as the input is part of the email
            // message
            if self.state == SessionState::AwaitingMailInput {
                if buffer.starts_with(".\r\n") {
                    self.process_mail().await?;
                    continue;
                }

                self.data.push_str(&buffer);
                continue;
            }

            let command = match Command::try_from(buffer.as_str()) {
                Err(_) => {
                    writer.write_all(Reply::InvalidCommand.to_string().as_bytes())?;
                    continue;
                }
                Ok(v) => v,
            };
            let res = self.apply(command).await?;
            writer.write_all(res.to_string().as_bytes())?;
        }

        log::info!("Connection closed by {}", self.addr);
        Ok(())
    }

    async fn apply(&mut self, command: Command) -> Result<Reply> {
        match &command {
            Command::HELO(hostname) => self.helo(hostname).await,
            Command::MAIL(from) => self.mail(from).await,
            Command::RCPT(to) => self.rcpt(to).await,
            Command::DATA => self.start_data().await,
        }
    }

    /// HELO resets buffers
    async fn helo(&mut self, hostname: &str) -> Result<Reply> {
        self.client_host = String::from(hostname);
        self.state = SessionState::Normal;
        self.from = None;
        self.to.clear();
        Ok(Reply::Completed(String::from(SERVER_NAME)))
    }

    /// TODO: Validate email address
    /// MAIL resets buffers
    async fn mail(&mut self, from: &str) -> Result<Reply> {
        if from.is_empty() {
            return Ok(Reply::InvalidParameter);
        }

        self.from = Some(String::from(from));
        self.to.clear();
        self.state = SessionState::MailTransaction;

        Ok(Reply::Completed(String::from("Ok")))
    }

    /// TODO: Validate email addresses
    /// Only after having started a mail transaction
    async fn rcpt(&mut self, to: &str) -> Result<Reply> {
        if self.state != SessionState::MailTransaction {
            return Ok(Reply::BadSequence);
        }

        if to.is_empty() {
            return Ok(Reply::InvalidParameter);
        }

        self.to.push(String::from(to));
        Ok(Reply::Completed(String::from("Ok")))
    }

    /// Only after having started a mail transaction
    async fn start_data(&mut self) -> std::result::Result<Reply, anyhow::Error> {
        if self.state != SessionState::MailTransaction {
            return Ok(Reply::BadSequence);
        }

        // Should I return a more meaningful string with this error code?
        if self.from.is_none() || self.to.is_empty() {
            return Ok(Reply::InvalidCommand);
        }

        self.state = SessionState::AwaitingMailInput;

        Ok(Reply::StartMailInput)
    }

    async fn process_mail(&mut self) -> Result<Reply> {
        if self.from.is_none() || self.to.is_empty() || self.data.trim().is_empty() {
            return Ok(Reply::InvalidCommand);
        }

        let mail = Mail {
            from: self.from.as_ref().unwrap().clone(),
            recipient: self.to.clone(),
            data: self.data.clone(),
        };

        self.tx_processor.send(mail).await?;

        self.state = SessionState::Normal;
        self.from = None;
        self.to.clear();
        self.data.clear();

        Ok(Reply::Completed(String::from("Ok")))
    }
}

#[derive(PartialEq, Default)]
enum SessionState {
    #[default]
    WaitingHelo,
    MailTransaction,
    AwaitingMailInput,
    Normal,
}

enum Reply {
    Ready(String),
    Completed(String),
    StartMailInput,
    InvalidCommand,
    InvalidParameter,
    BadSequence,
}

impl std::fmt::Display for Reply {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Reply::Ready(hostname) => write!(f, "220 {}", hostname),
            Reply::Completed(hostname) => write!(f, "250 {}", hostname),
            Reply::InvalidCommand => write!(f, "500 Invalid command"),
            Reply::InvalidParameter => write!(f, "501 Parameter error"),
            Reply::BadSequence => write!(f, "503 Bad sequence of commands"),
            Reply::StartMailInput => write!(f, "354 <CRLF>.<CRLF>"),
        }
    }
}

enum Command {
    HELO(String),
    MAIL(String),
    RCPT(String),
    DATA,
}

impl TryFrom<&str> for Command {
    type Error = anyhow::Error;

    fn try_from(value: &str) -> std::result::Result<Self, Self::Error> {
        if value.len() >= 5 && value.to_lowercase().starts_with("helo ") {
            return Ok(Command::HELO(String::from(&value[5..])));
        }

        if value.len() >= 11
            && value.to_lowercase().starts_with("mail from:<")
            && value.contains('>')
        {
            let from = &value[11..value.find('>').unwrap()];
            return Ok(Command::MAIL(String::from(from)));
        }

        if value.len() >= 9 && value.to_lowercase().starts_with("rcpt to:<") && value.contains('>')
        {
            let to = &value[9..value.find('>').unwrap()];
            return Ok(Command::RCPT(String::from(to)));
        }

        if value.to_lowercase().starts_with("data") {
            return Ok(Command::DATA);
        }

        Err(anyhow::format_err!("Invalid command"))
    }
}