From 8a1e484d0c7a86adb6c5d261966842d2faae10fe Mon Sep 17 00:00:00 2001
From: Lorenzo Blasa <lblasa@fb.com>
Date: Thu, 16 Jun 2022 10:35:17 -0700
Subject: [PATCH] Handle ws operations in a safe manner

Summary:
After doing some tests, any error thrown during ws events can result in flipper-server shutting down.

To avoid this, effectively, errors throughout should be properly handled.

The problem is that there's no guarantees this is or will be the case.

Instead, wrap the event handlers in a safe function execution wrapper. Any errors will be caught and logged.

Reviewed By: aigoncharov

Differential Revision: D37206923

fbshipit-source-id: 6f7cadc297ac39768030962c6eaadde55048fd21
---
 .../src/server/attachSocketServer.tsx         | 44 ++++++++++++++-----
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/desktop/flipper-server-core/src/server/attachSocketServer.tsx b/desktop/flipper-server-core/src/server/attachSocketServer.tsx
index 12d1001bd..342a11334 100644
--- a/desktop/flipper-server-core/src/server/attachSocketServer.tsx
+++ b/desktop/flipper-server-core/src/server/attachSocketServer.tsx
@@ -19,13 +19,23 @@ import {
   CompanionEventWebSocketMessage,
 } from 'flipper-common';
 import {FlipperServerImpl} from '../FlipperServerImpl';
-import {WebSocketServer} from 'ws';
+import {RawData, WebSocketServer} from 'ws';
 import {
   FlipperServerCompanion,
   FlipperServerCompanionEnv,
 } from 'flipper-server-companion';
 import {URLSearchParams} from 'url';
 
+const safe = (f: () => void) => {
+  try {
+    f();
+  } catch (error) {
+    if (error instanceof Error) {
+      console.error(error.name, error.stack);
+    }
+  }
+};
+
 /**
  * Attach and handle incoming messages from clients.
  * @param flipperServer A FlipperServer instance.
@@ -42,7 +52,7 @@ export function attachSocketServer(
         ` ${req.socket.remoteAddress}:${req.socket.remotePort}`) ||
       '';
 
-    console.log(`Client connected${clientAddress}`);
+    console.log('Client connected', clientAddress);
 
     let connected = true;
 
@@ -117,7 +127,7 @@ export function attachSocketServer(
 
     flipperServerCompanion?.onAny(onServerCompanionEvent);
 
-    client.on('message', (data) => {
+    async function onClientMessage(data: RawData) {
       let [event, payload]: [event: string | null, payload: any | null] = [
         null,
         null,
@@ -205,20 +215,30 @@ export function attachSocketServer(
             });
         }
       }
+    }
+
+    client.on('message', (data) => {
+      safe(() => onClientMessage(data));
     });
 
+    async function onClientClose(error: Error | undefined = undefined) {
+      if (error) {
+        console.error(`Client disconnected ${clientAddress} with error`, error);
+      } else {
+        console.log(`Client disconnected ${clientAddress}`);
+      }
+
+      connected = false;
+      flipperServer.offAny(onServerEvent);
+      flipperServerCompanion?.destroyAll();
+    }
+
     client.on('close', () => {
-      console.log(`Client disconnected ${clientAddress}`);
-      connected = false;
-      flipperServer.offAny(onServerEvent);
-      flipperServerCompanion?.destroyAll();
+      safe(() => onClientClose());
     });
 
-    client.on('error', (e) => {
-      console.error(`Socket error ${clientAddress}`, e);
-      connected = false;
-      flipperServer.offAny(onServerEvent);
-      flipperServerCompanion?.destroyAll();
+    client.on('error', (error) => {
+      safe(() => onClientClose(error));
     });
   });
 }