From 182db5a2db95e27759dfdb1278a6c8204675cd71 Mon Sep 17 00:00:00 2001
From: Marcin Cichocki <marcincichocki1993@gmail.com>
Date: Tue, 19 Oct 2021 14:04:44 +0200
Subject: [PATCH] fix(client-electron): prevent memory leak in preload

---
 src/electron/renderer/common.ts               | 22 +++++++--------
 .../renderer/components/StatusBar.tsx         |  4 +--
 src/electron/renderer/preload/ipc.ts          | 27 +++++++++----------
 3 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/src/electron/renderer/common.ts b/src/electron/renderer/common.ts
index a05706cf..18f88801 100644
--- a/src/electron/renderer/common.ts
+++ b/src/electron/renderer/common.ts
@@ -8,7 +8,6 @@ import {
   WorkerStatus,
 } from '@/electron/common';
 import { format, formatDistanceToNow } from 'date-fns';
-import type { IpcRendererEvent } from 'electron';
 import { useContext, useEffect, useState } from 'react';
 import { useParams } from 'react-router-dom';
 import { ScreenshotDisplayOutput } from 'screenshot-desktop';
@@ -50,17 +49,13 @@ export function transformTimestamp(timestamp: number) {
 
 export function useIpcEvent<T>(
   channels: IpcOnChannels[],
-  callback: (event: IpcRendererEvent, value: T) => void
+  callback: (value: T) => void
 ) {
   useEffect(() => {
-    channels.forEach((c) => {
-      api.on(c, callback);
-    });
+    const unsubscriptions = channels.map((c) => api.on(c, callback));
 
     return () => {
-      channels.forEach((c) => {
-        api.removeListener(c, callback);
-      });
+      unsubscriptions.forEach((u) => u());
     };
   }, []);
 }
@@ -70,7 +65,7 @@ export function useIpcEventDialog<T>(channel: IpcOnChannels) {
   const [data, setData] = useState<T>(null);
   const close = () => setIsOpen(false);
 
-  useIpcEvent([channel], (e, data: T) => {
+  useIpcEvent([channel], (data: T) => {
     setData(data);
     setIsOpen(true);
   });
@@ -81,7 +76,7 @@ export function useIpcEventDialog<T>(channel: IpcOnChannels) {
 export function useIpcState() {
   const [state, setState] = useState<State>(api.getState());
 
-  function handleEvent(e: IpcRendererEvent, { payload }: Action<State>) {
+  function handleEvent({ payload }: Action<State>) {
     setState(payload);
   }
 
@@ -116,16 +111,17 @@ export function dispatchAsyncRequest<TRes, TReq = any>(
   return new Promise<TRes>((resolve) => {
     const uuid = uuidv4();
     const req: Request<TReq> = { ...action, uuid };
+    let removeListener: () => void = null;
 
-    function onAsyncResponse(e: IpcRendererEvent, res: Response<TRes>) {
+    function onAsyncResponse(res: Response<TRes>) {
       if (res.uuid !== uuid) return;
 
-      api.removeListener('renderer:async-response', onAsyncResponse);
+      removeListener();
 
       resolve(res.data);
     }
 
-    api.on('renderer:async-response', onAsyncResponse);
+    removeListener = api.on('renderer:async-response', onAsyncResponse);
     api.send('main:async-request', req);
   });
 }
diff --git a/src/electron/renderer/components/StatusBar.tsx b/src/electron/renderer/components/StatusBar.tsx
index 3ed1e92a..63a5a991 100644
--- a/src/electron/renderer/components/StatusBar.tsx
+++ b/src/electron/renderer/components/StatusBar.tsx
@@ -68,7 +68,7 @@ function useSettingsChangeListener(delay = 2000) {
   const [show, setShow] = useState(false);
   let id: any = null;
 
-  useIpcEvent([ActionTypes.UPDATE_SETTINGS], (e, { meta }: any) => {
+  useIpcEvent([ActionTypes.UPDATE_SETTINGS], ({ meta }: any) => {
     if (meta?.notify === false) {
       return;
     }
@@ -90,7 +90,7 @@ function useSettingsChangeListener(delay = 2000) {
 function useDownloadProgress() {
   const [progress, setProgress] = useState(0);
 
-  useIpcEvent(['renderer:download-progress'], (e, info: ProgressInfo) => {
+  useIpcEvent(['renderer:download-progress'], (info: ProgressInfo) => {
     setProgress(info.percent);
   });
 
diff --git a/src/electron/renderer/preload/ipc.ts b/src/electron/renderer/preload/ipc.ts
index 12ee691f..ae317b53 100644
--- a/src/electron/renderer/preload/ipc.ts
+++ b/src/electron/renderer/preload/ipc.ts
@@ -46,13 +46,20 @@ function validateChannel(input: string, whitelist: readonly string[]) {
 }
 
 /** Wrapper around {@link ipcRenderer.on}. Accepts only curated list of channels. */
-export function on(
-  channel: IpcOnChannels,
-  listener: (event: Electron.IpcRendererEvent, ...args: any[]) => void
-) {
+export function on(channel: IpcOnChannels, listener: (...args: any[]) => void) {
   validateChannel(channel, onChannels);
 
-  ipcRenderer.on(channel, listener);
+  // Intentionally remove event from parameters.
+  const safeListener = (e: any, ...args: any[]) => listener(...args);
+
+  ipcRenderer.on(channel, safeListener);
+
+  return () => {
+    // Function passed to preload will be wrapped in a proxy and its
+    // reference will be lost. Creating local version of the listener will
+    // allow to free memory later on.
+    ipcRenderer.removeListener(channel, safeListener);
+  };
 }
 
 /** Wrapper around {@link ipcRenderer.send}. Accepts only curated list of channels. */
@@ -69,16 +76,6 @@ export function invoke<T = any>(channel: IpcInvokeChannels, ...args: any[]) {
   return ipcRenderer.invoke(channel, ...args) as Promise<T>;
 }
 
-/** Wrapper around {@link ipcRenderer.removeListener}. Accepts only curated list of channels. */
-export function removeListener(
-  channel: IpcOnChannels,
-  listener: (...args: any[]) => void
-) {
-  validateChannel(channel, onChannels);
-
-  ipcRenderer.removeListener(channel, listener);
-}
-
 export function getState(): State {
   return ipcRenderer.sendSync('main:get-state');
 }