Require an initial fallback address to be used

2 files changed, 100 insertions, 42 deletions

diff --git a/mullvad-rpc/src/cached_dns_resolver.rs b/mullvad-rpc/src/cached_dns_resolver.rs
index 318c3dc2d8..9e25bff511 100644
--- a/mullvad-rpc/src/cached_dns_resolver.rs
+++ b/mullvad-rpc/src/cached_dns_resolver.rs
@@ -25,30 +25,48 @@ pub struct CachedDnsResolver<R: DnsResolver = SystemDnsResolver> {
     hostname: String,
     dns_resolver: R,
     cache_file: PathBuf,
-    cached_address: Option<IpAddr>,
+    cached_address: IpAddr,
+    should_update: bool,
 }
 
 impl CachedDnsResolver<SystemDnsResolver> {
-    pub fn new(hostname: String, cache_file: PathBuf) -> Self {
-        Self::with_dns_resolver(SystemDnsResolver, hostname, cache_file)
+    pub fn new(hostname: String, cache_file: PathBuf, fallback_address: IpAddr) -> Self {
+        Self::with_dns_resolver(SystemDnsResolver, hostname, cache_file, fallback_address)
     }
 }
 
 impl<R: DnsResolver> CachedDnsResolver<R> {
-    pub fn with_dns_resolver(dns_resolver: R, hostname: String, cache_file: PathBuf) -> Self {
-        let cached_address = Self::load_from_file(&cache_file).ok();
+    pub fn with_dns_resolver(
+        dns_resolver: R,
+        hostname: String,
+        cache_file: PathBuf,
+        fallback_address: IpAddr,
+    ) -> Self {
+        let (cached_address, should_update) =
+            Self::load_initial_cached_address(&cache_file, fallback_address);
 
         CachedDnsResolver {
             hostname,
             dns_resolver,
             cache_file,
             cached_address,
+            should_update,
         }
     }
 
-    pub fn resolve(&self) -> Option<IpAddr> {
+    pub fn resolve(&mut self) -> IpAddr {
+        if self.should_update {
+            self.resolve_into_cache();
+        }
+
         self.cached_address
-            .or_else(|| self.resolve_into_cache().ok())
+    }
+
+    fn load_initial_cached_address(cache_file: &Path, fallback_address: IpAddr) -> (IpAddr, bool) {
+        match Self::load_from_file(cache_file) {
+            Ok(previously_cached_address) => (previously_cached_address, false),
+            Err(_) => (fallback_address, true),
+        }
     }
 
     fn load_from_file(file_path: &Path) -> io::Result<IpAddr> {
@@ -63,22 +81,18 @@ impl<R: DnsResolver> CachedDnsResolver<R> {
             .map_err(|_| io::Error::new(io::ErrorKind::Other, "Invalid address data"))
     }
 
-    fn resolve_into_cache(&self) -> io::Result<IpAddr> {
-        let address = self.resolve_address()?;
-
-        let _ = self.store_in_cache(address);
-
-        Ok(address)
-    }
-
-    fn resolve_address(&self) -> io::Result<IpAddr> {
-        self.dns_resolver.resolve(&self.hostname)
+    fn resolve_into_cache(&mut self) {
+        if let Ok(address) = self.dns_resolver.resolve(&self.hostname) {
+            self.cached_address = address;
+            self.should_update = false;
+            self.update_cache_file();
+        }
     }
 
-    fn store_in_cache(&self, address: IpAddr) -> io::Result<()> {
-        let mut cache_file = File::create(&self.cache_file)?;
-
-        writeln!(cache_file, "{}", address)
+    fn update_cache_file(&mut self) {
+        if let Ok(mut cache_file) = File::create(&self.cache_file) {
+            let _ = writeln!(cache_file, "{}", self.cached_address);
+        }
     }
 }
 
@@ -103,8 +117,8 @@ mod tests {
 
         write_address(&cache_dir, cached_address);
 
-        let cache = create_cached_dns_resolver(mock_resolver, &cache_dir);
-        let address = cache.resolve().unwrap();
+        let mut cache = create_cached_dns_resolver(mock_resolver, &cache_dir, None);
+        let address = cache.resolve();
 
         assert!(!mock_resolver_was_called.load(Ordering::Acquire));
         assert_eq!(address, cached_address);
@@ -115,9 +129,9 @@ mod tests {
         let (_temp_dir, cache_dir) = create_test_dirs();
         let mock_address = "192.168.1.206".parse().unwrap();
         let mock_resolver = MockDnsResolver::with_address(mock_address);
-        let cache = create_cached_dns_resolver(mock_resolver, &cache_dir);
 
-        let address = cache.resolve().unwrap();
+        let mut cache = create_cached_dns_resolver(mock_resolver, &cache_dir, None);
+        let address = cache.resolve();
 
         assert_eq!(address, mock_address);
         assert_eq!(get_cached_address(&cache_dir), address.to_string());
@@ -128,11 +142,41 @@ mod tests {
         let (temp_dir, cache_dir) = create_test_dirs();
         let mock_address = "192.168.1.206".parse().unwrap();
         let mock_resolver = MockDnsResolver::with_address(mock_address);
-        let cache = create_cached_dns_resolver(mock_resolver, &cache_dir);
+
+        let mut cache = create_cached_dns_resolver(mock_resolver, &cache_dir, None);
 
         ::std::mem::drop(temp_dir);
 
-        assert_eq!(cache.resolve().unwrap(), mock_address);
+        assert_eq!(cache.resolve(), mock_address);
+    }
+
+    #[test]
+    fn uses_fallback_address() {
+        let (_temp_dir, cache_dir) = create_test_dirs();
+        let fallback_address = "192.168.1.31".parse().unwrap();
+        let mock_resolver = MockDnsResolver::that_fails();
+        let mock_resolver_was_called = mock_resolver.was_called_handle();
+
+        let mut cache =
+            create_cached_dns_resolver(mock_resolver, &cache_dir, Some(fallback_address));
+        let address = cache.resolve();
+
+        assert!(mock_resolver_was_called.load(Ordering::Acquire));
+        assert_eq!(address, fallback_address);
+    }
+
+    #[test]
+    fn ignores_fallback_address_if_resolution_succeeds() {
+        let (_temp_dir, cache_dir) = create_test_dirs();
+        let fallback_address = "192.168.1.31".parse().unwrap();
+        let mock_address = "192.168.1.206".parse().unwrap();
+        let mock_resolver = MockDnsResolver::with_address(mock_address);
+
+        let mut cache =
+            create_cached_dns_resolver(mock_resolver, &cache_dir, Some(fallback_address));
+        let address = cache.resolve();
+
+        assert_eq!(address, mock_address);
     }
 
     fn create_test_dirs() -> (TempDir, PathBuf) {
@@ -169,12 +213,14 @@ mod tests {
     fn create_cached_dns_resolver(
         mock_resolver: MockDnsResolver,
         cache_dir: &Path,
+        fallback_address: Option<IpAddr>,
     ) -> CachedDnsResolver<MockDnsResolver> {
         let hostname = "dummy.host".to_owned();
         let filename = "api_ip_address.txt";
         let cache_file = cache_dir.join(filename);
+        let fallback_address = fallback_address.unwrap_or(IpAddr::from([10, 0, 109, 91]));
 
-        CachedDnsResolver::with_dns_resolver(mock_resolver, hostname, cache_file)
+        CachedDnsResolver::with_dns_resolver(mock_resolver, hostname, cache_file, fallback_address)
     }
 
     struct MockDnsResolver {
@@ -190,6 +236,13 @@ mod tests {
             }
         }
 
+        pub fn that_fails() -> Self {
+            MockDnsResolver {
+                address: None,
+                called: Arc::new(AtomicBool::new(false)),
+            }
+        }
+
         pub fn was_called_handle(&self) -> Arc<AtomicBool> {
             self.called.clone()
         }
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index a059620a6d..6a1a682a1b 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -37,6 +37,7 @@ use mullvad_types::relay_list::RelayList;
 use mullvad_types::version;
 
 use std::collections::HashMap;
+use std::net::IpAddr;
 use std::path::Path;
 
 pub mod event_loop;
@@ -61,12 +62,13 @@ impl MullvadRpcFactory {
         }
     }
 
-    /// Create a new `MullvadRpcFactory` using the specified resource directory.
-    pub fn with_resource_dir(resource_dir: &Path) -> Self {
+    /// Create a new `MullvadRpcFactory` using the specified cache directory.
+    pub fn with_cache_dir(cache_dir: &Path) -> Self {
         let hostname = MASTER_API_HOST.to_owned();
-        let cache_file = resource_dir.join("api_ip_address.txt");
+        let cache_file = cache_dir.join("api_ip_address.txt");
+        let fallback_address = IpAddr::from([193, 138, 219, 46]);
 
-        let cached_dns_resolver = CachedDnsResolver::new(hostname, cache_file);
+        let cached_dns_resolver = CachedDnsResolver::new(hostname, cache_file, fallback_address);
 
         MullvadRpcFactory {
             address_cache: Some(cached_dns_resolver),
@@ -74,16 +76,19 @@ impl MullvadRpcFactory {
     }
 
     /// Spawns a tokio core on a new thread and returns a `HttpHandle` running on that core.
-    pub fn new_connection(&self) -> Result<HttpHandle, HttpError> {
+    pub fn new_connection(&mut self) -> Result<HttpHandle, HttpError> {
         self.setup_connection(HttpTransport::new()?)
     }
 
     /// Create and returns a `HttpHandle` running on the given core handle.
-    pub fn new_connection_on_event_loop(&self, handle: &Handle) -> Result<HttpHandle, HttpError> {
+    pub fn new_connection_on_event_loop(
+        &mut self,
+        handle: &Handle,
+    ) -> Result<HttpHandle, HttpError> {
         self.setup_connection(HttpTransport::shared(handle)?)
     }
 
-    fn setup_connection(&self, transport: HttpTransport) -> Result<HttpHandle, HttpError> {
+    fn setup_connection(&mut self, transport: HttpTransport) -> Result<HttpHandle, HttpError> {
         let mut handle = transport.handle(&self.api_uri())?;
 
         handle.set_header(Host::new(MASTER_API_HOST, None));
@@ -91,12 +96,12 @@ impl MullvadRpcFactory {
         Ok(handle)
     }
 
-    fn api_uri(&self) -> String {
-        let address = self.address_cache
-            .as_ref()
-            .and_then(CachedDnsResolver::resolve)
-            .map(|ip| ip.to_string())
-            .unwrap_or_else(|| MASTER_API_HOST.to_owned());
+    fn api_uri(&mut self) -> String {
+        let address = if let Some(ref mut address_cache) = self.address_cache {
+            address_cache.resolve().to_string()
+        } else {
+            MASTER_API_HOST.to_owned()
+        };
 
         format!("https://{}/rpc/", address)
     }
@@ -117,7 +122,7 @@ jsonrpc_client!(pub struct ProblemReportProxy {
 });
 
 impl ProblemReportProxy<HttpHandle> {
-    pub fn connect(manager: &MullvadRpcFactory) -> Result<Self, HttpError> {
+    pub fn connect(manager: &mut MullvadRpcFactory) -> Result<Self, HttpError> {
         Ok(ProblemReportProxy::new(manager.new_connection()?))
     }
 }