Add non-blocking DNS resolver for Android API requests

1 files changed, 50 insertions, 2 deletions

diff --git a/mullvad-api/src/lib.rs b/mullvad-api/src/lib.rs
index 6b3ac3c951..1f47d600b3 100644
--- a/mullvad-api/src/lib.rs
+++ b/mullvad-api/src/lib.rs
@@ -1,4 +1,5 @@
 #![allow(rustdoc::private_intra_doc_links)]
+use async_trait::async_trait;
 #[cfg(target_os = "android")]
 use futures::channel::mpsc;
 #[cfg(target_os = "android")]
@@ -12,10 +13,11 @@ use std::{
     cell::Cell,
     collections::BTreeMap,
     future::Future,
+    io,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     ops::Deref,
     path::Path,
-    sync::OnceLock,
+    sync::{Arc, OnceLock},
 };
 use talpid_types::ErrorExt;
 
@@ -304,11 +306,43 @@ impl ApiEndpoint {
     }
 }
 
+#[async_trait]
+pub trait DnsResolver: 'static + Send + Sync {
+    async fn resolve(&self, host: String) -> io::Result<Vec<IpAddr>>;
+}
+
+/// DNS resolver that relies on `ToSocketAddrs` (`getaddrinfo`).
+pub struct DefaultDnsResolver;
+
+#[async_trait]
+impl DnsResolver for DefaultDnsResolver {
+    async fn resolve(&self, host: String) -> io::Result<Vec<IpAddr>> {
+        use std::net::ToSocketAddrs;
+        // Spawn a blocking thread, since `to_socket_addrs` relies on `libc::getaddrinfo`, which
+        // blocks and either has no timeout or a very long one.
+        let addrs = tokio::task::spawn_blocking(move || (host, 0).to_socket_addrs())
+            .await
+            .expect("DNS task panicked")?;
+        Ok(addrs.map(|addr| addr.ip()).collect())
+    }
+}
+
+/// DNS resolver that always returns no results
+pub struct NullDnsResolver;
+
+#[async_trait]
+impl DnsResolver for NullDnsResolver {
+    async fn resolve(&self, _host: String) -> io::Result<Vec<IpAddr>> {
+        Ok(vec![])
+    }
+}
+
 /// A type that helps with the creation of API connections.
 pub struct Runtime {
     handle: tokio::runtime::Handle,
     address_cache: AddressCache,
     api_availability: availability::ApiAvailability,
+    dns_resolver: Arc<dyn DnsResolver>,
     #[cfg(target_os = "android")]
     socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
 }
@@ -323,13 +357,20 @@ pub enum Error {
 
     #[error("API availability check failed")]
     ApiCheckError(#[from] availability::Error),
+
+    #[error("DNS resolution error")]
+    ResolutionFailed(#[from] std::io::Error),
 }
 
 impl Runtime {
     /// Create a new `Runtime`.
-    pub fn new(handle: tokio::runtime::Handle) -> Result<Self, Error> {
+    pub fn new(
+        handle: tokio::runtime::Handle,
+        dns_resolver: impl DnsResolver,
+    ) -> Result<Self, Error> {
         Self::new_inner(
             handle,
+            dns_resolver,
             #[cfg(target_os = "android")]
             None,
         )
@@ -346,12 +387,14 @@ impl Runtime {
 
     fn new_inner(
         handle: tokio::runtime::Handle,
+        dns_resolver: impl DnsResolver,
         #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
     ) -> Result<Self, Error> {
         Ok(Runtime {
             handle,
             address_cache: AddressCache::new(None)?,
             api_availability: ApiAvailability::default(),
+            dns_resolver: Arc::new(dns_resolver),
             #[cfg(target_os = "android")]
             socket_bypass_tx,
         })
@@ -360,15 +403,18 @@ impl Runtime {
     /// Create a new `Runtime` using the specified directories.
     /// Try to use the cache directory first, and fall back on the bundled address otherwise.
     pub async fn with_cache(
+        dns_resolver: impl DnsResolver,
         cache_dir: &Path,
         write_changes: bool,
         #[cfg(target_os = "android")] socket_bypass_tx: Option<mpsc::Sender<SocketBypassRequest>>,
     ) -> Result<Self, Error> {
         let handle = tokio::runtime::Handle::current();
+
         #[cfg(feature = "api-override")]
         if API.disable_address_cache {
             return Self::new_inner(
                 handle,
+                dns_resolver,
                 #[cfg(target_os = "android")]
                 socket_bypass_tx,
             );
@@ -402,6 +448,7 @@ impl Runtime {
             handle,
             address_cache,
             api_availability,
+            dns_resolver: Arc::new(dns_resolver),
             #[cfg(target_os = "android")]
             socket_bypass_tx,
         })
@@ -419,6 +466,7 @@ impl Runtime {
             self.api_availability.clone(),
             self.address_cache.clone(),
             connection_mode_provider,
+            self.dns_resolver.clone(),
             #[cfg(target_os = "android")]
             socket_bypass_tx,
         )