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

//! Per-connection SMTP session state.
//!
//! Tracks the envelope, authentication, TLS status, and counters
//! for a single SMTP connection.

use crate::state::SmtpState;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

/// Envelope accumulator for the current mail transaction.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Envelope {
    /// Sender address from MAIL FROM.
    pub mail_from: String,
    /// Recipient addresses from RCPT TO.
    pub rcpt_to: Vec<String>,
    /// Declared message size from MAIL FROM SIZE= param (if any).
    pub declared_size: Option<u64>,
    /// BODY parameter (e.g. "8BITMIME").
    pub body_type: Option<String>,
}

/// Authentication state for the session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AuthState {
    /// Not authenticated and not in progress.
    None,
    /// Waiting for AUTH credentials (LOGIN flow step).
    WaitingForUsername,
    /// Have username, waiting for password.
    WaitingForPassword { username: String },
    /// Successfully authenticated.
    Authenticated { username: String },
}

impl Default for AuthState {
    fn default() -> Self {
        AuthState::None
    }
}

/// Per-connection session state.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SmtpSession {
    /// Unique session identifier.
    pub id: String,
    /// Current protocol state.
    pub state: SmtpState,
    /// Client's EHLO/HELO hostname.
    pub client_hostname: Option<String>,
    /// Whether the client used EHLO (vs HELO).
    pub esmtp: bool,
    /// Whether the connection is using TLS.
    pub secure: bool,
    /// Authentication state.
    pub auth_state: AuthState,
    /// Current transaction envelope.
    pub envelope: Envelope,
    /// Remote IP address.
    pub remote_addr: String,
    /// Number of messages sent in this session.
    pub message_count: u32,
    /// Number of failed auth attempts.
    pub auth_failures: u32,
    /// Number of invalid commands.
    pub invalid_commands: u32,
    /// Maximum allowed invalid commands before disconnect.
    pub max_invalid_commands: u32,
}

impl SmtpSession {
    /// Create a new session for a connection.
    pub fn new(remote_addr: String, secure: bool) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            state: SmtpState::Connected,
            client_hostname: None,
            esmtp: false,
            secure,
            auth_state: AuthState::None,
            envelope: Envelope::default(),
            remote_addr,
            message_count: 0,
            auth_failures: 0,
            invalid_commands: 0,
            max_invalid_commands: 20,
        }
    }

    /// Reset the current transaction (RSET), preserving connection state.
    pub fn reset_transaction(&mut self) {
        self.envelope = Envelope::default();
        if self.state != SmtpState::Connected {
            self.state = SmtpState::Greeted;
        }
    }

    /// Reset session for a new EHLO (preserves counters and TLS).
    pub fn reset_for_ehlo(&mut self, hostname: String, esmtp: bool) {
        self.client_hostname = Some(hostname);
        self.esmtp = esmtp;
        self.envelope = Envelope::default();
        self.state = SmtpState::Greeted;
        // Auth state is reset on new EHLO per RFC
        self.auth_state = AuthState::None;
    }

    /// Check if the client is authenticated.
    pub fn is_authenticated(&self) -> bool {
        matches!(self.auth_state, AuthState::Authenticated { .. })
    }

    /// Get the authenticated username, if any.
    pub fn authenticated_user(&self) -> Option<&str> {
        match &self.auth_state {
            AuthState::Authenticated { username } => Some(username),
            _ => None,
        }
    }

    /// Record a completed message delivery.
    pub fn record_message(&mut self) {
        self.message_count += 1;
    }

    /// Record a failed auth attempt. Returns true if limit exceeded.
    pub fn record_auth_failure(&mut self, max_failures: u32) -> bool {
        self.auth_failures += 1;
        self.auth_failures >= max_failures
    }

    /// Record an invalid command. Returns true if limit exceeded.
    pub fn record_invalid_command(&mut self) -> bool {
        self.invalid_commands += 1;
        self.invalid_commands >= self.max_invalid_commands
    }
}

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

    #[test]
    fn test_new_session() {
        let session = SmtpSession::new("127.0.0.1".into(), false);
        assert_eq!(session.state, SmtpState::Connected);
        assert!(!session.secure);
        assert!(!session.is_authenticated());
        assert!(session.client_hostname.is_none());
    }

    #[test]
    fn test_reset_transaction() {
        let mut session = SmtpSession::new("127.0.0.1".into(), false);
        session.state = SmtpState::RcptTo;
        session.envelope.mail_from = "sender@example.com".into();
        session.envelope.rcpt_to.push("rcpt@example.com".into());

        session.reset_transaction();

        assert_eq!(session.state, SmtpState::Greeted);
        assert!(session.envelope.mail_from.is_empty());
        assert!(session.envelope.rcpt_to.is_empty());
    }

    #[test]
    fn test_reset_for_ehlo() {
        let mut session = SmtpSession::new("127.0.0.1".into(), true);
        session.auth_state = AuthState::Authenticated {
            username: "user".into(),
        };

        session.reset_for_ehlo("mail.example.com".into(), true);

        assert_eq!(session.state, SmtpState::Greeted);
        assert_eq!(session.client_hostname.as_deref(), Some("mail.example.com"));
        assert!(session.esmtp);
        assert!(!session.is_authenticated()); // Auth reset after EHLO
    }

    #[test]
    fn test_auth_failures() {
        let mut session = SmtpSession::new("127.0.0.1".into(), false);
        assert!(!session.record_auth_failure(3));
        assert!(!session.record_auth_failure(3));
        assert!(session.record_auth_failure(3)); // 3rd failure -> limit
    }

    #[test]
    fn test_invalid_commands() {
        let mut session = SmtpSession::new("127.0.0.1".into(), false);
        session.max_invalid_commands = 3;
        assert!(!session.record_invalid_command());
        assert!(!session.record_invalid_command());
        assert!(session.record_invalid_command()); // 3rd -> limit
    }

    #[test]
    fn test_message_count() {
        let mut session = SmtpSession::new("127.0.0.1".into(), false);
        assert_eq!(session.message_count, 0);
        session.record_message();
        session.record_message();
        assert_eq!(session.message_count, 2);
    }
}