From bbb29a0d1b84d8e38ca0d79a54e90dac8becf337 Mon Sep 17 00:00:00 2001
From: Tasos Boulis <tboulis@hotmail.com>
Date: Mon, 30 Jan 2023 16:14:57 +0200
Subject: [PATCH] [MM-48407] Include entry for websocket too when trusting a
 certificate for https (#2526)

* Inlcude entry for websocket too when trusting a certificate for https

* Improve condition for protocol
---
 src/main/certificateStore.test.js | 55 +++++++++++++++++++------------
 src/main/certificateStore.ts      | 10 +++++-
 2 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/src/main/certificateStore.test.js b/src/main/certificateStore.test.js
index 9c143302..9c72adb1 100644
--- a/src/main/certificateStore.test.js
+++ b/src/main/certificateStore.test.js
@@ -32,13 +32,13 @@ jest.mock('fs', () => ({
 
 const certificateData = {
     'https://server-1.com': {
-        data: 'somerandomdata',
-        issuerName: 'someissuer',
+        data: 'someRandomData',
+        issuerName: 'someIssuer',
         dontTrust: false,
     },
     'https://server-2.com': {
-        data: 'somerandomdata',
-        issuerName: 'someissuer',
+        data: 'someRandomData',
+        issuerName: 'someIssuer',
         dontTrust: true,
     },
 };
@@ -49,7 +49,7 @@ describe('main/certificateStore', () => {
 
         let certificateStore;
         expect(() => {
-            certificateStore = new CertificateStore('somefilename');
+            certificateStore = new CertificateStore('someFilename');
         }).not.toThrow(Error);
         expect(certificateStore.data).toStrictEqual({});
     });
@@ -59,36 +59,49 @@ describe('main/certificateStore', () => {
         beforeAll(() => {
             validateCertificateStore.mockImplementation((data) => JSON.parse(data));
             fs.readFileSync.mockImplementation(() => JSON.stringify(certificateData));
-            certificateStore = new CertificateStore('somefilename');
+            certificateStore = new CertificateStore('someFilename');
         });
 
         it('should return true for stored matching certificate', () => {
-            certificateStore = new CertificateStore('somefilename');
+            certificateStore = new CertificateStore('someFilename');
 
             expect(certificateStore.isTrusted('https://server-1.com', {
-                data: 'somerandomdata',
-                issuerName: 'someissuer',
+                data: 'someRandomData',
+                issuerName: 'someIssuer',
             })).toBe(true);
         });
 
         it('should return false for missing url', () => {
             expect(certificateStore.isTrusted('https://server-3.com', {
-                data: 'somerandomdata',
-                issuerName: 'someissuer',
+                data: 'someRandomData',
+                issuerName: 'someIssuer',
             })).toBe(false);
         });
 
-        it('should return false for unmatching cert', () => {
+        it('should return false for unmatched cert', () => {
             expect(certificateStore.isTrusted('https://server-1.com', {
-                data: 'someotherrandomdata',
-                issuerName: 'someissuer',
+                data: 'someOtherRandomData',
+                issuerName: 'someIssuer',
             })).toBe(false);
 
             expect(certificateStore.isTrusted('https://server-1.com', {
-                data: 'somerandomdata',
-                issuerName: 'someotherissuer',
+                data: 'someRandomData',
+                issuerName: 'someOtherIssuer',
             })).toBe(false);
         });
+
+        it('should add certificate for websocket too', () => {
+            const certOrigin = 'https://server-websocket.com';
+            const wssCertOrigin = certOrigin.replace('https', 'wss');
+            const certData = {
+                data: 'someRandomData',
+                issuerName: 'someIssuer',
+            };
+
+            certificateStore = new CertificateStore('someFilename');
+            certificateStore.add(certOrigin, certData);
+            expect(certificateStore.isTrusted(wssCertOrigin, certData)).toBe(true);
+        });
     });
 
     describe('isExplicitlyUntrusted', () => {
@@ -96,20 +109,20 @@ describe('main/certificateStore', () => {
         beforeAll(() => {
             validateCertificateStore.mockImplementation((data) => JSON.parse(data));
             fs.readFileSync.mockImplementation(() => JSON.stringify(certificateData));
-            certificateStore = new CertificateStore('somefilename');
+            certificateStore = new CertificateStore('someFilename');
         });
 
         it('should return true for explicitly untrusted cert', () => {
             expect(certificateStore.isExplicitlyUntrusted('https://server-2.com', {
-                data: 'somerandomdata',
-                issuerName: 'someissuer',
+                data: 'someRandomData',
+                issuerName: 'someIssuer',
             })).toBe(true);
         });
 
         it('should return false for trusted cert', () => {
             expect(certificateStore.isExplicitlyUntrusted('https://server-1.com', {
-                data: 'somerandomdata',
-                issuerName: 'someissuer',
+                data: 'someRandomData',
+                issuerName: 'someIssuer',
             })).toBe(false);
         });
     });
diff --git a/src/main/certificateStore.ts b/src/main/certificateStore.ts
index bf672710..218e08b3 100644
--- a/src/main/certificateStore.ts
+++ b/src/main/certificateStore.ts
@@ -58,7 +58,15 @@ export class CertificateStore {
     };
 
     add = (targetURL: string, certificate: Certificate, dontTrust = false) => {
-        this.data[urlUtils.getHost(targetURL)] = comparableCertificate(certificate, dontTrust);
+        const host = urlUtils.getHost(targetURL);
+        const comparableCert = comparableCertificate(certificate, dontTrust);
+        this.data[host] = comparableCert;
+
+        // Trust certificate for websocket connections on the same origin.
+        if (host.startsWith('https://')) {
+            const wssHost = host.replace('https', 'wss');
+            this.data[wssHost] = comparableCert;
+        }
     };
 
     isExisting = (targetURL: string) => {