Merge branch 'manage-wg-keys-in-gui'

3 files changed, 98 insertions, 9 deletions

diff --git a/gui/src/main/daemon-rpc.ts b/gui/src/main/daemon-rpc.ts
index 91f9e8771c..e317ffb794 100644
--- a/gui/src/main/daemon-rpc.ts
+++ b/gui/src/main/daemon-rpc.ts
@@ -7,6 +7,7 @@ import {
   ILocation,
   IRelayList,
   ISettings,
+  KeygenEvent,
   RelaySettingsUpdate,
   TunnelState,
 } from '../shared/daemon-rpc-types';
@@ -330,6 +331,13 @@ const settingsSchema = partialObject({
   tunnel_options: tunnelOptionsSchema,
 });
 
+const keygenEventSchema = oneOf(
+  enumeration('too_many_keys', 'generation_failure'),
+  object({
+    new_key: string,
+  }),
+);
+
 const daemonEventSchema = oneOf(
   object({
     tunnel_state: tunnelStateSchema,
@@ -341,12 +349,7 @@ const daemonEventSchema = oneOf(
     relay_list: relayListSchema,
   }),
   object({
-    wireguard_key: oneOf(
-      enumeration('too_many_keys', 'generation_failure'),
-      object({
-        new_key: string,
-      }),
-    ),
+    wireguard_key: keygenEventSchema,
   }),
 );
 
@@ -524,6 +527,40 @@ export class DaemonRpc {
     }
   }
 
+  public async generateWireguardKey(): Promise<KeygenEvent> {
+    const response = await this.transport.send('generate_wireguard_key');
+    try {
+      const validatedResponse: any = validate(keygenEventSchema, response);
+      switch (validatedResponse) {
+        case 'too_many_keys':
+        case 'generation_failure':
+          return validatedResponse;
+        default:
+          return camelCaseObjectKeys(validatedResponse as object) as KeygenEvent;
+      }
+    } catch (error) {
+      throw new ResponseParseError(`Invalid response from generate_wireguard_key ${error}`);
+    }
+  }
+
+  public async getWireguardKey(): Promise<string | undefined> {
+    const response = await this.transport.send('get_wireguard_key');
+    try {
+      return validate(maybe(string), response) || undefined;
+    } catch (error) {
+      throw new ResponseParseError('Invalid response from get_wireguard_key');
+    }
+  }
+
+  public async verifyWireguardKey(): Promise<boolean> {
+    const response = await this.transport.send('verify_wireguard_key');
+    try {
+      return validate(boolean, response);
+    } catch (error) {
+      throw new ResponseParseError('Invalid response from verify_wireguard_key');
+    }
+  }
+
   public async getVersionInfo(): Promise<IAppVersionInfo> {
     const response = await this.transport.send('get_version_info', [], NETWORK_CALL_TIMEOUT);
     try {
diff --git a/gui/src/main/index.ts b/gui/src/main/index.ts
index bdc1c4f5ec..a0dcbeb50b 100644
--- a/gui/src/main/index.ts
+++ b/gui/src/main/index.ts
@@ -13,6 +13,7 @@ import {
   IRelayList,
   IRelayListHostname,
   ISettings,
+  KeygenEvent,
   RelaySettings,
   RelaySettingsUpdate,
   TunnelState,
@@ -126,6 +127,8 @@ class ApplicationMain {
   // The UI locale which is set once from onReady handler
   private locale = 'en';
 
+  private wireguardPublicKey?: string;
+
   public run() {
     // Since electron's GPU blacklists are broken, GPU acceleration won't work on older distros
     if (process.platform === 'linux') {
@@ -487,9 +490,7 @@ class ApplicationMain {
         } else if ('relayList' in daemonEvent) {
           this.setRelays(daemonEvent.relayList, this.settings.relaySettings);
         } else if ('wireguardKey' in daemonEvent) {
-          /// TODO: handle wireguard key events properly.
-          log.info(`Received new key event`);
-          log.info(daemonEvent);
+          this.handleWireguardKeygenEvent(daemonEvent.wireguardKey);
         }
       },
       (error: Error) => {
@@ -512,6 +513,28 @@ class ApplicationMain {
     }
   }
 
+  private setWireguardKey(wireguardKey?: string) {
+    this.wireguardPublicKey = wireguardKey;
+    if (this.windowController) {
+      IpcMainEventChannel.wireguardKeys.notify(this.windowController.webContents, wireguardKey);
+    }
+  }
+
+  private handleWireguardKeygenEvent(event: KeygenEvent) {
+    switch (event) {
+      case 'too_many_keys':
+      case 'generation_failure':
+        this.notificationController.notifyKeyGenerationFailed();
+        this.wireguardPublicKey = undefined;
+        break;
+      default:
+        this.wireguardPublicKey = event.newKey;
+    }
+    if (this.windowController) {
+      IpcMainEventChannel.wireguardKeys.notifyKeygenEvent(this.windowController.webContents, event);
+    }
+  }
+
   private setTunnelState(newState: TunnelState) {
     this.tunnelState = newState;
     this.updateTrayIcon(newState, this.settings.blockWhenDisconnected);
@@ -534,6 +557,7 @@ class ApplicationMain {
 
     if (oldSettings.accountToken !== newSettings.accountToken) {
       this.updateAccountHistory();
+      this.fetchWireguardKey();
     }
 
     if (this.windowController) {
@@ -831,6 +855,7 @@ class ApplicationMain {
       currentVersion: this.currentVersion,
       upgradeVersion: this.upgradeVersion,
       guiSettings: this.guiSettings.state,
+      wireguardPublicKey: this.wireguardPublicKey,
     }));
 
     IpcMainEventChannel.settings.handleAllowLan((allowLan: boolean) =>
@@ -888,6 +913,15 @@ class ApplicationMain {
       this.updateAccountHistory();
     });
 
+    IpcMainEventChannel.wireguardKeys.handleGenerateKey(async () => {
+      try {
+        return await this.daemonRpc.generateWireguardKey();
+      } catch {
+        return 'generation_failure';
+      }
+    });
+    IpcMainEventChannel.wireguardKeys.handleVerifyKey(() => this.daemonRpc.verifyWireguardKey());
+
     ipcMain.on('show-window', () => {
       const windowController = this.windowController;
       if (windowController) {
@@ -970,6 +1004,14 @@ class ApplicationMain {
     }
   }
 
+  private async fetchWireguardKey(): Promise<void> {
+    try {
+      this.setWireguardKey(await this.daemonRpc.getWireguardKey());
+    } catch (error) {
+      log.error(`Failed to fetch wireguard key: ${error.message}`);
+    }
+  }
+
   private updateDaemonsAutoConnect() {
     const daemonAutoConnect = this.guiSettings.autoConnect && getOpenAtLogin();
     if (daemonAutoConnect !== this.settings.autoConnect) {
diff --git a/gui/src/main/notification-controller.ts b/gui/src/main/notification-controller.ts
index 771e788722..4b0f59fbf6 100644
--- a/gui/src/main/notification-controller.ts
+++ b/gui/src/main/notification-controller.ts
@@ -109,6 +109,16 @@ export default class NotificationController {
     });
   }
 
+  public notifyKeyGenerationFailed() {
+    const notification = new Notification({
+      title: this.notificationTitle,
+      body: messages.pgettext('notifications', 'Wireguard key generation failed'),
+      silent: true,
+      icon: this.notificationIcon,
+    });
+    this.scheduleNotification(notification);
+  }
+
   public cancelPendingNotifications() {
     for (const notification of this.pendingNotifications) {
       notification.close();