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::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, 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, ) -> 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), 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> { 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> { 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> { 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> { 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), } } }