Skip to content

Commit

Permalink
feat: parallelize IMAP and SMTP connection attempts (#5915)
Browse files Browse the repository at this point in the history 
Previously for each connection candidate (essentially host and port
pair) after resolving the host to a list of IPs Delta Chat iterated IP
addresses one by one. Now for IMAP and SMTP we try up to 5 IP addresses
in parallel. We start with one connection and add more connections
later. If some connection fails, e.g. we try to connect to IPv6 on IPv4
network and get "Network is unreachable" (ENETUNREACH) error, we replace
failed connection with another one immediately.

Co-authored-by: Hocuri <[email protected]>
  • Loading branch information
@link2xt @Hocuri
link2xt and Hocuri authored Aug 28, 2024
1 parent 84c1ffd commit aab8ef2
Show file tree
Hide file tree
Showing 4 changed files with 204 additions and 80 deletions.
2 changes: 1 addition & 1 deletion src/configure.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -576,7 +576,7 @@ async fn get_autoconfig(

async fn nicer_configuration_error(context: &Context, e: String) -> String {
if e.to_lowercase().contains("could not resolve")
|| e.to_lowercase().contains("no dns resolution results")
|| e.to_lowercase().contains("connection attempts")
|| e.to_lowercase()
.contains("temporary failure in name resolution")
|| e.to_lowercase().contains("name or service not known")
Expand Down
94 changes: 62 additions & 32 deletions src/imap/client.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::net::SocketAddr;
use std::ops::{Deref, DerefMut};

use anyhow::{format_err, Context as _, Result};
use anyhow::{Context as _, Result};
use async_imap::Client as ImapClient;
use async_imap::Session as ImapSession;
use fast_socks5::client::Socks5Stream;
Expand All @@ -14,7 +14,9 @@ use crate::login_param::{ConnectionCandidate, ConnectionSecurity};
use crate::net::dns::{lookup_host_with_cache, update_connect_timestamp};
use crate::net::session::SessionStream;
use crate::net::tls::wrap_tls;
use crate::net::{connect_tcp_inner, connect_tls_inner, update_connection_history};
use crate::net::{
connect_tcp_inner, connect_tls_inner, run_connection_attempts, update_connection_history,
};
use crate::socks::Socks5Config;
use crate::tools::time;

Expand Down Expand Up @@ -106,6 +108,53 @@ impl Client {
Ok(Session::new(session, capabilities))
}

async fn connection_attempt(
context: Context,
host: String,
security: ConnectionSecurity,
resolved_addr: SocketAddr,
strict_tls: bool,
) -> Result<Self> {
let context = &context;
let host = &host;
info!(
context,
"Attempting IMAP connection to {host} ({resolved_addr})."
);
let res = match security {
ConnectionSecurity::Tls => {
Client::connect_secure(resolved_addr, host, strict_tls).await
}
ConnectionSecurity::Starttls => {
Client::connect_starttls(resolved_addr, host, strict_tls).await
}
ConnectionSecurity::Plain => Client::connect_insecure(resolved_addr).await,
};
match res {
Ok(client) => {
let ip_addr = resolved_addr.ip().to_string();
let port = resolved_addr.port();

let save_cache = match security {
ConnectionSecurity::Tls | ConnectionSecurity::Starttls => strict_tls,
ConnectionSecurity::Plain => false,
};
if save_cache {
update_connect_timestamp(context, host, &ip_addr).await?;
}
update_connection_history(context, "imap", host, port, &ip_addr, time()).await?;
Ok(client)
}
Err(err) => {
warn!(
context,
"Failed to connect to {host} ({resolved_addr}): {err:#}."
);
Err(err)
}
}
}

pub async fn connect(
context: &Context,
socks5_config: Option<Socks5Config>,
Expand All @@ -131,40 +180,21 @@ impl Client {
};
Ok(client)
} else {
let mut first_error = None;
let load_cache = match security {
ConnectionSecurity::Tls | ConnectionSecurity::Starttls => strict_tls,
ConnectionSecurity::Plain => false,
};
for resolved_addr in
lookup_host_with_cache(context, host, port, "imap", load_cache).await?
{
let res = match security {
ConnectionSecurity::Tls => {
Client::connect_secure(resolved_addr, host, strict_tls).await
}
ConnectionSecurity::Starttls => {
Client::connect_starttls(resolved_addr, host, strict_tls).await
}
ConnectionSecurity::Plain => Client::connect_insecure(resolved_addr).await,
};
match res {
Ok(client) => {
let ip_addr = resolved_addr.ip().to_string();
if load_cache {
update_connect_timestamp(context, host, &ip_addr).await?;
}
update_connection_history(context, "imap", host, port, &ip_addr, time())
.await?;
return Ok(client);
}
Err(err) => {
warn!(context, "Failed to connect to {resolved_addr}: {err:#}.");
first_error.get_or_insert(err);
}
}
}
Err(first_error.unwrap_or_else(|| format_err!("no DNS resolution results for {host}")))

let connection_futures =
lookup_host_with_cache(context, host, port, "imap", load_cache)
.await?
.into_iter()
.map(|resolved_addr| {
let context = context.clone();
let host = host.to_string();
Self::connection_attempt(context, host, security, resolved_addr, strict_tls)
});
run_connection_attempts(connection_futures).await
}
}

Expand Down
103 changes: 85 additions & 18 deletions src/net.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
//! # Common network utilities.
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::time::Duration;

use anyhow::{format_err, Context as _, Result};
use async_native_tls::TlsStream;
use tokio::net::TcpStream;
use tokio::task::JoinSet;
use tokio::time::timeout;
use tokio_io_timeout::TimeoutStream;

Expand Down Expand Up @@ -130,6 +132,84 @@ pub(crate) async fn connect_tls_inner(
Ok(tls_stream)
}

/// Runs connection attempt futures.
///
/// Accepts iterator of connection attempt futures
/// and runs them until one of them succeeds
/// or all of them fail.
///
/// If all connection attempts fail, returns the first error.
///
/// This functions starts with one connection attempt and maintains
/// up to five parallel connection attempts if connecting takes time.
pub(crate) async fn run_connection_attempts<O, I, F>(mut futures: I) -> Result<O>
where
I: Iterator<Item = F>,
F: Future<Output = Result<O>> + Send + 'static,
O: Send + 'static,
{
let mut connection_attempt_set = JoinSet::new();

// Start additional connection attempts after 300 ms, 1 s, 5 s and 10 s.
// This way we can have up to 5 parallel connection attempts at the same time.
let mut delays = [
Duration::from_millis(300),
Duration::from_secs(1),
Duration::from_secs(5),
Duration::from_secs(10),
]
.into_iter();

let mut first_error = None;

let res = loop {
if let Some(fut) = futures.next() {
connection_attempt_set.spawn(fut);
}

let one_year = Duration::from_secs(60 * 60 * 24 * 365);
let delay = delays.next().unwrap_or(one_year); // one year can be treated as infinitely long here
let Ok(res) = timeout(delay, connection_attempt_set.join_next()).await else {
// The delay for starting the next connection attempt has expired.
// `continue` the loop to push the next connection into connection_attempt_set.
continue;
};

match res {
Some(res) => {
match res.context("Failed to join task") {
Ok(Ok(conn)) => {
// Successfully connected.
break Ok(conn);
}
Ok(Err(err)) => {
// Some connection attempt failed.
first_error.get_or_insert(err);
}
Err(err) => {
break Err(err);
}
}
}
None => {
// Out of connection attempts.
//
// Break out of the loop and return error.
break Err(
first_error.unwrap_or_else(|| format_err!("No connection attempts were made"))
);
}
}
};

// Abort remaining connection attempts and free resources
// such as OS sockets and `Context` references
// held by connection attempt tasks.
connection_attempt_set.shutdown().await;

res
}

/// If `load_cache` is true, may use cached DNS results.
/// Because the cache may be poisoned with incorrect results by networks hijacking DNS requests,
/// this option should only be used when connection is authenticated,
Expand All @@ -142,22 +222,9 @@ pub(crate) async fn connect_tcp(
port: u16,
load_cache: bool,
) -> Result<Pin<Box<TimeoutStream<TcpStream>>>> {
let mut first_error = None;

for resolved_addr in lookup_host_with_cache(context, host, port, "", load_cache).await? {
match connect_tcp_inner(resolved_addr).await {
Ok(stream) => {
return Ok(stream);
}
Err(err) => {
warn!(
context,
"Failed to connect to {}: {:#}.", resolved_addr, err
);
first_error.get_or_insert(err);
}
}
}

Err(first_error.unwrap_or_else(|| format_err!("no DNS resolution results for {host}")))
let connection_futures = lookup_host_with_cache(context, host, port, "", load_cache)
.await?
.into_iter()
.map(connect_tcp_inner);
run_connection_attempts(connection_futures).await
}
85 changes: 56 additions & 29 deletions src/smtp/connect.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
use std::net::SocketAddr;

use anyhow::{bail, format_err, Context as _, Result};
use anyhow::{bail, Context as _, Result};
use async_smtp::{SmtpClient, SmtpTransport};
use tokio::io::BufStream;

Expand All @@ -11,7 +11,9 @@ use crate::login_param::{ConnectionCandidate, ConnectionSecurity};
use crate::net::dns::{lookup_host_with_cache, update_connect_timestamp};
use crate::net::session::SessionBufStream;
use crate::net::tls::wrap_tls;
use crate::net::{connect_tcp_inner, connect_tls_inner, update_connection_history};
use crate::net::{
connect_tcp_inner, connect_tls_inner, run_connection_attempts, update_connection_history,
};
use crate::oauth2::get_oauth2_access_token;
use crate::socks::Socks5Config;
use crate::tools::time;
Expand Down Expand Up @@ -72,6 +74,49 @@ pub(crate) async fn connect_and_auth(
Ok(transport)
}

async fn connection_attempt(
context: Context,
host: String,
security: ConnectionSecurity,
resolved_addr: SocketAddr,
strict_tls: bool,
) -> Result<Box<dyn SessionBufStream>> {
let context = &context;
let host = &host;
info!(
context,
"Attempting SMTP connection to {host} ({resolved_addr})."
);
let res = match security {
ConnectionSecurity::Tls => connect_secure(resolved_addr, host, strict_tls).await,
ConnectionSecurity::Starttls => connect_starttls(resolved_addr, host, strict_tls).await,
ConnectionSecurity::Plain => connect_insecure(resolved_addr).await,
};
match res {
Ok(stream) => {
let ip_addr = resolved_addr.ip().to_string();
let port = resolved_addr.port();

let save_cache = match security {
ConnectionSecurity::Tls | ConnectionSecurity::Starttls => strict_tls,
ConnectionSecurity::Plain => false,
};
if save_cache {
update_connect_timestamp(context, host, &ip_addr).await?;
}
update_connection_history(context, "smtp", host, port, &ip_addr, time()).await?;
Ok(stream)
}
Err(err) => {
warn!(
context,
"Failed to connect to {host} ({resolved_addr}): {err:#}."
);
Err(err)
}
}
}

/// Returns TLS, STARTTLS or plaintext connection
/// using SOCKS5 or direct connection depending on the given configuration.
///
Expand Down Expand Up @@ -106,38 +151,20 @@ async fn connect_stream(
};
Ok(stream)
} else {
let mut first_error = None;
let load_cache = match security {
ConnectionSecurity::Tls | ConnectionSecurity::Starttls => strict_tls,
ConnectionSecurity::Plain => false,
};

for resolved_addr in lookup_host_with_cache(context, host, port, "smtp", load_cache).await?
{
let res = match security {
ConnectionSecurity::Tls => connect_secure(resolved_addr, host, strict_tls).await,
ConnectionSecurity::Starttls => {
connect_starttls(resolved_addr, host, strict_tls).await
}
ConnectionSecurity::Plain => connect_insecure(resolved_addr).await,
};
match res {
Ok(stream) => {
let ip_addr = resolved_addr.ip().to_string();
if load_cache {
update_connect_timestamp(context, host, &ip_addr).await?;
}
update_connection_history(context, "smtp", host, port, &ip_addr, time())
.await?;
return Ok(stream);
}
Err(err) => {
warn!(context, "Failed to connect to {resolved_addr}: {err:#}.");
first_error.get_or_insert(err);
}
}
}
Err(first_error.unwrap_or_else(|| format_err!("no DNS resolution results for {host}")))
let connection_futures = lookup_host_with_cache(context, host, port, "smtp", load_cache)
.await?
.into_iter()
.map(|resolved_addr| {
let context = context.clone();
let host = host.to_string();
connection_attempt(context, host, security, resolved_addr, strict_tls)
});
run_connection_attempts(connection_futures).await
}
}

Expand Down

0 comments on commit aab8ef2

Please sign in to comment.