From e3e5da10a9a3d637ba73017809f36501fff02cda Mon Sep 17 00:00:00 2001
From: Ryan Cao <70191398+ryanccn@users.noreply.github.com>
Date: Thu, 19 Oct 2023 05:44:29 +0800
Subject: [PATCH] fix: Content Security Policy patching (#1814)

Co-authored-by: Vendicated <vendicated@riseup.net>
---
 src/main/index.ts | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/main/index.ts b/src/main/index.ts
index d1b7c5d97..481736a98 100644
--- a/src/main/index.ts
+++ b/src/main/index.ts
@@ -62,6 +62,10 @@ if (IS_VESKTOP || !IS_VANILLA) {
         } catch { }
 
 
+        const findHeader = (headers: Record<string, string[]>, headerName: Lowercase<string>) => {
+            return Object.keys(headers).find(h => h.toLowerCase() === headerName);
+        };
+
         // Remove CSP
         type PolicyResult = Record<string, string[]>;
 
@@ -73,6 +77,7 @@ if (IS_VESKTOP || !IS_VANILLA) {
                     result[directiveKey] = directiveValue;
                 }
             });
+
             return result;
         };
         const stringifyPolicy = (policy: PolicyResult): string =>
@@ -81,31 +86,39 @@ if (IS_VESKTOP || !IS_VANILLA) {
                 .map(directive => directive.flat().join(" "))
                 .join("; ");
 
-        function patchCsp(headers: Record<string, string[]>, header: string) {
-            if (header in headers) {
+        const patchCsp = (headers: Record<string, string[]>) => {
+            const header = findHeader(headers, "content-security-policy");
+
+            if (header) {
                 const csp = parsePolicy(headers[header][0]);
 
                 for (const directive of ["style-src", "connect-src", "img-src", "font-src", "media-src", "worker-src"]) {
-                    csp[directive] = ["*", "blob:", "data:", "vencord:", "'unsafe-inline'"];
+                    csp[directive] ??= [];
+                    csp[directive].push("*", "blob:", "data:", "vencord:", "'unsafe-inline'");
                 }
+
                 // TODO: Restrict this to only imported packages with fixed version.
                 // Perhaps auto generate with esbuild
                 csp["script-src"] ??= [];
                 csp["script-src"].push("'unsafe-eval'", "https://unpkg.com", "https://cdnjs.cloudflare.com");
                 headers[header] = [stringifyPolicy(csp)];
             }
-        }
+        };
 
         session.defaultSession.webRequest.onHeadersReceived(({ responseHeaders, resourceType }, cb) => {
             if (responseHeaders) {
                 if (resourceType === "mainFrame")
-                    patchCsp(responseHeaders, "content-security-policy");
+                    patchCsp(responseHeaders);
 
                 // Fix hosts that don't properly set the css content type, such as
                 // raw.githubusercontent.com
-                if (resourceType === "stylesheet")
-                    responseHeaders["content-type"] = ["text/css"];
+                if (resourceType === "stylesheet") {
+                    const header = findHeader(responseHeaders, "content-type");
+                    if (header)
+                        responseHeaders[header] = ["text/css"];
+                }
             }
+
             cb({ cancel: false, responseHeaders });
         });