From c395f0b05b1ba820aa172995c19bb89ecd0fc9ba Mon Sep 17 00:00:00 2001
From: extremeheat <extreme@protonmail.ch>
Date: Wed, 24 Aug 2022 00:41:12 -0400
Subject: [PATCH] relay: Add multi-user login support (#258)

* relay: add multi-user login support

* relay: Fix close handling, connect username
---
 docs/API.md  |  2 +-
 index.d.ts   | 12 ++++++++++++
 src/relay.js | 29 +++++++++++++++++++++++++----
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/docs/API.md b/docs/API.md
index 8d0b545..2f23b2d 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -170,7 +170,7 @@ For documentation on the protocol, and packets/fields see the [the protocol doc]
 
 ### Proxy docs
 
-You can create a proxy ("Relay") to create a machine-in-the-middle (MITM) connection to a server. You can observe and intercept packets as they go through. The Relay is a server+client combo with some special packet handling and forwarding that takes care of the authentication and encryption on the server side. You'll be asked to login if `offline` is not specified once you connect.
+You can create a proxy ("Relay") to create a machine-in-the-middle (MITM) connection to a server. You can observe and intercept packets as they go through. The Relay is a server+client combo with some special packet handling and forwarding that takes care of the authentication and encryption on the server side. Clients will be asked to login if `offline` is not specified on connection.
 
 ```js
 const { Relay } = require('bedrock-protocol')
diff --git a/index.d.ts b/index.d.ts
index e03ef61..140a05a 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -131,6 +131,10 @@ declare module "bedrock-protocol" {
      * Close the connection. Already called by disconnect. Call this to manually close RakNet connection. 
      */
     close()
+
+    on(event: 'login', cb: () => void)
+    on(event: 'join', cb: () => void)
+    on(event: 'close', cb: (reason: string) => void)
   }
 
   export class Server extends EventEmitter {
@@ -161,6 +165,14 @@ declare module "bedrock-protocol" {
     }
     // Whether to enable chunk caching (default: false)
     enableChunkCaching?: boolean
+
+    // Only allow one client to connect at a time (default: false)
+    forceSinge: boolean
+
+    // Dispatched when a new client has logged in, and we need authentication
+    // tokens to join the backend server. Cached after the first login.
+    // If this is not specified, the client will be disconnected with a login prompt.
+    onMsaCode(data, client)
   }
 
   export class Relay extends Server {
diff --git a/src/relay.js b/src/relay.js
index 8f8ba1b..385ebf3 100644
--- a/src/relay.js
+++ b/src/relay.js
@@ -159,7 +159,7 @@ class Relay extends Server {
   constructor (options) {
     super(options)
     this.RelayPlayer = options.relayPlayer || RelayPlayer
-    this.forceSingle = true
+    this.forceSingle = options.forceSingle
     this.upstreams = new Map()
     this.conLog = debug
     this.enableChunkCaching = options.enableChunkCaching
@@ -175,12 +175,18 @@ class Relay extends Server {
     const options = {
       authTitle: this.options.authTitle,
       offline: this.options.destination.offline ?? this.options.offline,
-      username: this.options.offline ? ds.profile.name : null,
+      username: this.options.offline ? ds.profile.name : ds.profile.xuid,
       version: this.options.version,
       realms: this.options.destination.realms,
       host: this.options.destination.host,
       port: this.options.destination.port,
-      onMsaCode: this.options.onMsaCode,
+      onMsaCode: (code) => {
+        if (this.options.onMsaCode) {
+          this.options.onMsaCode(code, ds)
+        } else {
+          ds.disconnect("It's your first time joining. Please sign in and reconnect to join this server:\n\n" + code.message)
+        }
+      },
       profilesFolder: this.options.profilesFolder,
       backend: this.options.backend,
       autoInitPlayer: false
@@ -213,6 +219,16 @@ class Relay extends Server {
 
       this.emit('join', /* client connected to proxy */ ds, /* backend server */ client)
     })
+    client.on('error', (err) => {
+      ds.disconnect('Server error: ' + err.message)
+      debug(clientAddr, 'was disconnected because of error', err)
+      this.upstreams.delete(clientAddr.hash)
+    })
+    client.on('close', (reason) => {
+      ds.disconnect('Backend server closed connection')
+      this.upstreams.delete(clientAddr.hash)
+    })
+
     this.upstreams.set(clientAddr.hash, client)
   }
 
@@ -225,7 +241,7 @@ class Relay extends Server {
     this.conLog('closed upstream connection', clientAddr)
   }
 
-  // Called when a new player connects to our proxy server. Once the player has authenticted,
+  // Called when a new player connects to our proxy server. Once the player has authenticated,
   // we can open an upstream connection to the backend server.
   onOpenConnection = (conn) => {
     if (this.forceSingle && this.clientCount > 0) {
@@ -240,6 +256,11 @@ class Relay extends Server {
       player.on('login', () => {
         this.openUpstreamConnection(player, conn.address)
       })
+      player.on('close', (reason) => {
+        this.conLog('player disconnected', conn.address, reason)
+        this.clientCount--
+        delete this.clients[conn.address]
+      })
     }
   }