RustBox/src/imap/connection.rs

use super::error::ImapError;
use super::types::ImapCredentials;

// Counts connection attempts by credential identity for isolated regression tests.
#[cfg(test)]
static CONNECTION_COUNTS: std::sync::LazyLock<
    std::sync::Mutex<std::collections::HashMap<[u8; 32], usize>>,
> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));

#[cfg(test)]
pub(super) fn connection_count(creds: &ImapCredentials) -> usize {
    *CONNECTION_COUNTS
        .lock()
        .unwrap()
        .get(&super::pool::identity(creds))
        .unwrap_or(&0)
}

// ---- Connection helper ----------------------------------------------------

/// Establish an authenticated IMAP session.
///
/// Returns a `Session` over a TLS stream (when `creds.tls` is true) or a
/// plain TCP stream.  Because the two stream types are different concrete
/// types we use an enum wrapper that implements the traits `async-imap`
/// requires (`tokio::io::AsyncRead + AsyncWrite + Unpin + Debug`).
pub(crate) async fn connect(
    creds: &ImapCredentials,
) -> Result<async_imap::Session<ImapStream>, ImapError> {
    #[cfg(test)]
    {
        *CONNECTION_COUNTS
            .lock()
            .unwrap()
            .entry(super::pool::identity(creds))
            .or_default() += 1;
    }
    let connect_future = tokio::net::TcpStream::connect((creds.host.as_str(), creds.port));
    // 10 second timeout for the initial TCP connection
    let tcp = tokio::time::timeout(std::time::Duration::from_secs(10), connect_future)
        .await
        .map_err(|_| ImapError::ConnectionFailed("connection timed out".to_string()))?
        .map_err(|e| ImapError::ConnectionFailed(e.to_string()))?;

    if creds.tls {
        let tls = async_native_tls::TlsConnector::new();
        let tls_stream = tls
            .connect(&creds.host, tcp)
            .await
            .map_err(|e| ImapError::ConnectionFailed(e.to_string()))?;
        let client = async_imap::Client::new(ImapStream::Tls(tls_stream));
        let mut session = client
            .login(&creds.email, &creds.password)
            .await
            .map_err(|(e, _)| classify_login_error(e))?;
        maybe_enable_utf8_accept(&mut session).await?;
        Ok(session)
    } else {
        let client = async_imap::Client::new(ImapStream::Plain(tcp));
        let mut session = client
            .login(&creds.email, &creds.password)
            .await
            .map_err(|(e, _)| classify_login_error(e))?;
        maybe_enable_utf8_accept(&mut session).await?;
        Ok(session)
    }
}

/// Enable RFC 6855 UTF8=ACCEPT once per freshly authenticated session when
/// the server explicitly advertises both ENABLE and UTF8=ACCEPT. A tagged
/// NO/BAD is treated as a capability mismatch and leaves the session usable
/// in legacy mode; transport loss still invalidates the connection.
async fn maybe_enable_utf8_accept(
    session: &mut async_imap::Session<ImapStream>,
) -> Result<(), ImapError> {
    use async_imap::imap_proto::{Response, Status};
    let capabilities = session.capabilities().await.map_err(map_imap_error)?;
    let has_atom = |needle: &str| {
        capabilities.iter().any(|capability| match capability {
            async_imap::types::Capability::Atom(name) => name.eq_ignore_ascii_case(needle),
            _ => false,
        })
    };
    if !has_atom("ENABLE") || !has_atom("UTF8=ACCEPT") {
        return Ok(());
    }

    let id = session
        .run_command("ENABLE UTF8=ACCEPT")
        .await
        .map_err(map_imap_error)?;
    while let Some(response) = session.read_response().await {
        let response = response.map_err(|error| ImapError::ConnectionFailed(error.to_string()))?;
        match response.parsed() {
            Response::Done { tag, status, .. } if *tag == id => {
                if *status != Status::Ok {
                    tracing::warn!(
                        ?status,
                        "IMAP server advertised UTF8=ACCEPT but rejected ENABLE; continuing without it"
                    );
                }
                return Ok(());
            }
            Response::Data {
                status: Status::Bye,
                ..
            } => {
                return Err(ImapError::ConnectionFailed(
                    "IMAP server closed the connection while enabling UTF8=ACCEPT".to_string(),
                ));
            }
            _ => {}
        }
    }
    Err(ImapError::ConnectionFailed(
        "connection lost while enabling UTF8=ACCEPT".to_string(),
    ))
}

/// Classify an `async_imap::error::Error` that occurred during LOGIN.
pub(crate) fn classify_login_error(err: async_imap::error::Error) -> ImapError {
    match err {
        async_imap::error::Error::No(_) => ImapError::AuthenticationFailed,
        async_imap::error::Error::Io(e) => ImapError::ConnectionFailed(e.to_string()),
        async_imap::error::Error::ConnectionLost => {
            ImapError::ConnectionFailed("connection lost".to_string())
        }
        other => ImapError::ProtocolError(other.to_string()),
    }
}

/// Map a generic `async_imap` error to our `ImapError`.
pub(crate) fn map_imap_error(err: async_imap::error::Error) -> ImapError {
    match err {
        async_imap::error::Error::No(msg) => ImapError::ProtocolError(format!("NO: {msg}")),
        async_imap::error::Error::Io(e) => ImapError::ConnectionFailed(e.to_string()),
        async_imap::error::Error::ConnectionLost => {
            ImapError::ConnectionFailed("connection lost".to_string())
        }
        other => ImapError::ProtocolError(other.to_string()),
    }
}

// ---- Stream enum ----------------------------------------------------------

/// A wrapper enum so that `Session` can be generic over a single type
/// regardless of whether TLS is used.
#[derive(Debug)]
pub(crate) enum ImapStream {
    Tls(async_native_tls::TlsStream<tokio::net::TcpStream>),
    Plain(tokio::net::TcpStream),
}

impl tokio::io::AsyncRead for ImapStream {
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            ImapStream::Tls(s) => std::pin::Pin::new(s).poll_read(cx, buf),
            ImapStream::Plain(s) => std::pin::Pin::new(s).poll_read(cx, buf),
        }
    }
}

impl tokio::io::AsyncWrite for ImapStream {
    fn poll_write(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        match self.get_mut() {
            ImapStream::Tls(s) => std::pin::Pin::new(s).poll_write(cx, buf),
            ImapStream::Plain(s) => std::pin::Pin::new(s).poll_write(cx, buf),
        }
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            ImapStream::Tls(s) => std::pin::Pin::new(s).poll_flush(cx),
            ImapStream::Plain(s) => std::pin::Pin::new(s).poll_flush(cx),
        }
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            ImapStream::Tls(s) => std::pin::Pin::new(s).poll_shutdown(cx),
            ImapStream::Plain(s) => std::pin::Pin::new(s).poll_shutdown(cx),
        }
    }
}