From e4814e00131a996876e26d34d0890f61092969f1 Mon Sep 17 00:00:00 2001
From: Josh Seltzer <jyseltz@gmail.com>
Date: Sun, 18 Dec 2022 15:12:39 -0500
Subject: [PATCH] fix for iframe css injection

---
 src/css/inject-style.ts | 24 ++++++++++++++++++++++--
 src/inject-css/index.ts | 20 +++++++-------------
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/src/css/inject-style.ts b/src/css/inject-style.ts
index 4131fdcd..9b08bfd2 100644
--- a/src/css/inject-style.ts
+++ b/src/css/inject-style.ts
@@ -2,8 +2,10 @@ import * as postcss from 'postcss';
 import { appendImportantToDeclarations } from './declaration';
 import { getCssWithExpandedImports } from './import';
 
-const getStylesheetId = (id: string) => {
-  return `stylebot-css-${id}`;
+const getStylesheetId = (id: string, iframeIdx?: number) => {
+  return iframeIdx !== undefined
+    ? `stylebot-css-${id}-iframe-${iframeIdx}`
+    : `stylebot-css-${id}`;
 };
 
 export const injectCSSIntoDocument = async (
@@ -12,6 +14,24 @@ export const injectCSSIntoDocument = async (
 ): Promise<void> => {
   const cssWithExpandedImports = await getCssWithExpandedImports(css);
 
+  // inject CSS into all iframes
+  const iframeElems = document.getElementsByTagName('iframe');
+  for (const [iframeIdx, iframeElem] of Array.from(iframeElems).entries()) {
+    const iframeStylesheetId = getStylesheetId(id, iframeIdx);
+    const el = document.getElementById(iframeStylesheetId);
+    if (el) {
+      el.innerHTML = cssWithExpandedImports;
+      return;
+    }
+    const iframeDoc = (iframeElem as any).contentWindow.document;
+    const iframeStyle = iframeDoc.createElement('style');
+    iframeStyle.type = 'text/css';
+    iframeStyle.setAttribute('id', iframeStylesheetId);
+    iframeStyle.appendChild(iframeDoc.createTextNode(cssWithExpandedImports));
+    iframeDoc.head.appendChild(iframeStyle);
+  }
+
+  // then inject CSS into top document
   const stylesheetId = getStylesheetId(id);
   const el = document.getElementById(stylesheetId);
 
diff --git a/src/inject-css/index.ts b/src/inject-css/index.ts
index e9703395..dee6d937 100644
--- a/src/inject-css/index.ts
+++ b/src/inject-css/index.ts
@@ -44,18 +44,12 @@ const injectCss = (
 };
 
 const run = () => {
-  if (window === window.top) {
-    injectCss({
-      name: 'GetStylesForPage',
-      important: true,
-    });
-  } else {
-    injectCss({
-      name: 'GetStylesForIframe',
-      url: window.location.href,
-      important: true,
-    });
-  }
+  injectCss({
+    name: 'GetStylesForPage',
+    important: true,
+  });
 };
 
-run();
+document.addEventListener('DOMContentLoaded', () => {
+  run();
+});