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

//! SMTP protocol state machine.
//!
//! Defines valid states and transitions for an SMTP session.

use serde::{Deserialize, Serialize};

/// SMTP session states following RFC 5321.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum SmtpState {
    /// Initial state — waiting for server greeting.
    Connected,
    /// After successful EHLO/HELO.
    Greeted,
    /// After MAIL FROM accepted.
    MailFrom,
    /// After at least one RCPT TO accepted.
    RcptTo,
    /// In DATA mode — accumulating message body.
    Data,
    /// Transaction completed — can start a new one or QUIT.
    Finished,
}

/// State transition errors.
#[derive(Debug, Clone, PartialEq, thiserror::Error)]
pub enum TransitionError {
    #[error("cannot {action} in state {state:?}")]
    InvalidTransition {
        state: SmtpState,
        action: &'static str,
    },
}

impl SmtpState {
    /// Check whether EHLO/HELO is valid in the current state.
    /// EHLO/HELO can be issued at any time to reset the session.
    pub fn can_ehlo(&self) -> bool {
        true
    }

    /// Check whether MAIL FROM is valid in the current state.
    pub fn can_mail_from(&self) -> bool {
        matches!(self, SmtpState::Greeted | SmtpState::Finished)
    }

    /// Check whether RCPT TO is valid in the current state.
    pub fn can_rcpt_to(&self) -> bool {
        matches!(self, SmtpState::MailFrom | SmtpState::RcptTo)
    }

    /// Check whether DATA is valid in the current state.
    pub fn can_data(&self) -> bool {
        matches!(self, SmtpState::RcptTo)
    }

    /// Check whether STARTTLS is valid in the current state.
    /// Only before a transaction starts.
    pub fn can_starttls(&self) -> bool {
        matches!(self, SmtpState::Connected | SmtpState::Greeted | SmtpState::Finished)
    }

    /// Check whether AUTH is valid in the current state.
    /// Only after EHLO and before a transaction starts.
    pub fn can_auth(&self) -> bool {
        matches!(self, SmtpState::Greeted | SmtpState::Finished)
    }

    /// Transition to Greeted state (after EHLO/HELO).
    pub fn transition_ehlo(&self) -> Result<SmtpState, TransitionError> {
        // EHLO is always valid — it resets the session.
        Ok(SmtpState::Greeted)
    }

    /// Transition to MailFrom state (after MAIL FROM accepted).
    pub fn transition_mail_from(&self) -> Result<SmtpState, TransitionError> {
        if self.can_mail_from() {
            Ok(SmtpState::MailFrom)
        } else {
            Err(TransitionError::InvalidTransition {
                state: *self,
                action: "MAIL FROM",
            })
        }
    }

    /// Transition to RcptTo state (after RCPT TO accepted).
    pub fn transition_rcpt_to(&self) -> Result<SmtpState, TransitionError> {
        if self.can_rcpt_to() {
            Ok(SmtpState::RcptTo)
        } else {
            Err(TransitionError::InvalidTransition {
                state: *self,
                action: "RCPT TO",
            })
        }
    }

    /// Transition to Data state (after DATA command accepted).
    pub fn transition_data(&self) -> Result<SmtpState, TransitionError> {
        if self.can_data() {
            Ok(SmtpState::Data)
        } else {
            Err(TransitionError::InvalidTransition {
                state: *self,
                action: "DATA",
            })
        }
    }

    /// Transition to Finished state (after end-of-data).
    pub fn transition_finished(&self) -> Result<SmtpState, TransitionError> {
        if *self == SmtpState::Data {
            Ok(SmtpState::Finished)
        } else {
            Err(TransitionError::InvalidTransition {
                state: *self,
                action: "finish DATA",
            })
        }
    }

    /// Reset to Greeted state (after RSET command).
    pub fn transition_rset(&self) -> Result<SmtpState, TransitionError> {
        match self {
            SmtpState::Connected => Err(TransitionError::InvalidTransition {
                state: *self,
                action: "RSET",
            }),
            _ => Ok(SmtpState::Greeted),
        }
    }
}

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

    #[test]
    fn test_initial_state() {
        let state = SmtpState::Connected;
        assert!(!state.can_mail_from());
        assert!(!state.can_rcpt_to());
        assert!(!state.can_data());
        assert!(state.can_starttls());
        assert!(state.can_ehlo());
    }

    #[test]
    fn test_ehlo_always_valid() {
        for state in [
            SmtpState::Connected,
            SmtpState::Greeted,
            SmtpState::MailFrom,
            SmtpState::RcptTo,
            SmtpState::Data,
            SmtpState::Finished,
        ] {
            assert!(state.can_ehlo());
            assert!(state.transition_ehlo().is_ok());
        }
    }

    #[test]
    fn test_normal_flow() {
        let state = SmtpState::Connected;
        let state = state.transition_ehlo().unwrap();
        assert_eq!(state, SmtpState::Greeted);

        let state = state.transition_mail_from().unwrap();
        assert_eq!(state, SmtpState::MailFrom);

        let state = state.transition_rcpt_to().unwrap();
        assert_eq!(state, SmtpState::RcptTo);

        // Multiple RCPT TO
        let state = state.transition_rcpt_to().unwrap();
        assert_eq!(state, SmtpState::RcptTo);

        let state = state.transition_data().unwrap();
        assert_eq!(state, SmtpState::Data);

        let state = state.transition_finished().unwrap();
        assert_eq!(state, SmtpState::Finished);

        // New transaction
        let state = state.transition_mail_from().unwrap();
        assert_eq!(state, SmtpState::MailFrom);
    }

    #[test]
    fn test_invalid_transitions() {
        assert!(SmtpState::Connected.transition_mail_from().is_err());
        assert!(SmtpState::Connected.transition_rcpt_to().is_err());
        assert!(SmtpState::Connected.transition_data().is_err());
        assert!(SmtpState::Greeted.transition_rcpt_to().is_err());
        assert!(SmtpState::Greeted.transition_data().is_err());
        assert!(SmtpState::MailFrom.transition_data().is_err());
    }

    #[test]
    fn test_rset() {
        let state = SmtpState::RcptTo;
        let state = state.transition_rset().unwrap();
        assert_eq!(state, SmtpState::Greeted);

        // RSET from Connected is invalid (no EHLO yet)
        assert!(SmtpState::Connected.transition_rset().is_err());
    }

    #[test]
    fn test_starttls_validity() {
        assert!(SmtpState::Connected.can_starttls());
        assert!(SmtpState::Greeted.can_starttls());
        assert!(!SmtpState::MailFrom.can_starttls());
        assert!(!SmtpState::RcptTo.can_starttls());
        assert!(!SmtpState::Data.can_starttls());
        assert!(SmtpState::Finished.can_starttls());
    }
}