From a9a7320a1ef8854670eee619277736d4770cf682 Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Sun, 9 Jan 2022 17:46:43 +0000
Subject: [PATCH] Make disconnects graceful

---
 src/lib/agent/Peer.mjs       | 13 +++++++++++++
 src/lib/agent/PeerServer.mjs |  6 +++---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/lib/agent/Peer.mjs b/src/lib/agent/Peer.mjs
index 255a7ab..c1feda8 100644
--- a/src/lib/agent/Peer.mjs
+++ b/src/lib/agent/Peer.mjs
@@ -39,6 +39,10 @@ class Peer extends EventEmitter {
 		 * @type {{address:string,port:number}[]}
 		 */
 		this.known_peers = [];
+		
+		// If a message with the event name "end" is received, close our side
+		// of the connection
+		this.once(`message-end`, this.destroy);
 	}
 	
 	///////////////////////////////////////////////////////////////////////////
@@ -90,9 +94,16 @@ class Peer extends EventEmitter {
 			l.warn(`Our id (${this.server.our_id}) is equal to that of the remote (${this.id}), killing connection`);
 			return false;
 		}
+		
+		this.connection.on("message", this.__handle_message.bind(this));
 		return true;
 	}
 	
+	__handle_message(event_name, msg) {
+		this.emit("message", event_name, msg);
+		this.emit(`message-${event_name}`, msg);
+	}
+	
 	/**
 	 * Sends a hello message to this peer.
 	 * @return {Promise}	A Promise that resolves when the sending of the message is complete.
@@ -123,6 +134,8 @@ class Peer extends EventEmitter {
 	 * @return	{Promise}	A Promise that resolves once the connection is closed.
 	 */
 	async destroy() {
+		// TODO: If this takes too long, we should just ignore it and move on
+		await this.send("end", "goodbye");
 		await this.connection.destroy();
 		this.emit("destroy");
 		this.removeAllListeners();
diff --git a/src/lib/agent/PeerServer.mjs b/src/lib/agent/PeerServer.mjs
index 0aeeb66..f56e002 100644
--- a/src/lib/agent/PeerServer.mjs
+++ b/src/lib/agent/PeerServer.mjs
@@ -86,9 +86,9 @@ class PeerServer extends EventEmitter {
 		this.emit("peer", peer);
 	}
 	
-	async handle_message(peer, message) {
-		this.emit("message", peer, message);
-		this.emit(`message-${message.event}`, peer, message.message);
+	async handle_message(peer, event_name, msg) {
+		this.emit("message", peer, event_name, msg);
+		this.emit(`message-${event_name}`, peer, msg);
 	}
 	
 	async handle_destroy(peer) {