Merge pull request #1606 from avinxshKD/fix/codeembed-postmessage-origin

Restrict CodeFrame postMessage origin

Author: doradocodes

src/components/CodeEmbed/frame.tsx

@@ -7,6 +7,7 @@ interface CodeBundle {
   js?: string;
   base?: string;
   scripts?: string[];
+  parentOrigin: string;
 }
 
 /*
@@ -26,7 +27,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>
@@ -51,6 +52,8 @@ ${(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}') {
@@ -139,7 +142,7 @@ export const CodeFrame = (props: CodeFrameProps) => {
             sender: cdnLibraryUrl,
             message: p5ScriptText,
           },
-          "*",
+          window.location.origin,
         );
       } catch (e) {
         console.error(`Error loading ${p5ScriptTag.src}`);
@@ -161,6 +164,7 @@ export const CodeFrame = (props: CodeFrameProps) => {
           htmlBody: props.htmlBodyCode,
           base: props.base,
           scripts: props.scripts,
+          parentOrigin: window.location.origin,
         }) : ""}
         sandbox="allow-scripts allow-popups allow-modals allow-forms allow-same-origin"
         aria-label="Code Preview"

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();
+    }
+  });
 });