Pass relay selector handle directly to "connection mode provider"

2 files changed, 11 insertions, 47 deletions

diff --git a/mullvad-daemon/src/api.rs b/mullvad-daemon/src/api.rs
index 41ac64ea32..563a53fab4 100644
--- a/mullvad-daemon/src/api.rs
+++ b/mullvad-daemon/src/api.rs
@@ -34,9 +34,7 @@ use talpid_types::{
 pub struct ApiConnectionModeProvider {
     cache_dir: PathBuf,
 
-    selector_rx: Option<oneshot::Receiver<RelaySelector>>,
-
-    relay_selector: Option<RelaySelector>,
+    relay_selector: RelaySelector,
     retry_attempt: u32,
 
     current_task: Option<Pin<Box<dyn Future<Output = ApiConnectionMode> + Send>>>,
@@ -49,18 +47,6 @@ impl Stream for ApiConnectionModeProvider {
         mut self: Pin<&mut Self>,
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Option<Self::Item>> {
-        // First, obtain a relay selector handle
-        // Until we have obtained the handle, just return `ApiConnectionMode::Direct`.
-        if let Some(rx) = self.selector_rx.as_mut() {
-            match rx.try_recv() {
-                Ok(Some(selector)) => {
-                    self.selector_rx = None;
-                    self.relay_selector = Some(selector);
-                }
-                Ok(None) | Err(_) => return Poll::Ready(Some(ApiConnectionMode::Direct)),
-            }
-        }
-
         // Poll the current task
         if let Some(task) = self.current_task.as_mut() {
             return match task.as_mut().poll(cx) {
@@ -73,11 +59,8 @@ impl Stream for ApiConnectionModeProvider {
         }
 
         // Create a new task.
-        // Unwrapping is safe since we must have a handle at this point.
-        let selector = self.relay_selector.as_ref().unwrap();
-
         let config = if Self::should_use_bridge(self.retry_attempt) {
-            selector
+            self.relay_selector
                 .get_bridge_forced()
                 .map(|settings| match settings {
                     ProxySettings::Shadowsocks(ss_settings) => {
@@ -111,22 +94,15 @@ impl Stream for ApiConnectionModeProvider {
 }
 
 impl ApiConnectionModeProvider {
-    pub(crate) fn new(cache_dir: PathBuf) -> (Self, ApiConnectionModeProviderHandle) {
-        let (selector_tx, selector_rx) = oneshot::channel();
-
-        (
-            Self {
-                cache_dir,
-
-                selector_rx: Some(selector_rx),
+    pub(crate) fn new(cache_dir: PathBuf, relay_selector: RelaySelector) -> Self {
+        Self {
+            cache_dir,
 
-                relay_selector: None,
-                retry_attempt: 0,
+            relay_selector,
+            retry_attempt: 0,
 
-                current_task: None,
-            },
-            ApiConnectionModeProviderHandle { tx: selector_tx },
-        )
+            current_task: None,
+        }
     }
 
     fn should_use_bridge(retry_attempt: u32) -> bool {
@@ -134,17 +110,6 @@ impl ApiConnectionModeProvider {
     }
 }
 
-/// Used to initialize [ApiConnectionModeProvider]'s relay selector handle.
-pub(crate) struct ApiConnectionModeProviderHandle {
-    tx: oneshot::Sender<RelaySelector>,
-}
-
-impl ApiConnectionModeProviderHandle {
-    pub fn set_relay_selector(self, selector: RelaySelector) {
-        let _ = self.tx.send(selector);
-    }
-}
-
 /// Notifies the tunnel state machine that the API (real or proxied) endpoint has
 /// changed. [ApiEndpointUpdaterHandle::callback()] creates a callback that may
 /// be passed to the `mullvad-api` runtime.
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index b69a7ba2b4..b98c6b5321 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -639,9 +639,8 @@ where
         let initial_selector_config = new_selector_config(&settings);
         let relay_selector = RelaySelector::new(initial_selector_config, &resource_dir, &cache_dir);
 
-        let (proxy_provider, proxy_provider_handle) =
-            api::ApiConnectionModeProvider::new(cache_dir.clone());
-        proxy_provider_handle.set_relay_selector(relay_selector.clone());
+        let proxy_provider =
+            api::ApiConnectionModeProvider::new(cache_dir.clone(), relay_selector.clone());
         let api_handle = api_runtime
             .mullvad_rest_handle(proxy_provider, endpoint_updater.callback())
             .await;