From b901ce82a7dd4cb5da546ac4868d88414c383a2e Mon Sep 17 00:00:00 2001
From: Alexander Zinchuk <alx.zinchuk@gmail.com>
Date: Tue, 12 Dec 2023 12:34:18 +0100
Subject: [PATCH] Electron: Disable navigation to external URLs (#3965)

---
 src/electron/localStorage.ts | 34 ++++++++++++++++++++++++++++------
 src/electron/utils.ts        | 18 ++++++++++++++++++
 src/electron/window.ts       | 12 +++++++++---
 3 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/src/electron/localStorage.ts b/src/electron/localStorage.ts
index 4c170d07c..b32a9e513 100644
--- a/src/electron/localStorage.ts
+++ b/src/electron/localStorage.ts
@@ -1,17 +1,39 @@
-import { getLastWindow } from './utils';
+import { checkIsWebContentsUrlAllowed, getLastWindow } from './utils';
 
 let localStorage: Record<string, any> | undefined;
 
 export async function captureLocalStorage(): Promise<void> {
-  localStorage = await (getLastWindow())?.webContents.executeJavaScript('({ ...localStorage });');
-}
+  const lastWindow = getLastWindow();
 
-export async function restoreLocalStorage(): Promise<void> {
-  if (!localStorage) {
+  if (!lastWindow) {
     return;
   }
 
-  await getLastWindow()?.webContents.executeJavaScript(
+  const contents = lastWindow.webContents;
+  const contentsUrl = contents.getURL();
+
+  if (!checkIsWebContentsUrlAllowed(contentsUrl)) {
+    return;
+  }
+
+  localStorage = await contents.executeJavaScript('({ ...localStorage });');
+}
+
+export async function restoreLocalStorage(): Promise<void> {
+  const lastWindow = getLastWindow();
+
+  if (!lastWindow || !localStorage) {
+    return;
+  }
+
+  const contents = lastWindow.webContents;
+  const contentsUrl = contents.getURL();
+
+  if (!checkIsWebContentsUrlAllowed(contentsUrl)) {
+    return;
+  }
+
+  await contents.executeJavaScript(
     Object.keys(localStorage).map(
       (key: string) => `localStorage.setItem('${key}', JSON.stringify(${localStorage![key]}))`,
     ).join(';'),
diff --git a/src/electron/utils.ts b/src/electron/utils.ts
index e754b4dea..1cb0705b5 100644
--- a/src/electron/utils.ts
+++ b/src/electron/utils.ts
@@ -5,6 +5,10 @@ import fs from 'fs';
 
 import type { TrafficLightPosition } from '../types/electron';
 
+import { BASE_URL, PRODUCTION_URL } from '../config';
+
+const ALLOWED_URL_ORIGINS = [BASE_URL!, PRODUCTION_URL].map((url) => (new URL(url).origin));
+
 export const IS_MAC_OS = process.platform === 'darwin';
 export const IS_WINDOWS = process.platform === 'win32';
 export const IS_LINUX = process.platform === 'linux';
@@ -56,6 +60,20 @@ export function getAppTitle(chatTitle?: string): string {
   return `${chatTitle} · ${appName}`;
 }
 
+export function checkIsWebContentsUrlAllowed(url: string): boolean {
+  if (!app.isPackaged) {
+    return true;
+  }
+
+  const parsedUrl = new URL(url);
+
+  if (parsedUrl.pathname === encodeURI(`${__dirname}/index.html`)) {
+    return true;
+  }
+
+  return ALLOWED_URL_ORIGINS.includes(parsedUrl.origin);
+}
+
 export const TRAFFIC_LIGHT_POSITION: Record<TrafficLightPosition, Point> = {
   standard: { x: 10, y: 20 },
   lowered: { x: 10, y: 52 },
diff --git a/src/electron/window.ts b/src/electron/window.ts
index 2bec304c6..e0ac186e5 100644
--- a/src/electron/window.ts
+++ b/src/electron/window.ts
@@ -12,8 +12,8 @@ import { processDeeplink } from './deeplink';
 import { captureLocalStorage, restoreLocalStorage } from './localStorage';
 import tray from './tray';
 import {
-  forceQuit, getAppTitle, getCurrentWindow, getLastWindow, hasExtraWindows, IS_FIRST_RUN, IS_MAC_OS,
-  IS_PREVIEW, IS_WINDOWS, reloadWindows, store, TRAFFIC_LIGHT_POSITION, windows,
+  checkIsWebContentsUrlAllowed, forceQuit, getAppTitle, getCurrentWindow, getLastWindow, hasExtraWindows,
+  IS_FIRST_RUN, IS_MAC_OS, IS_PREVIEW, IS_WINDOWS, reloadWindows, store, TRAFFIC_LIGHT_POSITION, windows,
 } from './utils';
 import windowStateKeeper from './windowState';
 
@@ -80,6 +80,12 @@ export function createWindow(url?: string) {
     return deviceType === 'hid' && ALLOWED_DEVICE_ORIGINS.includes(origin);
   });
 
+  window.webContents.on('will-navigate', (event, newUrl) => {
+    if (!checkIsWebContentsUrlAllowed(newUrl)) {
+      event.preventDefault();
+    }
+  });
+
   window.on('page-title-updated', (event: Event) => {
     event.preventDefault();
   });
@@ -139,7 +145,7 @@ export function createWindow(url?: string) {
 }
 
 function loadWindowUrl(window: BrowserWindow, url?: string, hash?: string): void {
-  if (url) {
+  if (url && checkIsWebContentsUrlAllowed(url)) {
     window.loadURL(url);
   } else if (!app.isPackaged) {
     window.loadURL(`http://localhost:1234${hash}`);