[0.11>master] Prevent reusing the socket after receiving a server_dis…

…connect (#647)
microsoft · Nov 23, 2019 · 9be786c · 9be786c
2 parents 42a349a + 84700b7
commit 9be786c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/packages/drivers/odsp-socket-storage/src/OdspDocumentDeltaConnection.ts b/packages/drivers/odsp-socket-storage/src/OdspDocumentDeltaConnection.ts
@@ -12,7 +12,6 @@ import {
     IClient,
     IConnect,
 } from "@microsoft/fluid-protocol-definitions";
-import * as assert from "assert";
 import { IOdspSocketError } from "./contracts";
 import { debug } from "./debug";
 import { errorObjectFromOdspError } from "./OdspUtils";
@@ -112,6 +111,14 @@ export class OdspDocumentDeltaConnection extends DocumentDeltaConnection impleme
         documentId: string,
         telemetryLogger: ITelemetryLogger): ISocketReference {
         let socketReference = OdspDocumentDeltaConnection.socketIoSockets.get(key);
+
+        // verify the socket is healthy before reusing it
+        if (socketReference && (!socketReference.socket || !socketReference.socket.connected)) {
+            // the socket is in a bad state. fully remove the reference
+            socketReference = undefined;
+            OdspDocumentDeltaConnection.removeSocketIoReference(key, true);
+        }
+
         if (socketReference) {
             telemetryLogger.sendTelemetryEvent({
                 references: socketReference.references,
@@ -146,6 +153,10 @@ export class OdspDocumentDeltaConnection extends DocumentDeltaConnection impleme
                 });
 
             socket.on("server_disconnect", (socketError: IOdspSocketError) => {
+                // the server always closes the socket after sending this message
+                // fully remove the socket reference now
+                OdspDocumentDeltaConnection.removeSocketIoReference(key, true);
+
                 // Raise it as disconnect.
                 // That produces cleaner telemetry (no errors) and keeps protocol simpler (and not driver-specific).
                 socket.emit("disconnect", errorObjectFromOdspError(socketError));
@@ -172,12 +183,11 @@ export class OdspDocumentDeltaConnection extends DocumentDeltaConnection impleme
     private static removeSocketIoReference(key: string, isFatalError?: boolean) {
         const socketReference = OdspDocumentDeltaConnection.socketIoSockets.get(key);
         if (!socketReference) {
-            // this is expected to happens if we removed the reference due the socket not being connected
+            // this is expected to happen if we removed the reference due the socket not being connected
             return;
         }
 
         socketReference.references--;
-        assert(socketReference.delayDeleteTimeout === undefined);
 
         debug(`Removed socketio reference for ${key}. Remaining references: ${socketReference.references}.`);
 
@@ -188,12 +198,19 @@ export class OdspDocumentDeltaConnection extends DocumentDeltaConnection impleme
                 socketReference.socket = undefined;
             }
 
+            // clear the pending deletion if there is one
+            if (socketReference.delayDeleteTimeout !== undefined) {
+                clearTimeout(socketReference.delayDeleteTimeout);
+                socketReference.delayDeleteTimeout = undefined;
+                socketReference.delayDeleteTimeoutSetTime = undefined;
+            }
+
             OdspDocumentDeltaConnection.socketIoSockets.delete(key);
             debug(`Deleted socketio reference for ${key}. Is fatal error: ${isFatalError}.`);
             return;
         }
 
-        if (socketReference.references === 0) {
+        if (socketReference.references === 0 && socketReference.delayDeleteTimeout === undefined) {
             socketReference.delayDeleteTimeout = setTimeout(() => {
                 OdspDocumentDeltaConnection.socketIoSockets.delete(key);
 
@@ -229,7 +246,7 @@ export class OdspDocumentDeltaConnection extends DocumentDeltaConnection impleme
             throw new Error("Invalid socket reference key");
         }
 
-        OdspDocumentDeltaConnection.removeSocketIoReference(this.socketReferenceKey);
+        OdspDocumentDeltaConnection.removeSocketIoReference(this.socketReferenceKey, socketProtocolError);
         this.socketReferenceKey = undefined;
 
         this.emit("disconnect", "client closing connection");

diff --git a/packages/drivers/socket-storage-shared/src/documentDeltaConnection.ts b/packages/drivers/socket-storage-shared/src/documentDeltaConnection.ts
@@ -383,6 +383,12 @@ export class DocumentDeltaConnection extends EventEmitter implements IDocumentDe
                 reject(createErrorObject("connect_timeout", "Socket connection timed out"));
             });
 
+            // Listen for disconnects
+            this.addConnectionListener("disconnect", () => {
+                this.disconnect(true);
+                reject(createErrorObject("disconnect", "Socket disconnectd"));
+            });
+
             this.addConnectionListener("connect_document_success", (response: IConnected) => {
                 this.removeTrackedListeners(true);
                 resolve(response);