Allow RelaySelector to take cache_dir by argument

2 files changed, 22 insertions, 16 deletions

diff --git a/mullvad-daemon/src/main.rs b/mullvad-daemon/src/main.rs
index 42e595b57a..87e88498a3 100644
--- a/mullvad-daemon/src/main.rs
+++ b/mullvad-daemon/src/main.rs
@@ -237,7 +237,8 @@ impl Daemon {
         let rpc_handle = rpc_handle.chain_err(|| "Unable to create RPC client")?;
         let http_handle = http_handle.chain_err(|| "Unable to create HTTP client")?;
 
-        let relay_selector = Self::create_relay_selector(rpc_handle.clone(), &resource_dir);
+        let relay_selector =
+            Self::create_relay_selector(rpc_handle.clone(), &resource_dir, &cache_dir);
 
         let (tx, rx) = mpsc::channel();
         let management_interface_broadcaster = Self::start_management_interface(tx.clone())?;
@@ -273,8 +274,9 @@ impl Daemon {
     fn create_relay_selector(
         rpc_handle: mullvad_rpc::HttpHandle,
         resource_dir: &Path,
+        cache_dir: &Path,
     ) -> relays::RelaySelector {
-        let mut relay_selector = relays::RelaySelector::new(rpc_handle, &resource_dir);
+        let mut relay_selector = relays::RelaySelector::new(rpc_handle, &resource_dir, cache_dir);
         if let Ok(elapsed) = relay_selector.get_last_updated().elapsed() {
             if elapsed > *MAX_RELAY_CACHE_AGE {
                 if let Err(e) = relay_selector.update(*RELAY_CACHE_UPDATE_TIMEOUT) {
diff --git a/mullvad-daemon/src/relays.rs b/mullvad-daemon/src/relays.rs
index 358fab2251..ef0455ff52 100644
--- a/mullvad-daemon/src/relays.rs
+++ b/mullvad-daemon/src/relays.rs
@@ -23,6 +23,7 @@ use rand::distributions::{IndependentSample, Range};
 use rand::{self, Rng, ThreadRng};
 use tokio_timer::{TimeoutError, Timer};
 
+const RELAYS_FILENAME: &str = "relays.json";
 
 error_chain! {
     errors {
@@ -46,13 +47,15 @@ pub struct RelaySelector {
     last_updated: SystemTime,
     rng: ThreadRng,
     rpc_client: RelayListProxy<HttpHandle>,
+    cache_path: PathBuf,
 }
 
 impl RelaySelector {
     /// Returns a new `RelaySelector` backed by relays cached on disk. Use the `update` method
     /// to refresh the relay list from the internet.
-    pub fn new(rpc_handle: HttpHandle, resource_dir: &Path) -> Self {
-        let (last_updated, relay_list) = match Self::read_cached_relays(resource_dir) {
+    pub fn new(rpc_handle: HttpHandle, resource_dir: &Path, cache_dir: &Path) -> Self {
+        let cache_path = cache_dir.join(RELAYS_FILENAME);
+        let (last_updated, relay_list) = match Self::read_cached_relays(&cache_path, resource_dir) {
             Ok(value) => value,
             Err(error) => {
                 let error = error.chain_err(|| "Unable to load cached relays");
@@ -72,6 +75,7 @@ impl RelaySelector {
             last_updated,
             rng: rand::thread_rng(),
             rpc_client: RelayListProxy::new(rpc_handle),
+            cache_path,
         }
     }
 
@@ -257,7 +261,7 @@ impl RelaySelector {
             .relay_list()
             .map_err(|e| Error::with_chain(e, ErrorKind::DownloadError));
         let relay_list = Timer::default().timeout(download_future, timeout).wait()?;
-        if let Err(e) = Self::cache_relays(&relay_list) {
+        if let Err(e) = self.cache_relays(&relay_list) {
             error!("Unable to save relays to cache: {}", e.display_chain());
         }
         let (locations, relays) = Self::process_relay_list(relay_list);
@@ -297,16 +301,21 @@ impl RelaySelector {
     }
 
     /// Write a `RelayList` to the cache file.
-    fn cache_relays(relays: &RelayList) -> Result<()> {
-        let file = File::create(Self::get_cache_path()?).chain_err(|| ErrorKind::RelayCacheError)?;
-        serde_json::to_writer_pretty(file, relays).chain_err(|| ErrorKind::SerializationError)
+    fn cache_relays(&self, relays: &RelayList) -> Result<()> {
+        debug!("Writing relays cache to {}", self.cache_path.display());
+        let file = File::create(&self.cache_path).chain_err(|| ErrorKind::RelayCacheError)?;
+        serde_json::to_writer_pretty(io::BufWriter::new(file), relays)
+            .chain_err(|| ErrorKind::SerializationError)
     }
 
     /// Try to read the relays, first from cache and if that fails from the `resource_dir`.
-    fn read_cached_relays(resource_dir: &Path) -> Result<(SystemTime, RelayList)> {
-        match Self::get_cache_path().and_then(Self::read_relays) {
+    fn read_cached_relays(
+        cache_path: &Path,
+        resource_dir: &Path,
+    ) -> Result<(SystemTime, RelayList)> {
+        match Self::read_relays(cache_path) {
             Ok(value) => Ok(value),
-            Err(read_cache_error) => match Self::read_relays(resource_dir.join("relays.json")) {
+            Err(read_cache_error) => match Self::read_relays(resource_dir.join(RELAYS_FILENAME)) {
                 Ok(value) => Ok(value),
                 Err(read_resource_error) => Err(read_cache_error.chain_err(|| read_resource_error)),
             },
@@ -332,9 +341,4 @@ impl RelaySelector {
         let last_modified = file.metadata()?.modified()?;
         Ok((last_modified, file))
     }
-
-    fn get_cache_path() -> Result<PathBuf> {
-        let dir = ::cache::get_cache_dir().chain_err(|| ErrorKind::RelayCacheError)?;
-        Ok(dir.join("relays.json"))
-    }
 }