Remove the use of wildcard postMessage (codesandbox#1088)

Author: CompuIves

packages/app/src/app/components/Preview/index.js

@@ -113,7 +113,7 @@ class BasePreview extends React.Component<Props, State> {
   handleMessage = (data: Object, source: any) => {
     if (data && data.codesandbox) {
       if (data.type === 'initialized' && source) {
-        registerFrame(source);
+        registerFrame(source, frameUrl(this.props.sandbox.id));
 
         if (!this.state.frameInitialized && this.props.onInitialized) {
           this.disposeInitializer = this.props.onInitialized(this);

packages/codesandbox-api/src/dispatcher/index.ts

@@ -1,36 +1,52 @@
 // import * as debug from 'debug';
 import host from './host';
 
-const bundlers: Window[] = [];
+const bundlers: Map<Window, string> = new Map();
 
 // Whether the tab has a connection with the editor
 export const isStandalone =
   typeof window === 'undefined' || (!window.opener && window.parent === window);
 
+let parentOrigin: string | null = null;
+
+const parentOriginListener = (e: MessageEvent) => {
+  if (e.data.type === 'register-frame') {
+    parentOrigin = e.data.origin;
+
+    self.removeEventListener('message', parentOriginListener);
+  }
+};
+
+self.addEventListener('message', parentOriginListener);
+
 /**
  * Send a message to the editor, this is most probably an action you generated
  *
  * @export
  * @param {*} message
  * @returns
  */
-export function dispatch(message: Object) {
+export function dispatch(message: any) {
   if (!message) return;
 
   const newMessage = { ...message, codesandbox: true };
   notifyListeners(newMessage);
   notifyFrames(newMessage);
 
   if (isStandalone) return;
+  if (parentOrigin === null && message.type !== 'initialized') return;
 
   if (window.opener) {
-    window.opener.postMessage(newMessage, '*');
+    window.opener.postMessage(newMessage, parentOrigin === null ? '*' : parentOrigin);
   } else {
-    window.parent.postMessage(newMessage, '*');
+    window.parent.postMessage(newMessage, parentOrigin === null ? '*' : parentOrigin);
   }
 }
 
-export type Callback = (message: Object, source?: Window | null | undefined) => void;
+export type Callback = (
+  message: Object,
+  source?: MessageEvent['source'] | null | undefined
+) => void;
 
 const listeners: { [id: string]: Callback } = {};
 let listenerId = 0;
@@ -58,17 +74,17 @@ export function notifyListeners(data: Object, source?: MessageEvent['source']) {
 
 function notifyFrames(message: Object) {
   const rawMessage = JSON.parse(JSON.stringify(message));
-  bundlers.forEach(frame => {
+  bundlers.forEach((origin, frame) => {
     if (frame && frame.postMessage) {
-      frame.postMessage({ ...rawMessage, codesandbox: true }, '*');
+      frame.postMessage({ ...rawMessage, codesandbox: true }, origin);
     }
   });
 }
 
 function eventListener(e: MessageEvent) {
   const { data } = e;
 
-  if (data && data.codesandbox) {
+  if (data && data.codesandbox && (parentOrigin === null || e.origin === parentOrigin)) {
     notifyListeners(data, e.source);
   }
 }
@@ -78,9 +94,17 @@ function eventListener(e: MessageEvent) {
  *
  * @param frame
  */
-export function registerFrame(frame: Window) {
-  if (bundlers.indexOf(frame) === -1) {
-    bundlers.push(frame);
+export function registerFrame(frame: Window, origin: string) {
+  if (!bundlers.has(frame)) {
+    bundlers.set(frame, origin);
+
+    frame.postMessage(
+      {
+        type: 'register-frame',
+        origin: document.location.origin,
+      },
+      origin
+    );
   }
 }