From 4c26fcee3a37a140404707d65fd90f02d9375980 Mon Sep 17 00:00:00 2001
From: Hornwitser <github@hornwitser.no>
Date: Fri, 3 Apr 2020 20:35:20 +0200
Subject: [PATCH] Fix connection shutdown logic

Fix the WebSocket link not being properly teared down on the master
side causing the handshake to deadlock when the slave tried to connect
again.
---
 lib/link/connectors.js | 20 +++++++++++++-------
 lib/link/link.js       |  1 +
 master.js              | 31 ++++++++++++++++---------------
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/lib/link/connectors.js b/lib/link/connectors.js
index 55a33436..276d2922 100644
--- a/lib/link/connectors.js
+++ b/lib/link/connectors.js
@@ -90,14 +90,24 @@ class WebSocketBaseConnector extends events.EventEmitter {
 	 * @returns the sequence number of the message sent
 	 */
 	send(type, data) {
-		this._check(["handshake", "connected"]);
+		this._check(["handshake", "connected", "closing"]);
 		let message = { seq: this._seq, type, data };
 		this._sendBuffer.push(message);
-		if (this._state === "connected") {
+		if (["connected", "closing"].includes(this._state)) {
 			this._socket.send(JSON.stringify(message));
 		}
 		return this._seq++;
 	}
+
+	/**
+	 * Set connection state to closing
+	 *
+	 * Signal the connection is about tho close and that once the close
+	 * frame is received it should be considered finished.
+	 */
+	setClosing() {
+		this._state = "closing";
+	}
 }
 
 /**
@@ -156,12 +166,8 @@ class WebSocketClientConnector extends WebSocketBaseConnector {
 		}
 
 		this.stopHeartbeat();
-
-		if (["handshake", "connected"].includes(this._state)) {
-			this._socket.close(code, reason);
-		}
-
 		this._state = "closing";
+		this._socket.close(code, reason);
 	}
 
 	/**
diff --git a/lib/link/link.js b/lib/link/link.js
index f387f8f8..20377652 100644
--- a/lib/link/link.js
+++ b/lib/link/link.js
@@ -224,6 +224,7 @@ class Link {
 		}
 
 		await Promise.all(promises);
+		this.connector.setClosing();
 	}
 }
 
diff --git a/master.js b/master.js
index ef08616b..839af8fc 100644
--- a/master.js
+++ b/master.js
@@ -347,6 +347,12 @@ class BaseConnection extends link.Link {
 			event.send(slaveConnection, message.data);
 		}
 	}
+
+	async shutdown(code, reason) {
+		await link.messages.shutdownConnection.send(this);
+		this.connector.close(code, reason);
+		await events.once(this.connector, "close");
+	}
 }
 
 let controlConnections = new Array();
@@ -357,7 +363,7 @@ class ControlConnection extends BaseConnection {
 		this._agent = registerData.agent;
 		this._version = registerData.version;
 
-		this.connector.on('disconnect', () => {
+		this.connector.on("close", () => {
 			let index = controlConnections.indexOf(this);
 			if (index !== -1) {
 				controlConnections.splice(index, 1);
@@ -498,7 +504,7 @@ class SlaveConnection extends BaseConnection {
 			// XXX prometheusWsUsageCounter.labels('message', "other").inc();
 		});
 
-		this.connector.on('disconnect', () => {
+		this.connector.on("close", () => {
 			if (slaveConnections.get(this._id) === this) {
 				slaveConnections.delete(this._id);
 			}
@@ -541,12 +547,6 @@ class SlaveConnection extends BaseConnection {
 			}
 		}
 	}
-
-	async shutdown(code, reason) {
-		await link.messages.shutdownConnection.send(this);
-		this.connector.close(code, reason);
-		await events.once(this.connector, "close");
-	}
 }
 
 const wss = new WebSocket.Server({
@@ -582,6 +582,7 @@ class WebSocketServerConnector extends link.WebSocketBaseConnector {
 		// closing: Connection is in the process of being closed.
 		// closed: Connection has been closed
 		this._state = "handshake";
+		this._connected = false;
 	}
 
 	_check(expectedStates) {
@@ -609,7 +610,9 @@ class WebSocketServerConnector extends link.WebSocketBaseConnector {
 		}));
 
 		this._state = "connected";
+		this._connected = true;
 		this._attachSocketHandlers();
+		this.emit("connected");
 	}
 
 	/**
@@ -633,11 +636,13 @@ class WebSocketServerConnector extends link.WebSocketBaseConnector {
 		}));
 
 		this._state = "connected";
+		this._connected = true;
 		this._attachSocketHandlers();
 		this._dropSendBufferSeq(lastSeq);
 		for (let message of this._sendBuffer) {
 			this._socket.send(JSON.stringify(message));
 		}
+		this.emit("connected");
 	}
 
 	_attachSocketHandlers() {
@@ -649,7 +654,7 @@ class WebSocketServerConnector extends link.WebSocketBaseConnector {
 				this._lastReceivedSeq = null;
 				this._sendBuffer.length = 0;
 				this._state = "closed";
-				activeConnections.delete(this._sessionId);
+				activeConnectors.delete(this._sessionId);
 
 				if (this._connected) {
 					this.emit("close");
@@ -715,12 +720,8 @@ class WebSocketServerConnector extends link.WebSocketBaseConnector {
 		}
 
 		this.stopHeartbeat();
-
-		if (["handshake", "connected"].includes(this._state)) {
-			this._socket.close(code, reason);
-		}
-
 		this._state = "closing";
+		this._socket.close(code, reason);
 	}
 }
 
@@ -826,7 +827,7 @@ async function handleHandshake(message, socket, req, attachHandler) {
 	if (type === "register_slave") {
 		let connection = slaveConnections.get(data.id);
 		if (connection) {
-			console.log(`SOCKET | disconecting existing connection for slave ${data.id}`);
+			console.log(`SOCKET | disconnecting existing connection for slave ${data.id}`);
 			await connection.shutdown();
 		}