Fix issue where messages were missed after handshake

After the first message was handled by handshake, we were not switching
onmessage handler synchronously. This caused messages to be missed.

Provide a callback to be called while handling the first message.
Switch onmessage handler in that callback, without waiting for promise resolution.

src/client/RemoteAdbDevice.ts

@@ -96,11 +96,12 @@ export class RemoteAdbDevice extends EventEmitter {
             this.ws.send(JSON.stringify(handshakeData));
 
             // Wait for the handshake response
-            this.serverHandshake = await getRemoteHandshake<ServerHandshake>(this.ws);
+            await getRemoteHandshake<ServerHandshake>(this.ws, async (handshake) => {
+                this.serverHandshake = handshake;
+                logger.log(this.backend.serial, `Connected as ${this.remoteSerial}`);
 
-            logger.log(this.backend.serial, `Connected as ${this.remoteSerial}`);
-
-            await this.backend.pipe(this.ws);
+                await this.backend.pipe(this.ws);
+            });
 
             this.emit("connected", this);
         } catch(e) {

src/common/handshake.ts

@@ -11,28 +11,40 @@ export type ServerHandshake = {
     serial: string,
 }
 
-export function getRemoteHandshake<T extends ClientHandshake|ServerHandshake>(ws: WebSocket): Promise<T> {
+export function getRemoteHandshake<T extends ClientHandshake|ServerHandshake>(ws: WebSocket, callback?: (handshakeData: T) => Promise<void>): Promise<T> {
     return new Promise<T>((resolve, reject) => {
         const onWebsocketClose = () => {
             reject(new Error("WebSocket closed while waiting for handshake"));
         };
 
         ws.addEventListener("close", onWebsocketClose, {once: true});
 
-        ws.addEventListener("message", (message) => {
+        ws.onmessage = async (message) => {
+            let handshakeData: T;
+
             try {
-                const handshakeData: T = JSON.parse(message.data);
+                handshakeData = JSON.parse(message.data);
 
                 if (handshakeData.type !== "handshake") {
                     throw new Error("Unexpected handshake message");
                 }
-
-                resolve(handshakeData);
             } catch(e) {
                 reject(new Error(`Handshake failed: ${e.message}`));
+                return;
             } finally {
                 ws.removeEventListener("close", onWebsocketClose);
             }
-        }, {once: true});
+
+            try {
+                // Reset onmessage and call the callback synchronously.
+                // This is required as onmessage needs to be updated synchronously to prevent missed messages.
+                ws.onmessage = undefined;
+                await callback?.(handshakeData);
+
+                resolve(handshakeData);
+            } catch(e) {
+                reject(new Error(`Handshake callback failed: ${e.message}`));
+            }
+        };
     });
 }

src/server/server.ts

@@ -8,6 +8,7 @@ import session from 'express-session';
 import bodyParser from 'body-parser';
 import WebSocket from 'ws';
 import stoppable from 'stoppable';
+import type { Duplex } from 'stream';
 
 import logger from '../common/logger';
 import { monitorAdbServer, addAdbDevice, removeAdbDevice } from './adbConnection';
@@ -208,12 +209,14 @@ export class Server {
         logger.log("Got new web socket connection. Waiting for handshake...");
 
         try {
-            let handshakeData: ClientHandshake = await getRemoteHandshake<ClientHandshake>(ws);
+            let wsStream: Duplex;
 
-            let wsStream = WebSocket.createWebSocketStream(ws);
-            wsStream.on("error", () => {
-                // Do nothing
-                // Sometimes this can happen when trying to write to the socket after it is closed in the process of closing the connection. Ignore.
+            let handshakeData: ClientHandshake = await getRemoteHandshake<ClientHandshake>(ws, async (handshake: ClientHandshake) => {
+                wsStream = WebSocket.createWebSocketStream(ws);
+                wsStream.on("error", () => {
+                    // Do nothing
+                    // Sometimes this can happen when trying to write to the socket after it is closed in the process of closing the connection. Ignore.
+                });
             });
 
             let port: number;