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

//! SMTP TCP/TLS server.
//!
//! Listens on configured ports, accepts connections, and dispatches
//! them to per-connection handlers.

use crate::config::SmtpServerConfig;
use crate::connection::{
    self, CallbackRegistry, ConnectionEvent, SmtpStream,
};
use crate::rate_limiter::{RateLimitConfig, RateLimiter};

use rustls_pki_types::{CertificateDer, PrivateKeyDer};
use std::io::BufReader;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use tokio::io::BufReader as TokioBufReader;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tracing::{error, info, warn};

/// Handle for a running SMTP server.
pub struct SmtpServerHandle {
    /// Shutdown signal.
    shutdown: Arc<AtomicBool>,
    /// Join handles for the listener tasks.
    handles: Vec<tokio::task::JoinHandle<()>>,
    /// Active connection count.
    pub active_connections: Arc<AtomicU32>,
}

impl SmtpServerHandle {
    /// Signal shutdown and wait for all listeners to stop.
    pub async fn shutdown(self) {
        self.shutdown.store(true, Ordering::SeqCst);
        for handle in self.handles {
            let _ = handle.await;
        }
        info!("SMTP server shut down");
    }

    /// Check if the server is running.
    pub fn is_running(&self) -> bool {
        !self.shutdown.load(Ordering::SeqCst)
    }
}

/// Start the SMTP server with the given configuration.
///
/// Returns a handle that can be used to shut down the server,
/// and an event receiver for connection events (emailReceived, authRequest).
pub async fn start_server(
    config: SmtpServerConfig,
    callback_registry: Arc<dyn CallbackRegistry + Send + Sync>,
    rate_limit_config: Option<RateLimitConfig>,
) -> Result<(SmtpServerHandle, mpsc::Receiver<ConnectionEvent>), Box<dyn std::error::Error + Send + Sync>>
{
    let config = Arc::new(config);
    let shutdown = Arc::new(AtomicBool::new(false));
    let active_connections = Arc::new(AtomicU32::new(0));
    let rate_limiter = Arc::new(RateLimiter::new(
        rate_limit_config.unwrap_or_default(),
    ));

    let (event_tx, event_rx) = mpsc::channel::<ConnectionEvent>(1024);

    // Build TLS acceptor if configured
    let tls_acceptor = if config.has_tls() {
        Some(Arc::new(build_tls_acceptor(&config)?))
    } else {
        None
    };

    let mut handles = Vec::new();

    // Start listeners on each port
    for &port in &config.ports {
        let listener = TcpListener::bind(format!("0.0.0.0:{port}")).await?;
        info!(port = port, "SMTP server listening (STARTTLS)");

        let handle = tokio::spawn(accept_loop(
            listener,
            config.clone(),
            shutdown.clone(),
            active_connections.clone(),
            rate_limiter.clone(),
            event_tx.clone(),
            callback_registry.clone(),
            tls_acceptor.clone(),
            false, // not implicit TLS
        ));
        handles.push(handle);
    }

    // Start implicit TLS listener if configured
    if let Some(secure_port) = config.secure_port {
        if tls_acceptor.is_some() {
            let listener =
                TcpListener::bind(format!("0.0.0.0:{secure_port}")).await?;
            info!(port = secure_port, "SMTP server listening (implicit TLS)");

            let handle = tokio::spawn(accept_loop(
                listener,
                config.clone(),
                shutdown.clone(),
                active_connections.clone(),
                rate_limiter.clone(),
                event_tx.clone(),
                callback_registry.clone(),
                tls_acceptor.clone(),
                true, // implicit TLS
            ));
            handles.push(handle);
        } else {
            warn!("Secure port configured but TLS certificates not provided");
        }
    }

    // Spawn periodic rate limiter cleanup
    {
        let rate_limiter = rate_limiter.clone();
        let shutdown = shutdown.clone();
        tokio::spawn(async move {
            let mut interval =
                tokio::time::interval(tokio::time::Duration::from_secs(60));
            loop {
                interval.tick().await;
                if shutdown.load(Ordering::SeqCst) {
                    break;
                }
                rate_limiter.cleanup();
            }
        });
    }

    Ok((
        SmtpServerHandle {
            shutdown,
            handles,
            active_connections,
        },
        event_rx,
    ))
}

/// Accept loop for a single listener.
async fn accept_loop(
    listener: TcpListener,
    config: Arc<SmtpServerConfig>,
    shutdown: Arc<AtomicBool>,
    active_connections: Arc<AtomicU32>,
    rate_limiter: Arc<RateLimiter>,
    event_tx: mpsc::Sender<ConnectionEvent>,
    callback_registry: Arc<dyn CallbackRegistry + Send + Sync>,
    tls_acceptor: Option<Arc<tokio_rustls::TlsAcceptor>>,
    implicit_tls: bool,
) {
    loop {
        if shutdown.load(Ordering::SeqCst) {
            break;
        }

        // Use a short timeout to check shutdown periodically
        let accept_result = tokio::time::timeout(
            tokio::time::Duration::from_secs(1),
            listener.accept(),
        )
        .await;

        let (tcp_stream, peer_addr) = match accept_result {
            Ok(Ok((stream, addr))) => (stream, addr),
            Ok(Err(e)) => {
                error!(error = %e, "Accept error");
                continue;
            }
            Err(_) => continue, // timeout, check shutdown
        };

        // Check max connections
        let current = active_connections.load(Ordering::SeqCst);
        if current >= config.max_connections {
            warn!(
                current = current,
                max = config.max_connections,
                "Max connections reached, rejecting"
            );
            drop(tcp_stream);
            continue;
        }

        let remote_addr = peer_addr.ip().to_string();
        let config = config.clone();
        let rate_limiter = rate_limiter.clone();
        let event_tx = event_tx.clone();
        let callback_registry = callback_registry.clone();
        let tls_acceptor = tls_acceptor.clone();
        let active_connections = active_connections.clone();

        active_connections.fetch_add(1, Ordering::SeqCst);

        tokio::spawn(async move {
            let stream = if implicit_tls {
                // Implicit TLS: wrap immediately
                if let Some(acceptor) = &tls_acceptor {
                    match acceptor.accept(tcp_stream).await {
                        Ok(tls_stream) => {
                            SmtpStream::Tls(TokioBufReader::new(tls_stream))
                        }
                        Err(e) => {
                            warn!(
                                remote_addr = %remote_addr,
                                error = %e,
                                "Implicit TLS handshake failed"
                            );
                            active_connections.fetch_sub(1, Ordering::SeqCst);
                            return;
                        }
                    }
                } else {
                    active_connections.fetch_sub(1, Ordering::SeqCst);
                    return;
                }
            } else {
                SmtpStream::Plain(TokioBufReader::new(tcp_stream))
            };

            connection::handle_connection(
                stream,
                config,
                rate_limiter,
                event_tx,
                callback_registry,
                tls_acceptor,
                remote_addr,
                implicit_tls,
            )
            .await;

            active_connections.fetch_sub(1, Ordering::SeqCst);
        });
    }
}

/// Build a TLS acceptor from PEM cert/key strings.
fn build_tls_acceptor(
    config: &SmtpServerConfig,
) -> Result<tokio_rustls::TlsAcceptor, Box<dyn std::error::Error + Send + Sync>> {
    let cert_pem = config
        .tls_cert_pem
        .as_ref()
        .ok_or("TLS cert not configured")?;
    let key_pem = config
        .tls_key_pem
        .as_ref()
        .ok_or("TLS key not configured")?;

    // Parse certificates
    let certs: Vec<CertificateDer<'static>> = {
        let mut reader = BufReader::new(cert_pem.as_bytes());
        rustls_pemfile::certs(&mut reader)
            .collect::<Result<Vec<_>, _>>()?
    };

    if certs.is_empty() {
        return Err("No certificates found in PEM".into());
    }

    // Parse private key
    let key: PrivateKeyDer<'static> = {
        let mut reader = BufReader::new(key_pem.as_bytes());
        // Try PKCS8 first, then RSA, then EC
        let mut keys = Vec::new();
        for item in rustls_pemfile::read_all(&mut reader) {
            match item? {
                rustls_pemfile::Item::Pkcs8Key(key) => {
                    keys.push(PrivateKeyDer::Pkcs8(key));
                }
                rustls_pemfile::Item::Pkcs1Key(key) => {
                    keys.push(PrivateKeyDer::Pkcs1(key));
                }
                rustls_pemfile::Item::Sec1Key(key) => {
                    keys.push(PrivateKeyDer::Sec1(key));
                }
                _ => {}
            }
        }
        keys.into_iter()
            .next()
            .ok_or("No private key found in PEM")?
    };

    let tls_config = rustls::ServerConfig::builder()
        .with_no_client_auth()
        .with_single_cert(certs, key)?;

    Ok(tokio_rustls::TlsAcceptor::from(Arc::new(tls_config)))
}

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

    #[test]
    fn test_server_config_defaults() {
        let config = SmtpServerConfig::default();
        assert!(!config.has_tls());
        assert_eq!(config.ports, vec![25]);
    }
}