smartmta/rust/crates/mailer-smtp/src/client/protocol.rs

//! SMTP client protocol engine.
//!
//! Implements the SMTP command/response flow for sending outbound email.

use super::config::SmtpAuthConfig;
use super::connection::ClientSmtpStream;
use super::error::SmtpClientError;
use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine;
use serde::{Deserialize, Serialize};
use tokio::time::{timeout, Duration};
use tracing::debug;

/// Parsed SMTP response (from the remote server).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SmtpClientResponse {
    pub code: u16,
    pub lines: Vec<String>,
}

impl SmtpClientResponse {
    pub fn is_success(&self) -> bool {
        self.code >= 200 && self.code < 300
    }

    pub fn is_positive_intermediate(&self) -> bool {
        self.code >= 300 && self.code < 400
    }

    pub fn is_temp_error(&self) -> bool {
        self.code >= 400 && self.code < 500
    }

    pub fn is_perm_error(&self) -> bool {
        self.code >= 500
    }

    /// Full response text (all lines joined).
    pub fn full_message(&self) -> String {
        self.lines.join(" ")
    }

    /// Convert to a protocol error if this is an error response.
    pub fn to_error(&self) -> SmtpClientError {
        SmtpClientError::ProtocolError {
            code: self.code,
            message: self.full_message(),
        }
    }
}

/// Server capabilities parsed from EHLO response.
#[derive(Debug, Clone, Default)]
pub struct EhloCapabilities {
    pub extensions: Vec<String>,
    pub max_size: Option<u64>,
    pub starttls: bool,
    pub auth_methods: Vec<String>,
    pub pipelining: bool,
    pub eight_bit_mime: bool,
}

/// Read a multi-line SMTP response from the server.
pub async fn read_response(
    stream: &mut ClientSmtpStream,
    timeout_secs: u64,
) -> Result<SmtpClientResponse, SmtpClientError> {
    let mut lines = Vec::new();
    let mut code: u16;

    loop {
        let mut line = String::new();
        let n = timeout(
            Duration::from_secs(timeout_secs),
            stream.read_line(&mut line),
        )
        .await
        .map_err(|_| SmtpClientError::TimeoutError {
            message: format!("Timeout reading SMTP response after {timeout_secs}s"),
        })??;

        if n == 0 {
            return Err(SmtpClientError::ConnectionError {
                message: "Connection closed while reading response".into(),
            });
        }

        // Guard against unbounded lines from malicious servers (RFC 5321 §4.5.3.1.4 says 512 max)
        if line.len() > 4096 {
            return Err(SmtpClientError::ProtocolError {
                code: 0,
                message: format!("Response line too long ({} bytes, max 4096)", line.len()),
            });
        }

        let line = line.trim_end_matches('\n').trim_end_matches('\r');

        if line.len() < 3 {
            return Err(SmtpClientError::ProtocolError {
                code: 0,
                message: format!("Invalid response line: {line}"),
            });
        }

        // Parse the 3-digit code
        let parsed_code: u16 = line[..3].parse().map_err(|_| SmtpClientError::ProtocolError {
            code: 0,
            message: format!("Invalid response code in: {line}"),
        })?;
        code = parsed_code;

        // Text after the code (skip the separator character)
        let text = if line.len() > 4 { &line[4..] } else { "" };
        lines.push(text.to_string());

        // Check for continuation: "250-" means more lines, "250 " means last line
        if line.len() >= 4 && line.as_bytes()[3] == b'-' {
            continue;
        } else {
            break;
        }
    }

    debug!("SMTP response: {} {}", code, lines.join(" | "));
    Ok(SmtpClientResponse { code, lines })
}

/// Read the server greeting (first response after connect).
pub async fn read_greeting(
    stream: &mut ClientSmtpStream,
    timeout_secs: u64,
) -> Result<SmtpClientResponse, SmtpClientError> {
    let resp = read_response(stream, timeout_secs).await?;
    if resp.code == 220 {
        Ok(resp)
    } else {
        Err(SmtpClientError::ProtocolError {
            code: resp.code,
            message: format!("Unexpected greeting: {}", resp.full_message()),
        })
    }
}

/// Send a raw command and read the response.
async fn send_command(
    stream: &mut ClientSmtpStream,
    command: &str,
    timeout_secs: u64,
) -> Result<SmtpClientResponse, SmtpClientError> {
    debug!("SMTP C: {}", command);
    stream
        .write_all(format!("{command}\r\n").as_bytes())
        .await?;
    stream.flush().await?;
    read_response(stream, timeout_secs).await
}

/// Send EHLO and parse capabilities.
pub async fn send_ehlo(
    stream: &mut ClientSmtpStream,
    domain: &str,
    timeout_secs: u64,
) -> Result<EhloCapabilities, SmtpClientError> {
    let resp = send_command(stream, &format!("EHLO {domain}"), timeout_secs).await?;

    if !resp.is_success() {
        // Fall back to HELO
        let helo_resp = send_command(stream, &format!("HELO {domain}"), timeout_secs).await?;
        if !helo_resp.is_success() {
            return Err(helo_resp.to_error());
        }
        return Ok(EhloCapabilities::default());
    }

    let mut caps = EhloCapabilities::default();

    // First line is the greeting, remaining lines are capabilities
    for line in resp.lines.iter().skip(1) {
        let upper = line.to_uppercase();
        if upper.starts_with("SIZE ") {
            caps.max_size = upper[5..].trim().parse().ok();
        } else if upper == "STARTTLS" {
            caps.starttls = true;
        } else if upper.starts_with("AUTH ") {
            caps.auth_methods = upper[5..]
                .split_whitespace()
                .map(|s| s.to_string())
                .collect();
        } else if upper == "PIPELINING" {
            caps.pipelining = true;
        } else if upper == "8BITMIME" {
            caps.eight_bit_mime = true;
        }
        caps.extensions.push(line.clone());
    }

    Ok(caps)
}

/// Send STARTTLS command (does not perform the TLS handshake itself).
pub async fn send_starttls(
    stream: &mut ClientSmtpStream,
    timeout_secs: u64,
) -> Result<(), SmtpClientError> {
    let resp = send_command(stream, "STARTTLS", timeout_secs).await?;
    if resp.code != 220 {
        return Err(SmtpClientError::ProtocolError {
            code: resp.code,
            message: format!("STARTTLS rejected: {}", resp.full_message()),
        });
    }
    Ok(())
}

/// Authenticate using AUTH PLAIN.
pub async fn send_auth_plain(
    stream: &mut ClientSmtpStream,
    user: &str,
    pass: &str,
    timeout_secs: u64,
) -> Result<(), SmtpClientError> {
    // AUTH PLAIN sends \0user\0pass in base64
    let credentials = format!("\x00{user}\x00{pass}");
    let encoded = BASE64.encode(credentials.as_bytes());
    let resp = send_command(stream, &format!("AUTH PLAIN {encoded}"), timeout_secs).await?;

    if resp.code != 235 {
        return Err(SmtpClientError::AuthenticationError {
            message: format!("AUTH PLAIN failed ({}): {}", resp.code, resp.full_message()),
        });
    }
    Ok(())
}

/// Authenticate using AUTH LOGIN.
pub async fn send_auth_login(
    stream: &mut ClientSmtpStream,
    user: &str,
    pass: &str,
    timeout_secs: u64,
) -> Result<(), SmtpClientError> {
    // Step 1: Send AUTH LOGIN
    let resp = send_command(stream, "AUTH LOGIN", timeout_secs).await?;
    if resp.code != 334 {
        return Err(SmtpClientError::AuthenticationError {
            message: format!(
                "AUTH LOGIN challenge failed ({}): {}",
                resp.code,
                resp.full_message()
            ),
        });
    }

    // Step 2: Send base64 username
    let user_b64 = BASE64.encode(user.as_bytes());
    let resp = send_command(stream, &user_b64, timeout_secs).await?;
    if resp.code != 334 {
        return Err(SmtpClientError::AuthenticationError {
            message: format!(
                "AUTH LOGIN username rejected ({}): {}",
                resp.code,
                resp.full_message()
            ),
        });
    }

    // Step 3: Send base64 password
    let pass_b64 = BASE64.encode(pass.as_bytes());
    let resp = send_command(stream, &pass_b64, timeout_secs).await?;
    if resp.code != 235 {
        return Err(SmtpClientError::AuthenticationError {
            message: format!(
                "AUTH LOGIN password rejected ({}): {}",
                resp.code,
                resp.full_message()
            ),
        });
    }

    Ok(())
}

/// Authenticate using the configured method.
pub async fn authenticate(
    stream: &mut ClientSmtpStream,
    auth: &SmtpAuthConfig,
    _caps: &EhloCapabilities,
    timeout_secs: u64,
) -> Result<(), SmtpClientError> {
    match auth.method.to_uppercase().as_str() {
        "LOGIN" => send_auth_login(stream, &auth.user, &auth.pass, timeout_secs).await,
        _ => send_auth_plain(stream, &auth.user, &auth.pass, timeout_secs).await,
    }
}

/// Send MAIL FROM.
pub async fn send_mail_from(
    stream: &mut ClientSmtpStream,
    sender: &str,
    timeout_secs: u64,
) -> Result<SmtpClientResponse, SmtpClientError> {
    let resp = send_command(stream, &format!("MAIL FROM:<{sender}>"), timeout_secs).await?;
    if !resp.is_success() {
        return Err(resp.to_error());
    }
    Ok(resp)
}

/// Send RCPT TO. Returns per-recipient success/failure.
pub async fn send_rcpt_to(
    stream: &mut ClientSmtpStream,
    recipient: &str,
    timeout_secs: u64,
) -> Result<SmtpClientResponse, SmtpClientError> {
    let resp = send_command(stream, &format!("RCPT TO:<{recipient}>"), timeout_secs).await?;
    // We don't fail the entire send on per-recipient errors;
    // the caller decides based on the response code.
    Ok(resp)
}

/// Send MAIL FROM + RCPT TO commands in a single pipelined batch.
///
/// Writes all envelope commands at once, then reads responses in order.
/// Returns `(mail_from_ok, accepted_recipients, rejected_recipients)`.
pub async fn send_pipelined_envelope(
    stream: &mut ClientSmtpStream,
    sender: &str,
    recipients: &[String],
    timeout_secs: u64,
) -> Result<(bool, Vec<String>, Vec<String>), SmtpClientError> {
    // Build the full pipelined command batch
    let mut batch = format!("MAIL FROM:<{sender}>\r\n");
    for rcpt in recipients {
        batch.push_str(&format!("RCPT TO:<{rcpt}>\r\n"));
    }

    // Send all commands at once
    debug!("SMTP C (pipelined): MAIL FROM + {} RCPT TO", recipients.len());
    stream.write_all(batch.as_bytes()).await?;
    stream.flush().await?;

    // Read MAIL FROM response
    let mail_resp = read_response(stream, timeout_secs).await?;
    if !mail_resp.is_success() {
        return Err(mail_resp.to_error());
    }

    // Read RCPT TO responses
    let mut accepted = Vec::new();
    let mut rejected = Vec::new();
    for rcpt in recipients {
        match read_response(stream, timeout_secs).await {
            Ok(resp) => {
                if resp.is_success() {
                    accepted.push(rcpt.clone());
                } else {
                    rejected.push(rcpt.clone());
                }
            }
            Err(_) => {
                rejected.push(rcpt.clone());
            }
        }
    }

    Ok((true, accepted, rejected))
}

/// Send DATA command, followed by the message body with dot-stuffing.
pub async fn send_data(
    stream: &mut ClientSmtpStream,
    message: &[u8],
    timeout_secs: u64,
) -> Result<SmtpClientResponse, SmtpClientError> {
    // Send DATA command
    let resp = send_command(stream, "DATA", timeout_secs).await?;
    if !resp.is_positive_intermediate() {
        return Err(resp.to_error());
    }

    // Send the message body with dot-stuffing
    let stuffed = dot_stuff(message);
    stream.write_all(&stuffed).await?;

    // Send terminator: CRLF.CRLF
    // If the message doesn't end with CRLF, add one
    if !stuffed.ends_with(b"\r\n") {
        stream.write_all(b"\r\n").await?;
    }
    stream.write_all(b".\r\n").await?;
    stream.flush().await?;

    // Read final response
    let final_resp = read_response(stream, timeout_secs).await?;
    if !final_resp.is_success() {
        return Err(final_resp.to_error());
    }

    Ok(final_resp)
}

/// Send RSET command to reset the server state between messages on a reused connection.
pub async fn send_rset(
    stream: &mut ClientSmtpStream,
    timeout_secs: u64,
) -> Result<(), SmtpClientError> {
    let resp = send_command(stream, "RSET", timeout_secs).await?;
    if !resp.is_success() {
        return Err(resp.to_error());
    }
    Ok(())
}

/// Send QUIT command.
pub async fn send_quit(
    stream: &mut ClientSmtpStream,
    timeout_secs: u64,
) -> Result<(), SmtpClientError> {
    // Best-effort QUIT — ignore errors since we're closing anyway
    let _ = send_command(stream, "QUIT", timeout_secs).await;
    Ok(())
}

/// Apply SMTP dot-stuffing to a message body.
///
/// Any line starting with a period gets an extra period prepended.
/// Also normalizes bare LF to CRLF.
pub fn dot_stuff(data: &[u8]) -> Vec<u8> {
    let mut result = Vec::with_capacity(data.len() + data.len() / 40);
    let mut at_line_start = true;

    for i in 0..data.len() {
        let byte = data[i];

        // Normalize bare LF to CRLF
        if byte == b'\n' && (i == 0 || data[i - 1] != b'\r') {
            result.push(b'\r');
            result.push(b'\n');
            at_line_start = true;
            continue;
        }

        // Dot-stuff: add extra dot at start of line
        if at_line_start && byte == b'.' {
            result.push(b'.');
        }

        result.push(byte);
        at_line_start = byte == b'\n';
    }

    result
}

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

    #[test]
    fn test_dot_stuffing_basic() {
        assert_eq!(
            dot_stuff(b"Hello\r\n.World\r\n"),
            b"Hello\r\n..World\r\n"
        );
    }

    #[test]
    fn test_dot_stuffing_leading_dot() {
        assert_eq!(dot_stuff(b".starts with dot\r\n"), b"..starts with dot\r\n");
    }

    #[test]
    fn test_dot_stuffing_multiple_dots() {
        assert_eq!(
            dot_stuff(b"ok\r\n.line1\r\n..line2\r\n"),
            b"ok\r\n..line1\r\n...line2\r\n"
        );
    }

    #[test]
    fn test_dot_stuffing_bare_lf() {
        assert_eq!(
            dot_stuff(b"line1\nline2\n"),
            b"line1\r\nline2\r\n"
        );
    }

    #[test]
    fn test_dot_stuffing_bare_lf_with_dot() {
        assert_eq!(
            dot_stuff(b"ok\n.dotline\n"),
            b"ok\r\n..dotline\r\n"
        );
    }

    #[test]
    fn test_dot_stuffing_no_change() {
        assert_eq!(
            dot_stuff(b"Hello World\r\nNo dots here\r\n"),
            b"Hello World\r\nNo dots here\r\n"
        );
    }

    #[test]
    fn test_dot_stuffing_empty() {
        assert_eq!(dot_stuff(b""), b"");
    }

    #[test]
    fn test_response_is_success() {
        let resp = SmtpClientResponse {
            code: 250,
            lines: vec!["OK".into()],
        };
        assert!(resp.is_success());
        assert!(!resp.is_temp_error());
        assert!(!resp.is_perm_error());
    }

    #[test]
    fn test_response_temp_error() {
        let resp = SmtpClientResponse {
            code: 450,
            lines: vec!["Mailbox busy".into()],
        };
        assert!(!resp.is_success());
        assert!(resp.is_temp_error());
    }

    #[test]
    fn test_response_perm_error() {
        let resp = SmtpClientResponse {
            code: 550,
            lines: vec!["No such user".into()],
        };
        assert!(!resp.is_success());
        assert!(resp.is_perm_error());
    }

    #[test]
    fn test_response_positive_intermediate() {
        let resp = SmtpClientResponse {
            code: 354,
            lines: vec!["Start mail input".into()],
        };
        assert!(resp.is_positive_intermediate());
        assert!(!resp.is_success());
    }

    #[test]
    fn test_response_full_message() {
        let resp = SmtpClientResponse {
            code: 250,
            lines: vec!["OK".into(), "SIZE 10485760".into()],
        };
        assert_eq!(resp.full_message(), "OK SIZE 10485760");
    }

    #[test]
    fn test_ehlo_capabilities_default() {
        let caps = EhloCapabilities::default();
        assert!(!caps.starttls);
        assert!(!caps.pipelining);
        assert!(!caps.eight_bit_mime);
        assert!(caps.auth_methods.is_empty());
        assert!(caps.max_size.is_none());
    }
}