Wait for WireGuard key generation before auto-connecting

1 files changed, 34 insertions, 1 deletions

diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index b911a19c6a..f6417bd89a 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -39,6 +39,7 @@ import {
   setupLogging,
 } from '../shared/logging';
 import consumePromise from '../shared/promise';
+import { Scheduler } from '../shared/scheduler';
 import AccountDataCache from './account-data-cache';
 import { getOpenAtLogin, setOpenAtLogin } from './autostart';
 import {
@@ -59,6 +60,8 @@ import WindowController from './window-controller';
 const DAEMON_RPC_PATH =
   process.platform === 'win32' ? '//./pipe/Mullvad VPN' : '/var/run/mullvad-vpn';
 
+const AUTO_CONNECT_FALLBACK_DELAY = 6000;
+
 enum AppQuitStage {
   unready,
   initiated,
@@ -179,6 +182,9 @@ class ApplicationMain {
     },
   );
 
+  private autoConnectOnWireguardKeyEvent = false;
+  private autoConnectFallbackScheduler = new Scheduler();
+
   public run() {
     // Since electron's GPU blacklists are broken, GPU acceleration won't work on older distros
     if (process.platform === 'linux') {
@@ -588,6 +594,10 @@ class ApplicationMain {
     if (this.windowController) {
       IpcMainEventChannel.wireguardKeys.notify(this.windowController.webContents, wireguardKey);
     }
+
+    if (wireguardKey) {
+      this.wireguardKeygenEventAutoConnect();
+    }
   }
 
   private handleWireguardKeygenEvent(event: KeygenEvent) {
@@ -599,9 +609,12 @@ class ApplicationMain {
       default:
         this.wireguardPublicKey = event.newKey;
     }
+
     if (this.windowController) {
       IpcMainEventChannel.wireguardKeys.notifyKeygenEvent(this.windowController.webContents, event);
     }
+
+    this.wireguardKeygenEventAutoConnect();
   }
 
   private setTunnelState(newState: TunnelState) {
@@ -1101,11 +1114,21 @@ class ApplicationMain {
         log.warn(`Failed to get account data, logging in anyway: ${verification.error.message}`);
       }
 
+      this.autoConnectOnWireguardKeyEvent = true;
       await this.daemonRpc.setAccount(accountToken);
-      consumePromise(this.autoConnect());
+
+      // Fallback if daemon doesn't send event.
+      if (this.autoConnectOnWireguardKeyEvent) {
+        this.autoConnectFallbackScheduler.schedule(
+          () => this.wireguardKeygenEventAutoConnect(),
+          AUTO_CONNECT_FALLBACK_DELAY,
+        );
+      }
     } catch (error) {
       log.error(`Failed to login: ${error.message}`);
 
+      this.autoConnectOnWireguardKeyEvent = false;
+
       if (error instanceof InvalidAccountError) {
         throw Error(messages.gettext('Invalid account number'));
       } else {
@@ -1114,6 +1137,14 @@ class ApplicationMain {
     }
   }
 
+  private wireguardKeygenEventAutoConnect() {
+    if (this.autoConnectOnWireguardKeyEvent) {
+      this.autoConnectOnWireguardKeyEvent = false;
+      this.autoConnectFallbackScheduler.cancel();
+      consumePromise(this.autoConnect());
+    }
+  }
+
   private async autoConnect() {
     if (
       !this.accountData ||
@@ -1136,6 +1167,8 @@ class ApplicationMain {
         global.clearTimeout(this.accountExpiryNotificationTimeout);
         this.accountExpiryNotificationTimeout = undefined;
       }
+
+      this.autoConnectFallbackScheduler.cancel();
     } catch (error) {
       log.info(`Failed to logout: ${error.message}`);