From 61b856c2c9a92c9e8f83a59fb13c38dd464446d6 Mon Sep 17 00:00:00 2001
From: Calvin Montgomery <calzoneman@gmail.com>
Date: Sun, 17 Jun 2018 16:25:40 -0700
Subject: [PATCH] Refactor socket.io middleware for future uws compatibility

---
 src/io/ioserver.js  | 39 +++++++++++++++++++++++++++------------
 test/io/ioserver.js |  7 +++----
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/io/ioserver.js b/src/io/ioserver.js
index d7227394..03095b69 100644
--- a/src/io/ioserver.js
+++ b/src/io/ioserver.js
@@ -34,6 +34,20 @@ const authFailureCount = new Counter({
     help: 'Number of failed authentications from session middleware'
 });
 
+class SocketIOContext {
+    constructor(socket) {
+        socket.handshake.connection = {
+            remoteAddress: socket.handshake.address
+        };
+
+        this.upgradeReq = socket.handshake;
+        this.ipAddress = null;
+        this.torConnection = null;
+        this.ipSessionFirstSeen = null;
+        this.user = null;
+    }
+}
+
 class IOServer {
     constructor(options = {
         proxyTrustFn: proxyaddr.compile('127.0.0.1')
@@ -49,21 +63,16 @@ class IOServer {
     // Map proxied sockets to the real IP address via X-Forwarded-For
     // If the resulting address is a known Tor exit, flag it as such
     ipProxyMiddleware(socket, next) {
-        if (!socket.context) socket.context = {};
-
         try {
-            socket.handshake.connection = {
-                remoteAddress: socket.handshake.address
-            };
-
             socket.context.ipAddress = proxyaddr(
-                socket.handshake,
+                socket.context.upgradeReq,
                 this.proxyTrustFn
             );
 
             if (!socket.context.ipAddress) {
                 throw new Error(
-                    `Assertion failed: unexpected IP ${socket.context.ipAddress}`
+                    'Could not determine IP address from ' +
+                    socket.context.upgradeReq.connection.remoteAddress
                 );
             }
         } catch (error) {
@@ -163,7 +172,7 @@ class IOServer {
 
     // Parse cookies
     cookieParsingMiddleware(socket, next) {
-        const req = socket.handshake;
+        const req = socket.context.upgradeReq;
         if (req.headers.cookie) {
             cookieParser(req, null, () => next());
         } else {
@@ -176,7 +185,7 @@ class IOServer {
     // Determine session age from ip-session cookie
     // (Used for restricting chat)
     ipSessionCookieMiddleware(socket, next) {
-        const cookie = socket.handshake.signedCookies['ip-session'];
+        const cookie = socket.context.upgradeReq.signedCookies['ip-session'];
         if (!cookie) {
             socket.context.ipSessionFirstSeen = new Date();
             next();
@@ -197,7 +206,7 @@ class IOServer {
         socket.context.aliases = [];
 
         const promises = [];
-        const auth = socket.handshake.signedCookies.auth;
+        const auth = socket.context.upgradeReq.signedCookies.auth;
         if (auth) {
             promises.push(verifySession(auth).then(user => {
                 socket.context.user = Object.assign({}, user);
@@ -245,6 +254,10 @@ class IOServer {
         patchTypecheckedFunctions();
 
         const io = this.io = sio.instance = sio();
+        io.use((socket, next) => {
+            socket.context = new SocketIOContext(socket);
+            next();
+        });
         io.use(this.ipProxyMiddleware.bind(this));
         io.use(this.ipBanMiddleware.bind(this));
         io.use(this.ipThrottleMiddleware.bind(this));
@@ -422,7 +435,9 @@ module.exports = {
         ioServer.bindTo(servers);
     },
 
-    IOServer: IOServer
+    IOServer: IOServer,
+
+    SocketIOContext: SocketIOContext
 };
 
 /* Clean out old rate limiters */
diff --git a/test/io/ioserver.js b/test/io/ioserver.js
index e343ecb7..2278e34d 100644
--- a/test/io/ioserver.js
+++ b/test/io/ioserver.js
@@ -1,5 +1,6 @@
 const assert = require('assert');
 const IOServer = require('../../lib/io/ioserver').IOServer;
+const SocketIOContext = require('../../lib/io/ioserver').SocketIOContext;
 
 describe('IOServer', () => {
     let server;
@@ -7,9 +8,6 @@ describe('IOServer', () => {
     beforeEach(() => {
         server = new IOServer();
         socket = {
-            context: {
-                ipAddress: '9.9.9.9'
-            },
             handshake: {
                 address: '127.0.0.1',
                 headers: {
@@ -17,6 +15,7 @@ describe('IOServer', () => {
                 }
             }
         };
+        socket.context = new SocketIOContext(socket);
     });
 
     describe('#ipProxyMiddleware', () => {
@@ -29,7 +28,7 @@ describe('IOServer', () => {
         });
 
         it('does not proxy from a non-trusted address', done => {
-            socket.handshake.address = '5.6.7.8';
+            socket.context.upgradeReq.connection.remoteAddress = '5.6.7.8';
             server.ipProxyMiddleware(socket, error => {
                 assert(!error);
                 assert.strictEqual(socket.context.ipAddress, '5.6.7.8');