2 files changed, 39 insertions, 40 deletions

diff --git a/gui/packages/desktop/src/renderer/app.js b/gui/packages/desktop/src/renderer/app.js
index b3b22b9f2e..6d68cf4e1b 100644
--- a/gui/packages/desktop/src/renderer/app.js
+++ b/gui/packages/desktop/src/renderer/app.js
@@ -14,7 +14,7 @@ import { createMemoryHistory } from 'history';
 
 import makeRoutes from './routes';
 import ReconnectionBackoff from './lib/reconnection-backoff';
-import { DaemonRpc } from './lib/daemon-rpc';
+import { DaemonRpc, ConnectionObserver } from './lib/daemon-rpc';
 import NotificationController from './lib/notification-controller';
 import setShutdownHandler from './lib/shutdown-handler';
 import { NoAccountError } from './errors';
@@ -34,7 +34,6 @@ import type {
   TunnelState,
   DaemonRpcProtocol,
   AccountData,
-  ConnectionObserver as DaemonConnectionObserver,
 } from './lib/daemon-rpc';
 import type { ReduxStore } from './redux/store';
 import type { TrayIconType } from '../main/tray-icon-controller';
@@ -42,9 +41,15 @@ import type { TrayIconType } from '../main/tray-icon-controller';
 export default class AppRenderer {
   _notificationController = new NotificationController();
   _daemonRpc: DaemonRpcProtocol = new DaemonRpc();
+  _connectionObserver = new ConnectionObserver(
+    () => {
+      this._onOpenConnection();
+    },
+    (error) => {
+      this._onCloseConnection(error);
+    },
+  );
   _reconnectBackoff = new ReconnectionBackoff();
-  _openConnectionObserver: ?DaemonConnectionObserver;
-  _closeConnectionObserver: ?DaemonConnectionObserver;
   _memoryHistory = createMemoryHistory();
   _reduxStore: ReduxStore;
   _reduxActions: *;
@@ -75,13 +80,7 @@ export default class AppRenderer {
       ),
     };
 
-    this._openConnectionObserver = this._daemonRpc.addOpenConnectionObserver(() => {
-      this._onOpenConnection();
-    });
-
-    this._closeConnectionObserver = this._daemonRpc.addCloseConnectionObserver((error) => {
-      this._onCloseConnection(error);
-    });
+    this._daemonRpc.addConnectionObserver(this._connectionObserver);
 
     setShutdownHandler(async () => {
       log.info('Executing a shutdown handler');
@@ -89,7 +88,7 @@ export default class AppRenderer {
         await this.disconnectTunnel();
         log.info('Disconnected the tunnel');
       } catch (e) {
-        log.error(`Failed to shutdown tunnel: ${e.message}`);
+        log.error(`Failed to disconnect the tunnel: ${e.message}`);
       }
     });
 
@@ -112,7 +111,7 @@ export default class AppRenderer {
   }
 
   connect() {
-    this._connectToDaemon();
+    this._daemonRpc.connect({ path: getIpcPath() });
   }
 
   disconnect() {
@@ -424,10 +423,6 @@ export default class AppRenderer {
     actions.version.updateLatest(latestVersionInfo);
   }
 
-  async _connectToDaemon(): Promise<void> {
-    this._daemonRpc.connect({ path: getIpcPath() });
-  }
-
   async _onOpenConnection() {
     this._connectedToDaemon = true;
 
@@ -474,7 +469,7 @@ export default class AppRenderer {
     }
   }
 
-  async _onCloseConnection(error: ?Error) {
+  _onCloseConnection(error: ?Error) {
     const actions = this._reduxActions;
 
     // recover connection on error
@@ -678,10 +673,10 @@ class AccountDataCache {
   }
 }
 
-const getIpcPath = (): string => {
+function getIpcPath(): string {
   if (process.platform === 'win32') {
     return '//./pipe/Mullvad VPN';
   } else {
     return '/var/run/mullvad-vpn';
   }
-};
+}
diff --git a/gui/packages/desktop/src/renderer/lib/daemon-rpc.js b/gui/packages/desktop/src/renderer/lib/daemon-rpc.js
index 9e4c2c1604..ca45c74af2 100644
--- a/gui/packages/desktop/src/renderer/lib/daemon-rpc.js
+++ b/gui/packages/desktop/src/renderer/lib/daemon-rpc.js
@@ -260,6 +260,24 @@ const AppVersionInfoSchema = object({
   }),
 });
 
+export class ConnectionObserver {
+  _openHandler: () => void;
+  _closeHandler: (error: ?Error) => void;
+
+  constructor(openHandler: () => void, closeHandler: (error: ?Error) => void) {
+    this._openHandler = openHandler;
+    this._closeHandler = closeHandler;
+  }
+
+  _onOpen = () => {
+    this._openHandler();
+  };
+
+  _onClose = (error: ?Error) => {
+    this._closeHandler(error);
+  };
+}
+
 export interface DaemonRpcProtocol {
   connect({ path: string }): void;
   disconnect(): void;
@@ -280,8 +298,8 @@ export interface DaemonRpcProtocol {
   getLocation(): Promise<Location>;
   getState(): Promise<TunnelStateTransition>;
   subscribeStateListener((state: ?TunnelStateTransition, error: ?Error) => void): Promise<void>;
-  addOpenConnectionObserver(() => void): ConnectionObserver;
-  addCloseConnectionObserver((error: ?Error) => void): ConnectionObserver;
+  addConnectionObserver(observer: ConnectionObserver): void;
+  removeConnectionObserver(observer: ConnectionObserver): void;
   getAccountHistory(): Promise<Array<AccountToken>>;
   removeAccountFromHistory(accountToken: AccountToken): Promise<void>;
   getCurrentVersion(): Promise<string>;
@@ -301,10 +319,6 @@ export class ResponseParseError extends Error {
   }
 }
 
-export type ConnectionObserver = {
-  unsubscribe: () => void,
-};
-
 // Timeout used for RPC calls that do networking
 const NETWORK_CALL_TIMEOUT = 10000;
 
@@ -319,22 +333,12 @@ export class DaemonRpc implements DaemonRpcProtocol {
     this._transport.disconnect();
   }
 
-  addOpenConnectionObserver(handler: () => void): ConnectionObserver {
-    this._transport.on('open', handler);
-    return {
-      unsubscribe: () => {
-        this._transport.off('open', handler);
-      },
-    };
+  addConnectionObserver(observer: ConnectionObserver) {
+    this._transport.on('open', observer._onOpen).on('close', observer._onClose);
   }
 
-  addCloseConnectionObserver(handler: (error: ?Error) => void): ConnectionObserver {
-    this._transport.on('close', handler);
-    return {
-      unsubscribe: () => {
-        this._transport.off('close', handler);
-      },
-    };
+  removeConnectionObserver(observer: ConnectionObserver) {
+    this._transport.off('open', observer._onOpen).off('close', observer._onClose);
   }
 
   async getAccountData(accountToken: AccountToken): Promise<AccountData> {