Skip to content

Commit

Permalink
feat: parallelize connection attempts
Browse files Browse the repository at this point in the history
  • Loading branch information
link2xt committed Aug 25, 2024
1 parent 94caf74 commit cdaefb4
Showing 1 changed file with 53 additions and 26 deletions.
79 changes: 53 additions & 26 deletions src/imap/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,40 @@ impl Client {
Ok(Session::new(session, capabilities))
}

async fn connection_attempt(
context: &Context,
host: &str,
security: ConnectionSecurity,
resolved_addr: SocketAddr,
strict_tls: bool,
load_cache: bool,
) -> Result<Self> {
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();
if load_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 {resolved_addr}: {err:#}.");
Err(err)
}
}
}

pub async fn connect(
context: &Context,
socks5_config: Option<Socks5Config>,
Expand All @@ -131,40 +165,33 @@ impl Client {
};
Ok(client)
} else {
let mut first_error = None;
let load_cache = match security {
ConnectionSecurity::Tls | ConnectionSecurity::Starttls => strict_tls,
ConnectionSecurity::Plain => false,
};

let mut connection_attempts = Vec::new();
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);
}
}
let fut = Box::pin(Self::connection_attempt(
context,
host,
security,
resolved_addr,
strict_tls,
load_cache,
));
connection_attempts.push(fut);
}

if connection_attempts.is_empty() {
return Err(format_err!("no DNS resolution results for {host}"));
}
Err(first_error.unwrap_or_else(|| format_err!("no DNS resolution results for {host}")))

connection_attempts.reverse();
let (res, _remaining_futures) = futures::future::select_ok(connection_attempts).await?;
Ok(res)
}
}

Expand Down

0 comments on commit cdaefb4

Please sign in to comment.