fix: read existing iframe src

BuilderIO · Feb 1, 2022 · 90ef609 · 90ef609 · vercel · Feb 1, 2022
1 parent 47db069
commit 90ef609
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 22 deletions.
diff --git a/src/lib/web-worker/worker-iframe.ts b/src/lib/web-worker/worker-iframe.ts
@@ -7,39 +7,24 @@ import type { Node } from './worker-node';
 import { SCRIPT_TYPE } from '../utils';
 import { setInstanceStateValue } from './worker-state';
 import { StateProp, WorkerMessageType } from '../types';
+import type { WorkerInstance } from './worker-instance';
 
 export const HTMLIFrameDescriptorMap: PropertyDescriptorMap & ThisType<Node> = {
   contentDocument: {
     get() {
-      return (this as any).contentWindow.document;
+      return getIframeEnv(this).$document$;
     },
   },
 
   contentWindow: {
     get() {
-      // the winId of an iframe's contentWindow is the same
-      // as the instanceId of the containing iframe element
-      const $winId$ = this[InstanceIdKey];
-
-      if (!environments[$winId$]) {
-        createEnvironment(
-          {
-            $winId$,
-            // iframe contentWindow parent winId is the iframe element's winId
-            $parentWinId$: this[WinIdKey],
-            $url$: getter(this, ['src']) || 'about:blank',
-          },
-          true
-        );
-      }
-
-      return environments[$winId$].$window$;
+      return getIframeEnv(this).$window$;
     },
   },
 
   src: {
     get() {
-      let src = environments[this[InstanceIdKey]].$location$.href;
+      let src = getIframeEnv(this).$location$.href;
       if (src.startsWith('about')) {
         src = '';
       }
@@ -48,8 +33,7 @@ export const HTMLIFrameDescriptorMap: PropertyDescriptorMap & ThisType<Node> = {
     set(src: string) {
       let xhr = new XMLHttpRequest();
       let xhrStatus: number;
-      let winId = this[InstanceIdKey];
-      let env = environments[winId];
+      let env = getIframeEnv(this);
 
       env.$location$.href = src = resolveUrl(getEnv(this), src);
       env.$isLoading$ = 1;
@@ -73,7 +57,7 @@ export const HTMLIFrameDescriptorMap: PropertyDescriptorMap & ThisType<Node> = {
         );
 
         sendToMain(true);
-        webWorkerCtx.$postMessage$([WorkerMessageType.InitializeNextScript, winId]);
+        webWorkerCtx.$postMessage$([WorkerMessageType.InitializeNextScript, env.$winId$]);
       } else {
         setInstanceStateValue(this, StateProp.loadErrorStatus, xhrStatus);
         env.$isLoading$ = 0;
@@ -83,3 +67,23 @@ export const HTMLIFrameDescriptorMap: PropertyDescriptorMap & ThisType<Node> = {
 
   ...HTMLSrcElementDescriptorMap,
 };
+
+const getIframeEnv = (iframe: WorkerInstance) => {
+  // the winId of an iframe's contentWindow is the same
+  // as the instanceId of the containing iframe element
+  const $winId$ = iframe[InstanceIdKey];
+
+  if (!environments[$winId$]) {
+    createEnvironment(
+      {
+        $winId$,
+        // iframe contentWindow parent winId is the iframe element's winId
+        $parentWinId$: iframe[WinIdKey],
+        $url$: getter(iframe, ['src']) || 'about:blank',
+      },
+      true
+    );
+  }
+
+  return environments[$winId$];
+};
diff --git a/tests/platform/iframe/iframe.spec.ts b/tests/platform/iframe/iframe.spec.ts
@@ -122,4 +122,8 @@ test('iframe', async ({ page }) => {
   await page.waitForSelector('.testGetByTagName');
   const testGetByTagName = page.locator('#testGetByTagName');
   await expect(testGetByTagName).toHaveText('#document');
+
+  await page.waitForSelector('.testGetIframeSrc');
+  const testGetIframeSrc = page.locator('#testGetIframeSrc');
+  await expect(testGetIframeSrc).toHaveText('/platform/iframe/content.html');
 });
diff --git a/tests/platform/iframe/index.html b/tests/platform/iframe/index.html
@@ -799,6 +799,20 @@ <h1>Iframe</h1>
           })();
         </script>
       </li>
+
+      <li>
+        <strong>getElementsById('existing-iframe2').getAttribute('src')</strong>
+        <code id="testGetIframeSrc"></code>
+        <iframe id="existing-iframe2" src="content.html"></iframe>
+        <script type="text/partytown">
+          (function () {
+            const elm = document.getElementById('testGetIframeSrc');
+            const iframe = document.getElementById('existing-iframe2');
+            elm.textContent = new URL(iframe.getAttribute('src'), location.href).pathname;
+            elm.className = 'testGetIframeSrc';
+          })();
+        </script>
+      </li>
     </ul>
 
     <hr />