Shuffle API address cache when loaded, and use bundled API address cache as fallback

3 files changed, 103 insertions, 52 deletions

diff --git a/mullvad-rpc/src/address_cache.rs b/mullvad-rpc/src/address_cache.rs
index 7c1859173e..ae0b916ad4 100644
--- a/mullvad-rpc/src/address_cache.rs
+++ b/mullvad-rpc/src/address_cache.rs
@@ -1,6 +1,8 @@
+use super::API_ADDRESS;
+use rand::seq::SliceRandom;
 use std::{
     io,
-    net::{IpAddr, SocketAddr},
+    net::SocketAddr,
     path::Path,
     sync::{Arc, Mutex},
 };
@@ -9,7 +11,18 @@ use tokio::{
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
 };
 
-const FALLBACK_API_ADDRESS: (IpAddr, u16) = (crate::API_IP, 443);
+#[derive(err_derive::Error, Debug)]
+#[error(no_from)]
+pub enum Error {
+    #[error(display = "Failed to open the address cache file")]
+    OpenAddressCache(#[error(source)] io::Error),
+
+    #[error(display = "Failed to read the address cache file")]
+    ReadAddressCache(#[error(source)] io::Error),
+
+    #[error(display = "The address cache is empty")]
+    EmptyAddressCache,
+}
 
 #[derive(Clone)]
 pub struct AddressCache {
@@ -18,20 +31,20 @@ pub struct AddressCache {
 }
 
 impl AddressCache {
-    pub fn new() -> Self {
-        Self {
-            inner: Arc::new(Mutex::new(Default::default())),
-            cache_path: None,
-        }
+    /// Initialize cache using the given list, and write changes to `cache_path`.
+    pub fn new(addresses: Vec<SocketAddr>, cache_path: Option<Box<Path>>) -> Result<Self, Error> {
+        log::trace!("Using API addresses: {:?}", addresses);
+        let cache = AddressCacheInner::from_addresses(addresses)?;
+        Ok(Self {
+            inner: Arc::new(Mutex::new(cache)),
+            cache_path: cache_path.map(|cache| Arc::from(cache)),
+        })
     }
 
-    pub async fn with_cache(cache_path: Box<Path>) -> Self {
-        let cache = AddressCacheInner::from_cache_file(&cache_path)
-            .await
-            .unwrap_or_default();
-        let inner = Arc::new(Mutex::new(cache));
-        let cache_path = Some(cache_path.into());
-        Self { inner, cache_path }
+    /// Initialize cache using `read_path`, and write changes to `cache_path`.
+    pub async fn from_file(read_path: &Path, cache_path: Option<Box<Path>>) -> Result<Self, Error> {
+        log::trace!("Loading API addresses from {:?}", read_path);
+        Self::new(read_address_file(read_path).await?, cache_path)
     }
 
     pub fn get_address(&self) -> SocketAddr {
@@ -43,12 +56,12 @@ impl AddressCache {
 
     fn get_address_inner(inner: &AddressCacheInner) -> SocketAddr {
         if inner.addresses.is_empty() {
-            return FALLBACK_API_ADDRESS.into();
+            return API_ADDRESS.into();
         }
         *inner
             .addresses
             .get(inner.choice % inner.addresses.len())
-            .unwrap_or(&FALLBACK_API_ADDRESS.into())
+            .unwrap_or(&API_ADDRESS.into())
     }
 
     pub fn register_failure(&self, failed_addr: SocketAddr, err: &dyn std::error::Error) {
@@ -67,11 +80,15 @@ impl AddressCache {
         }
     }
 
-    pub async fn set_addresses(&self, addresses: Vec<SocketAddr>) -> io::Result<()> {
+    pub async fn set_addresses(&self, mut addresses: Vec<SocketAddr>) -> io::Result<()> {
         let should_update = {
             let mut inner = self.inner.lock().unwrap();
-            if addresses != inner.addresses {
+            addresses.sort();
+            let mut current_sorted = inner.addresses.clone();
+            current_sorted.sort();
+            if addresses != current_sorted {
                 inner.addresses = addresses.clone();
+                inner.shuffle();
                 inner.choice = 0;
                 true
             } else {
@@ -120,37 +137,43 @@ struct AddressCacheInner {
 }
 
 impl AddressCacheInner {
-    async fn from_cache_file(path: &Path) -> io::Result<Self> {
-        let file = fs::File::open(path).await?;
-        let mut lines = BufReader::new(file).lines();
-        let mut addresses = vec![];
-        while let Some(line) = lines.next_line().await? {
-            // for line in lines.next_line() {
-            match line.trim().parse() {
-                Ok(address) => addresses.push(address),
-                Err(err) => {
-                    log::error!("Failed to parse cached address line: {}", err);
-                }
-            }
+    fn from_addresses(addresses: Vec<SocketAddr>) -> Result<Self, Error> {
+        if addresses.is_empty() {
+            return Err(Error::EmptyAddressCache);
         }
-
-        if !addresses.contains(&FALLBACK_API_ADDRESS.into()) {
-            addresses.push(FALLBACK_API_ADDRESS.into());
-        }
-
-        Ok(Self {
+        let mut cache = Self {
             addresses,
-            ..Default::default()
-        })
+            choice: 0,
+            last_try: None,
+        };
+        cache.shuffle();
+        Ok(cache)
+    }
+
+    fn shuffle(&mut self) {
+        let mut rng = rand::thread_rng();
+        (&mut self.addresses[..]).shuffle(&mut rng);
     }
 }
 
-impl Default for AddressCacheInner {
-    fn default() -> Self {
-        Self {
-            addresses: vec![FALLBACK_API_ADDRESS.into()],
-            choice: 0,
-            last_try: None,
+async fn read_address_file(path: &Path) -> Result<Vec<SocketAddr>, Error> {
+    let file = fs::File::open(path)
+        .await
+        .map_err(|error| Error::OpenAddressCache(error))?;
+    let mut lines = BufReader::new(file).lines();
+    let mut addresses = vec![];
+    while let Some(line) = lines
+        .next_line()
+        .await
+        .map_err(|error| Error::ReadAddressCache(error))?
+    {
+        // for line in lines.next_line() {
+        match line.trim().parse() {
+            Ok(address) => addresses.push(address),
+            Err(err) => {
+                log::error!("Failed to parse cached address line: {}", err);
+            }
         }
     }
+    Ok(addresses)
 }
diff --git a/mullvad-rpc/src/bin/address_cache.rs b/mullvad-rpc/src/bin/address_cache.rs
index 203c5e79bb..efb2501ba9 100644
--- a/mullvad-rpc/src/bin/address_cache.rs
+++ b/mullvad-rpc/src/bin/address_cache.rs
@@ -1,4 +1,4 @@
-/// Generate a first list of IP addresses for the Mullvad VPN to use to talk to the API.
+/// Generate a first list of IP addresses for Mullvad VPN to use to talk to the API.
 use mullvad_rpc::{rest::Error as RestError, ApiProxy, MullvadRpcRuntime};
 use std::process;
 use talpid_types::ErrorExt;
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index 323937ea9a..d92be6e8be 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -12,7 +12,7 @@ use std::{
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::Path,
 };
-use talpid_types::net::wireguard;
+use talpid_types::{net::wireguard, ErrorExt};
 
 
 pub mod rest;
@@ -35,6 +35,7 @@ pub const INVALID_VOUCHER: &str = "INVALID_VOUCHER";
 const API_HOST: &str = "api.mullvad.net";
 pub const API_IP_CACHE_FILENAME: &str = "api-ip-address.txt";
 const API_IP: IpAddr = IpAddr::V4(Ipv4Addr::new(193, 138, 218, 78));
+const API_ADDRESS: (IpAddr, u16) = (crate::API_IP, 443);
 
 
 /// A type that helps with the creation of RPC connections.
@@ -48,6 +49,9 @@ pub struct MullvadRpcRuntime {
 pub enum Error {
     #[error(display = "Failed to construct a rest client")]
     RestError(#[error(source)] rest::Error),
+
+    #[error(display = "Failed to load address cache")]
+    AddressCacheError(#[error(source)] address_cache::Error),
 }
 
 impl MullvadRpcRuntime {
@@ -56,17 +60,41 @@ impl MullvadRpcRuntime {
         Ok(MullvadRpcRuntime {
             https_connector: HttpsConnectorWithSni::new(),
             handle,
-            address_cache: AddressCache::new(),
+            address_cache: AddressCache::new(vec![API_ADDRESS.into()], None)?,
         })
     }
 
-    /// Create a new `MullvadRpcRuntime` using the specified cache directory.
-    pub async fn with_cache_dir(
+    /// Create a new `MullvadRpcRuntime` using the specified directories.
+    /// Try to use the cache directory first, and fall back on the resource directory
+    /// if it fails.
+    pub async fn with_cache(
         handle: tokio::runtime::Handle,
-        cache_dir: &Path,
+        resource_dir: &Path,
+        cache_dir: Option<&Path>,
     ) -> Result<Self, Error> {
-        let cache_file = cache_dir.join(API_IP_CACHE_FILENAME);
-        let address_cache = AddressCache::with_cache(cache_file.into_boxed_path()).await;
+        let resource_file = resource_dir.join(API_IP_CACHE_FILENAME);
+
+        let address_cache = if let Some(cache_dir) = cache_dir {
+            let cache_file = cache_dir.join(API_IP_CACHE_FILENAME);
+            let cache_file_boxed = cache_file.clone().into_boxed_path();
+
+            match AddressCache::from_file(&cache_file, Some(cache_file_boxed.clone())).await {
+                Ok(cache) => cache,
+                Err(error) => {
+                    if cache_file.exists() {
+                        log::error!(
+                            "{}",
+                            error.display_chain_with_msg(
+                                "Failed to load cached API addresses. Falling back on bundled list"
+                            )
+                        );
+                    }
+                    AddressCache::from_file(&resource_file, Some(cache_file_boxed)).await?
+                }
+            }
+        } else {
+            AddressCache::from_file(&resource_file, None).await?
+        };
 
         let https_connector = HttpsConnectorWithSni::new();