smartmta/rust/crates/mailer-smtp/src/command.rs

//! SMTP command parser.
//!
//! Parses raw SMTP command lines into structured `SmtpCommand` variants.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A parsed SMTP command.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum SmtpCommand {
    /// EHLO with client hostname/IP
    Ehlo(String),
    /// HELO with client hostname/IP
    Helo(String),
    /// MAIL FROM with sender address and optional parameters (e.g. SIZE=12345)
    MailFrom {
        address: String,
        params: HashMap<String, Option<String>>,
    },
    /// RCPT TO with recipient address and optional parameters
    RcptTo {
        address: String,
        params: HashMap<String, Option<String>>,
    },
    /// DATA command — begin message body
    Data,
    /// RSET — reset current transaction
    Rset,
    /// NOOP — no operation
    Noop,
    /// QUIT — close connection
    Quit,
    /// STARTTLS — upgrade to TLS
    StartTls,
    /// AUTH with mechanism and optional initial response
    Auth {
        mechanism: AuthMechanism,
        initial_response: Option<String>,
    },
    /// HELP with optional topic
    Help(Option<String>),
    /// VRFY with address or username
    Vrfy(String),
    /// EXPN with mailing list name
    Expn(String),
}

/// Supported AUTH mechanisms.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum AuthMechanism {
    Plain,
    Login,
}

/// Errors that can occur during command parsing.
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
pub enum ParseError {
    #[error("empty command line")]
    Empty,
    #[error("unrecognized command: {0}")]
    UnrecognizedCommand(String),
    #[error("syntax error in parameters: {0}")]
    SyntaxError(String),
    #[error("missing required argument for {0}")]
    MissingArgument(String),
}

/// Parse a raw SMTP command line (without trailing CRLF) into a `SmtpCommand`.
pub fn parse_command(line: &str) -> Result<SmtpCommand, ParseError> {
    let line = line.trim_end_matches('\r').trim_end_matches('\n');
    if line.is_empty() {
        return Err(ParseError::Empty);
    }

    // Split into verb and the rest
    let (verb, rest) = split_first_word(line);
    let verb_upper = verb.to_ascii_uppercase();

    match verb_upper.as_str() {
        "EHLO" => {
            let hostname = rest.trim();
            if hostname.is_empty() {
                return Err(ParseError::MissingArgument("EHLO".into()));
            }
            Ok(SmtpCommand::Ehlo(hostname.to_string()))
        }
        "HELO" => {
            let hostname = rest.trim();
            if hostname.is_empty() {
                return Err(ParseError::MissingArgument("HELO".into()));
            }
            Ok(SmtpCommand::Helo(hostname.to_string()))
        }
        "MAIL" => parse_mail_from(rest),
        "RCPT" => parse_rcpt_to(rest),
        "DATA" => Ok(SmtpCommand::Data),
        "RSET" => Ok(SmtpCommand::Rset),
        "NOOP" => Ok(SmtpCommand::Noop),
        "QUIT" => Ok(SmtpCommand::Quit),
        "STARTTLS" => Ok(SmtpCommand::StartTls),
        "AUTH" => parse_auth(rest),
        "HELP" => {
            let topic = rest.trim();
            if topic.is_empty() {
                Ok(SmtpCommand::Help(None))
            } else {
                Ok(SmtpCommand::Help(Some(topic.to_string())))
            }
        }
        "VRFY" => {
            let arg = rest.trim();
            if arg.is_empty() {
                return Err(ParseError::MissingArgument("VRFY".into()));
            }
            Ok(SmtpCommand::Vrfy(arg.to_string()))
        }
        "EXPN" => {
            let arg = rest.trim();
            if arg.is_empty() {
                return Err(ParseError::MissingArgument("EXPN".into()));
            }
            Ok(SmtpCommand::Expn(arg.to_string()))
        }
        _ => Err(ParseError::UnrecognizedCommand(verb_upper)),
    }
}

/// Parse `FROM:<addr> [PARAM=VALUE ...]` after "MAIL".
fn parse_mail_from(rest: &str) -> Result<SmtpCommand, ParseError> {
    // Expect "FROM:" prefix (case-insensitive, whitespace-flexible)
    let rest = rest.trim_start();
    let rest_upper = rest.to_ascii_uppercase();
    if !rest_upper.starts_with("FROM") {
        return Err(ParseError::SyntaxError(
            "expected FROM after MAIL".into(),
        ));
    }
    let rest = &rest[4..]; // skip "FROM"
    let rest = rest.trim_start();
    if !rest.starts_with(':') {
        return Err(ParseError::SyntaxError(
            "expected colon after MAIL FROM".into(),
        ));
    }
    let rest = &rest[1..]; // skip ':'
    let rest = rest.trim_start();

    parse_address_and_params(rest, "MAIL FROM").map(|(address, params)| SmtpCommand::MailFrom {
        address,
        params,
    })
}

/// Parse `TO:<addr> [PARAM=VALUE ...]` after "RCPT".
fn parse_rcpt_to(rest: &str) -> Result<SmtpCommand, ParseError> {
    let rest = rest.trim_start();
    let rest_upper = rest.to_ascii_uppercase();
    if !rest_upper.starts_with("TO") {
        return Err(ParseError::SyntaxError("expected TO after RCPT".into()));
    }
    let rest = &rest[2..]; // skip "TO"
    let rest = rest.trim_start();
    if !rest.starts_with(':') {
        return Err(ParseError::SyntaxError(
            "expected colon after RCPT TO".into(),
        ));
    }
    let rest = &rest[1..]; // skip ':'
    let rest = rest.trim_start();

    parse_address_and_params(rest, "RCPT TO").map(|(address, params)| SmtpCommand::RcptTo {
        address,
        params,
    })
}

/// Parse `<address> [PARAM=VALUE ...]` from the rest of a MAIL FROM or RCPT TO line.
fn parse_address_and_params(
    input: &str,
    context: &str,
) -> Result<(String, HashMap<String, Option<String>>), ParseError> {
    if !input.starts_with('<') {
        return Err(ParseError::SyntaxError(format!(
            "expected '<' in {context}"
        )));
    }
    let close_bracket = input.find('>').ok_or_else(|| {
        ParseError::SyntaxError(format!("missing '>' in {context}"))
    })?;
    let address = input[1..close_bracket].to_string();
    let remainder = &input[close_bracket + 1..];
    let params = parse_params(remainder)?;
    Ok((address, params))
}

/// Parse SMTP extension parameters like `SIZE=12345 BODY=8BITMIME`.
fn parse_params(input: &str) -> Result<HashMap<String, Option<String>>, ParseError> {
    let mut params = HashMap::new();
    for token in input.split_whitespace() {
        if let Some(eq_pos) = token.find('=') {
            let key = token[..eq_pos].to_ascii_uppercase();
            let value = token[eq_pos + 1..].to_string();
            params.insert(key, Some(value));
        } else {
            params.insert(token.to_ascii_uppercase(), None);
        }
    }
    Ok(params)
}

/// Parse AUTH command: `AUTH <mechanism> [initial-response]`.
fn parse_auth(rest: &str) -> Result<SmtpCommand, ParseError> {
    let rest = rest.trim();
    if rest.is_empty() {
        return Err(ParseError::MissingArgument("AUTH".into()));
    }
    let (mech_str, initial) = split_first_word(rest);
    let mechanism = match mech_str.to_ascii_uppercase().as_str() {
        "PLAIN" => AuthMechanism::Plain,
        "LOGIN" => AuthMechanism::Login,
        other => {
            return Err(ParseError::SyntaxError(format!(
                "unsupported AUTH mechanism: {other}"
            )));
        }
    };
    let initial_response = {
        let s = initial.trim();
        if s.is_empty() { None } else { Some(s.to_string()) }
    };
    Ok(SmtpCommand::Auth {
        mechanism,
        initial_response,
    })
}

/// Split a string into the first whitespace-delimited word and the remainder.
fn split_first_word(s: &str) -> (&str, &str) {
    match s.find(char::is_whitespace) {
        Some(pos) => (&s[..pos], &s[pos + 1..]),
        None => (s, ""),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ehlo() {
        let cmd = parse_command("EHLO mail.example.com").unwrap();
        assert_eq!(cmd, SmtpCommand::Ehlo("mail.example.com".into()));
    }

    #[test]
    fn test_ehlo_case_insensitive() {
        let cmd = parse_command("ehlo MAIL.EXAMPLE.COM").unwrap();
        assert_eq!(cmd, SmtpCommand::Ehlo("MAIL.EXAMPLE.COM".into()));
    }

    #[test]
    fn test_helo() {
        let cmd = parse_command("HELO example.com").unwrap();
        assert_eq!(cmd, SmtpCommand::Helo("example.com".into()));
    }

    #[test]
    fn test_ehlo_missing_arg() {
        let err = parse_command("EHLO").unwrap_err();
        assert!(matches!(err, ParseError::MissingArgument(_)));
    }

    #[test]
    fn test_mail_from() {
        let cmd = parse_command("MAIL FROM:<sender@example.com>").unwrap();
        assert_eq!(
            cmd,
            SmtpCommand::MailFrom {
                address: "sender@example.com".into(),
                params: HashMap::new(),
            }
        );
    }

    #[test]
    fn test_mail_from_with_params() {
        let cmd = parse_command("MAIL FROM:<sender@example.com> SIZE=12345 BODY=8BITMIME").unwrap();
        if let SmtpCommand::MailFrom { address, params } = cmd {
            assert_eq!(address, "sender@example.com");
            assert_eq!(params.get("SIZE"), Some(&Some("12345".into())));
            assert_eq!(params.get("BODY"), Some(&Some("8BITMIME".into())));
        } else {
            panic!("expected MailFrom");
        }
    }

    #[test]
    fn test_mail_from_empty_address() {
        let cmd = parse_command("MAIL FROM:<>").unwrap();
        assert_eq!(
            cmd,
            SmtpCommand::MailFrom {
                address: "".into(),
                params: HashMap::new(),
            }
        );
    }

    #[test]
    fn test_mail_from_flexible_spacing() {
        let cmd = parse_command("MAIL FROM: <user@example.com>").unwrap();
        if let SmtpCommand::MailFrom { address, .. } = cmd {
            assert_eq!(address, "user@example.com");
        } else {
            panic!("expected MailFrom");
        }
    }

    #[test]
    fn test_rcpt_to() {
        let cmd = parse_command("RCPT TO:<recipient@example.com>").unwrap();
        assert_eq!(
            cmd,
            SmtpCommand::RcptTo {
                address: "recipient@example.com".into(),
                params: HashMap::new(),
            }
        );
    }

    #[test]
    fn test_data() {
        assert_eq!(parse_command("DATA").unwrap(), SmtpCommand::Data);
    }

    #[test]
    fn test_rset() {
        assert_eq!(parse_command("RSET").unwrap(), SmtpCommand::Rset);
    }

    #[test]
    fn test_noop() {
        assert_eq!(parse_command("NOOP").unwrap(), SmtpCommand::Noop);
    }

    #[test]
    fn test_quit() {
        assert_eq!(parse_command("QUIT").unwrap(), SmtpCommand::Quit);
    }

    #[test]
    fn test_starttls() {
        assert_eq!(parse_command("STARTTLS").unwrap(), SmtpCommand::StartTls);
    }

    #[test]
    fn test_auth_plain() {
        let cmd = parse_command("AUTH PLAIN dGVzdAB0ZXN0AHBhc3N3b3Jk").unwrap();
        assert_eq!(
            cmd,
            SmtpCommand::Auth {
                mechanism: AuthMechanism::Plain,
                initial_response: Some("dGVzdAB0ZXN0AHBhc3N3b3Jk".into()),
            }
        );
    }

    #[test]
    fn test_auth_login_no_initial() {
        let cmd = parse_command("AUTH LOGIN").unwrap();
        assert_eq!(
            cmd,
            SmtpCommand::Auth {
                mechanism: AuthMechanism::Login,
                initial_response: None,
            }
        );
    }

    #[test]
    fn test_help() {
        assert_eq!(parse_command("HELP").unwrap(), SmtpCommand::Help(None));
        assert_eq!(
            parse_command("HELP MAIL").unwrap(),
            SmtpCommand::Help(Some("MAIL".into()))
        );
    }

    #[test]
    fn test_vrfy() {
        assert_eq!(
            parse_command("VRFY user@example.com").unwrap(),
            SmtpCommand::Vrfy("user@example.com".into())
        );
    }

    #[test]
    fn test_expn() {
        assert_eq!(
            parse_command("EXPN staff").unwrap(),
            SmtpCommand::Expn("staff".into())
        );
    }

    #[test]
    fn test_empty() {
        assert!(matches!(parse_command(""), Err(ParseError::Empty)));
    }

    #[test]
    fn test_unrecognized() {
        let err = parse_command("FOOBAR test").unwrap_err();
        assert!(matches!(err, ParseError::UnrecognizedCommand(_)));
    }

    #[test]
    fn test_crlf_stripped() {
        let cmd = parse_command("QUIT\r\n").unwrap();
        assert_eq!(cmd, SmtpCommand::Quit);
    }
}