Update postMessage handler for embedded query page to verify origin (#1676)

Author: jasoncoffman

docs_website/docs/configurations/infra_config.mdx

@@ -40,7 +40,7 @@ Otherwise you can also pass the environment variable directly when launching the
 
 ### Iframe Embedding
 
-`IFRAME_ALLOWED_ORIGINS`: This is the allowed list for embedding Querybook in an iframe. By default it can only be embedded by itself.
+`IFRAME_ALLOWED_ORIGINS`: This is the allowed list for embedding Querybook in an iframe. By default it can only be embedded by itself. When set, this also restricts which origins can send queries to the [embedded editor](/docs/integrations/embedded_iframe) via `postMessage`.
 
 ### Database
 

querybook/server/datasources/utils.py

@@ -1,5 +1,6 @@
 from app.flask_app import limiter
 from app.datasource import register
+from env import QuerybookSettings
 from lib.change_log import get_change_log_list, get_change_log_content_by_date
 
 
@@ -22,3 +23,8 @@ def test_ratelimit():
     Endpoint to ensure ratelimit works in prod
     """
     return "yes"
+
+
+@register("/utils/embedded/allowed_origins/", methods=["GET"])
+def get_embedded_allowed_origins():
+    return QuerybookSettings.IFRAME_ALLOWED_ORIGINS or []

querybook/webapp/components/EmbeddedQueryPage/EmbeddedQueryPage.tsx

@@ -2,14 +2,24 @@ import React from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 
 import QueryComposer from 'components/QueryComposer/QueryComposer';
+import { useResource } from 'hooks/useResource';
 import { IAdhocQuery } from 'const/adhocQuery';
 import { receiveAdhocQuery } from 'redux/adhocQuery/action';
 import { Dispatch, IStoreState } from 'redux/store/types';
+import { EmbeddedResource } from 'resource/embedded';
 import { Button } from 'ui/Button/Button';
 import { FullHeight } from 'ui/FullHeight/FullHeight';
 
 import './EmbeddedQueryPage.scss';
 
+function isOriginAllowed(origin: string, allowedOrigins: string[]): boolean {
+    if (allowedOrigins.length === 0 || origin === window.location.origin) {
+        return true;
+    }
+
+    return allowedOrigins.some((allowed) => origin === new URL(allowed).origin);
+}
+
 const EmbeddedQueryPage: React.FunctionComponent = () => {
     const environmentId = useSelector(
         (state: IStoreState) => state.environment.currentEnvironmentId
@@ -18,15 +28,30 @@ const EmbeddedQueryPage: React.FunctionComponent = () => {
         (state: IStoreState) => state.adhocQuery[environmentId]?.query ?? ''
     );
 
+    const { data: allowedOrigins } = useResource(
+        EmbeddedResource.getAllowedOrigins
+    );
+
     const dispatch: Dispatch = useDispatch();
     const setQuery = React.useCallback(
-        (query: IAdhocQuery) =>
-            dispatch(receiveAdhocQuery(query, environmentId)),
-        []
+        (adhocQuery: IAdhocQuery) =>
+            dispatch(receiveAdhocQuery(adhocQuery, environmentId)),
+        [dispatch, environmentId]
     );
 
-    const onMessage = React.useCallback((e) => {
-        if (e.data && e.data.type === 'SET_QUERY') {
+    const onMessage = React.useCallback(
+        (e: MessageEvent) => {
+            if (!e.data || e.data.type !== 'SET_QUERY') {
+                return;
+            }
+
+            if (!isOriginAllowed(e.origin, allowedOrigins)) {
+                console.warn(
+                    `[EmbeddedQueryPage] Blocked postMessage from untrusted origin: ${e.origin}`
+                );
+                return;
+            }
+
             const query: IAdhocQuery = {};
             if (e.data.value) {
                 query.query = e.data.value;
@@ -36,21 +61,30 @@ const EmbeddedQueryPage: React.FunctionComponent = () => {
             }
 
             setQuery(query);
-        }
-    }, []);
+        },
+        [allowedOrigins, setQuery]
+    );
 
     React.useEffect(() => {
-        // Tell the parent we are ready to receive query
+        if (!allowedOrigins) {
+            return;
+        }
+
+        // Tell the parent we are ready to receive query after allowed origins are loaded
         window.parent.postMessage({ type: 'SEND_QUERY' }, '*');
-    }, []);
+    }, [allowedOrigins]);
 
-    // Setup query receiver
     React.useEffect(() => {
+        if (!allowedOrigins) {
+            return;
+        }
+
+        // Attach event listener to receive query from parent
         window.addEventListener('message', onMessage, false);
         return () => {
             window.removeEventListener('message', onMessage);
         };
-    }, [onMessage]);
+    }, [allowedOrigins, onMessage]);
 
     return (
         <FullHeight flex={'column'} className="EmbeddedQueryPage">

querybook/webapp/resource/embedded.ts

@@ -0,0 +1,6 @@
+import ds from 'lib/datasource';
+
+export const EmbeddedResource = {
+    getAllowedOrigins: () =>
+        ds.fetch<string[]>(`/utils/embedded/allowed_origins/`),
+};