Apply CodeFrame postMessage origin fix to main

src/components/CodeEmbed/frame.tsx

@@ -8,6 +8,7 @@ interface CodeBundle {
   base?: string;
   scripts?: string[];
   instanceMode?: boolean;
+  parentOrigin: string;
 }
 
 /*
@@ -27,7 +28,7 @@ const wrapSketch = (sketchCode?: string) => {
 
 /*
  * Wraps the given code in a html document for display.
- * Single object argument, all properties optional:
+ * Single object argument:
  */
 const wrapInMarkup = (code: CodeBundle) =>
   `<!DOCTYPE html>
@@ -45,13 +46,18 @@ ${code.css || ""}
 </style>
 <!-- If we need an addon script, load p5 the usual way with no caching to make sure
 the import order doesn't get messed up. -->
-${(code.instanceMode || (code.scripts?.length ?? 0) > 0 ? [cdnLibraryUrl, ...(code.scripts ?? [])] : []).map((src) => `<script type="text/javascript" src="${src}"></script>`).join('\n')}
+${(code.instanceMode || (code.scripts?.length ?? 0) > 0 ? [cdnLibraryUrl, ...(code.scripts ?? [])] : []).map((src) => `<script type="text/javascript" src="${src}"></script>`).join("\n")}
 <body>${code.htmlBody || ""}</body>
 <script id="code" type="text/javascript">${wrapSketch(code.js) || ""}</script>
-${(code.scripts?.length ?? 0) > 0 ? '' : `
+${
+  (code.scripts?.length ?? 0) > 0
+    ? ""
+    : `
 <script type="text/javascript">
   // Listen for p5.min.js text content and include in iframe's head as script
   window.addEventListener("message", event => {
+    if (event.origin !== '${code.parentOrigin}') return;
+    if (event.source !== window.parent) return;
     // Include check to prevent p5.min.js from being loaded twice
     const scriptExists = !!document.getElementById("p5ScriptTagInIframe");
     if (!scriptExists && event.data?.sender === '${cdnLibraryUrl}') {
@@ -63,7 +69,8 @@ ${(code.scripts?.length ?? 0) > 0 ? '' : `
     }
   })
 </script>
-`}
+`
+}
 `.replace(/\u00A0/g, " ");
 
 export interface CodeFrameProps {
@@ -140,14 +147,14 @@ export const CodeFrame = (props: CodeFrameProps) => {
             sender: cdnLibraryUrl,
             message: p5ScriptText,
           },
-          "*",
+          window.location.origin,
         );
       } catch (e) {
         console.error(`Error loading ${p5ScriptTag.src}`);
         return;
       }
     })();
-  }, [props.jsCode, mounted,p5ScriptTag]);
+  }, [props.jsCode, mounted, p5ScriptTag]);
 
   return (
     <div
@@ -156,19 +163,24 @@ export const CodeFrame = (props: CodeFrameProps) => {
     >
       <iframe
         ref={iframeRef}
-        srcDoc={mounted ? wrapInMarkup({
-          js: props.jsCode,
-          css: props.cssCode,
-          htmlBody: props.htmlBodyCode,
-          base: props.base,
-          scripts: props.scripts,
-          instanceMode: props.jsCode.includes('new p5'),
-        }) : ""}
+        srcDoc={
+          mounted
+            ? wrapInMarkup({
+                js: props.jsCode,
+                css: props.cssCode,
+                htmlBody: props.htmlBodyCode,
+                base: props.base,
+                scripts: props.scripts,
+                instanceMode: props.jsCode.includes("new p5"),
+                parentOrigin: window.location.origin,
+              })
+            : ""
+        }
         sandbox="allow-scripts allow-popups allow-modals allow-forms allow-same-origin"
         aria-label="Code Preview"
         title="Code Preview"
         loading={props.lazyLoad ? "lazy" : "eager"}
-        style={{width: "100%", height: "100%"}}
+        style={{ width: "100%", height: "100%" }}
       />
     </div>
   );

test/components/CodeFrame.test.tsx

@@ -1,6 +1,14 @@
-import { it, expect, suite, assert } from "vitest";
-import { render, screen, act, cleanup, within } from "@testing-library/preact";
+import { it, expect, suite, assert, vi } from "vitest";
+import {
+  render,
+  screen,
+  act,
+  cleanup,
+  waitFor,
+  within,
+} from "@testing-library/preact";
 import { CodeFrame } from "@components/CodeEmbed/frame";
+import { cdnLibraryUrl } from "@/src/globals/globals";
 
 suite("CodeFrame", () => {
   it("defaults to empty DOM", async () => {
@@ -30,4 +38,65 @@ suite("CodeFrame", () => {
 
     cleanup();
   });
+
+  it("only sends the p5 script to its own iframe origin", async () => {
+    const postMessage = vi.fn();
+    const contentWindowSpy = vi
+      .spyOn(HTMLIFrameElement.prototype, "contentWindow", "get")
+      .mockReturnValue({ postMessage } as unknown as Window);
+    class ImmediateIntersectionObserver {
+      private callback: IntersectionObserverCallback;
+
+      constructor(callback: IntersectionObserverCallback) {
+        this.callback = callback;
+      }
+
+      observe() {
+        this.callback(
+          [{ isIntersecting: true } as IntersectionObserverEntry],
+          this as unknown as IntersectionObserver,
+        );
+      }
+
+      disconnect() {}
+    }
+    vi.stubGlobal("IntersectionObserver", ImmediateIntersectionObserver);
+    vi.stubGlobal(
+      "fetch",
+      vi.fn().mockResolvedValue({ text: () => Promise.resolve("p5 source") }),
+    );
+    const p5Script = document.createElement("script");
+    p5Script.id = "p5ScriptTag";
+    p5Script.src = "/p5.min.js";
+    document.body.appendChild(p5Script);
+
+    try {
+      act(() => {
+        render(<CodeFrame jsCode="function setup() {}" />);
+      });
+
+      const frame = screen.getByTitle("Code Preview") as HTMLIFrameElement;
+
+      await waitFor(() =>
+        expect(postMessage).toHaveBeenCalledWith(
+          {
+            sender: cdnLibraryUrl,
+            message: "p5 source",
+          },
+          window.location.origin,
+        ),
+      );
+      expect(frame.srcdoc).toContain(
+        `if (event.origin !== '${window.location.origin}') return;`,
+      );
+      expect(frame.srcdoc).toContain(
+        "if (event.source !== window.parent) return;",
+      );
+    } finally {
+      cleanup();
+      p5Script.remove();
+      contentWindowSpy.mockRestore();
+      vi.unstubAllGlobals();
+    }
+  });
 });