From 04dfc91c517d71959ca493d0fd3d876654abf217 Mon Sep 17 00:00:00 2001
From: Shachar Erez <sha1@fb.com>
Date: Thu, 24 Mar 2022 09:49:16 -0700
Subject: [PATCH] Fix origin validation check

Reviewed By: mweststrate

Differential Revision: D35080146

fbshipit-source-id: 3b8353c23b7c9c2f537801513e518b5b23a11520
---
 .../src/comms/BrowserServerWebSocket.tsx              | 11 +++++++++++
 .../comms/__tests__/BrowserServerWebSocket.node.tsx   |  6 +++---
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/desktop/flipper-server-core/src/comms/BrowserServerWebSocket.tsx b/desktop/flipper-server-core/src/comms/BrowserServerWebSocket.tsx
index ef7a41fd1..60a6d2fcc 100644
--- a/desktop/flipper-server-core/src/comms/BrowserServerWebSocket.tsx
+++ b/desktop/flipper-server-core/src/comms/BrowserServerWebSocket.tsx
@@ -18,6 +18,8 @@ import SecureServerWebSocket, {
 } from './SecureServerWebSocket';
 import {SecureClientQuery} from './ServerAdapter';
 import {ClientDescription, DeviceOS} from 'flipper-common';
+import {URL} from 'url';
+import {isFBBuild} from '../fb-stubs/constants';
 
 interface BrowserConnectionCtx extends SecureConnectionCtx {
   clientConnection?: BrowserClientConnection;
@@ -147,6 +149,15 @@ class BrowserServerWebSocket extends SecureServerWebSocket {
 
   protected verifyClient(): ws.VerifyClientCallbackSync {
     return (info: {origin: string; req: IncomingMessage; secure: boolean}) => {
+      if (isFBBuild) {
+        try {
+          const urlObj = new URL(info.origin);
+          if (urlObj.hostname.endsWith('.facebook.com')) {
+            return true;
+          }
+        } catch {}
+      }
+
       const ok = getFlipperServerConfig().validWebSocketOrigins.some(
         (validPrefix) => info.origin.startsWith(validPrefix),
       );
diff --git a/desktop/flipper-server-core/src/comms/__tests__/BrowserServerWebSocket.node.tsx b/desktop/flipper-server-core/src/comms/__tests__/BrowserServerWebSocket.node.tsx
index 3bb2807dd..6dbc602e6 100644
--- a/desktop/flipper-server-core/src/comms/__tests__/BrowserServerWebSocket.node.tsx
+++ b/desktop/flipper-server-core/src/comms/__tests__/BrowserServerWebSocket.node.tsx
@@ -23,7 +23,7 @@ import {createMockSEListener, WSMessageAccumulator} from './utils';
 
 jest.mock('../../FlipperServerConfig');
 (getFlipperServerConfig as jest.Mock).mockImplementation(() => ({
-  validWebSocketOrigins: ['localhost:'],
+  validWebSocketOrigins: ['http://localhost'],
 }));
 
 describe('BrowserServerWebSocket', () => {
@@ -62,7 +62,7 @@ describe('BrowserServerWebSocket', () => {
     const clientReceivedMessages = new WSMessageAccumulator();
     wsClient = new WebSocket(
       `ws://localhost:${port}?device_id=${deviceId}&device=${device}&app=${app}&os=${os}&sdk_version=${sdkVersion}`,
-      {origin: 'localhost:'},
+      {origin: 'http://localhost'},
     );
     wsClient.onmessage = ({data}) => clientReceivedMessages.add(data);
     await new Promise<void>((resolve, reject) => {
@@ -164,7 +164,7 @@ describe('BrowserServerWebSocket', () => {
     const clientReceivedMessages = new WSMessageAccumulator();
     wsClient = new WebSocket(
       `ws://localhost:${port}?deviceId=${deviceId}&device=${device}`,
-      {origin: 'localhost:'},
+      {origin: 'http://localhost'},
     );
     wsClient.onmessage = ({data}) => clientReceivedMessages.add(data);
     await new Promise<void>((resolve, reject) => {