Keep the per-user index in sync on reauthentication

lib/listener.js

@@ -441,26 +441,37 @@ internals.Sockets.prototype.add = function (socket) {
 };
 
 
-internals.Sockets.prototype.auth = function (socket) {
+internals.Sockets.prototype.auth = function (socket, previous, credentials) {
+
+    // Called before the socket credentials are replaced
+
+    if (!this._listener._settings.auth.index ||
+        !this._items[socket.id]) {                              // Closed while authenticating
 
-    if (!this._listener._settings.auth.index) {
         return;
     }
 
-    if (!socket.auth.credentials.user) {
-        return;
+    const user = credentials.user;
+    if (user === previous) {
+        return;                                                 // Same user as before, index unchanged
     }
 
-    const user = socket.auth.credentials.user;
-    if (this._listener._settings.auth.maxConnectionsPerUser &&
+    if (user &&
+        this._listener._settings.auth.maxConnectionsPerUser &&
         this._byUser[user] &&
         this._byUser[user].length >= this._listener._settings.auth.maxConnectionsPerUser) {
 
         throw Boom.serverUnavailable('Too many connections for the authenticated user');
     }
 
-    this._byUser[user] = this._byUser[user] ?? [];
-    this._byUser[user].push(socket);
+    if (previous) {
+        this._deindex(socket, previous);
+    }
+
+    if (user) {
+        this._byUser[user] = this._byUser[user] ?? [];
+        this._byUser[user].push(socket);
+    }
 };
 
 
@@ -471,12 +482,17 @@ internals.Sockets.prototype.remove = function (socket) {
     if (socket.auth.credentials &&
         socket.auth.credentials.user) {
 
-        const user = socket.auth.credentials.user;
-        if (this._byUser[user]) {
-            this._byUser[user] = this._byUser[user].filter((item) => item !== socket);
-            if (!this._byUser[user].length) {
-                delete this._byUser[user];
-            }
+        this._deindex(socket, socket.auth.credentials.user);
+    }
+};
+
+
+internals.Sockets.prototype._deindex = function (socket, user) {
+
+    if (this._byUser[user]) {
+        this._byUser[user] = this._byUser[user].filter((item) => item !== socket);
+        if (!this._byUser[user].length) {
+            delete this._byUser[user];
         }
     }
 };
@@ -582,15 +598,52 @@ internals.Subscribers.prototype._remove = async function (socket, path, params)
 };
 
 
+internals.Subscribers.prototype.reauth = function (socket, previous) {
+
+    // Called after the socket credentials are replaced
+
+    if (!this._settings.auth ||
+        !this._settings.auth.index) {
+
+        return;
+    }
+
+    const user = socket.auth.credentials.user;
+    if (user === previous) {
+        return;
+    }
+
+    const item = this._items[socket.id];
+    if (!item) {
+        return;                                                 // Removed while the socket was closing
+    }
+
+    if (previous) {
+        this._deindex(item, previous);
+    }
+
+    if (user) {
+        this._byUser[user] = this._byUser[user] ?? [];
+        this._byUser[user].push(item);
+    }
+};
+
+
 internals.Subscribers.prototype._cleanup = function (socket, item) {
 
     delete this._items[socket.id];
 
     if (socket.auth.credentials &&
-        socket.auth.credentials.user &&
-        this._byUser[socket.auth.credentials.user]) {
+        socket.auth.credentials.user) {
+
+        this._deindex(item, socket.auth.credentials.user);
+    }
+};
+
+
+internals.Subscribers.prototype._deindex = function (item, user) {
 
-        const user = socket.auth.credentials.user;
+    if (this._byUser[user]) {
         this._byUser[user] = this._byUser[user].filter((record) => record !== item);
         if (!this._byUser[user].length) {
             delete this._byUser[user];

lib/socket.js

@@ -598,14 +598,19 @@ internals.Socket.prototype._authByCookie = async function () {
 
 internals.Socket.prototype._setCredentials = function (auth) {
 
+    const previous = this.auth.credentials ? this.auth.credentials.user : undefined;
+    this._listener._sockets.auth(this, previous, auth.credentials);     // Throws before any state changes when the user limit is reached
+
     this.auth.isAuthenticated = true;
     this.auth.credentials = auth.credentials;
     this.auth.artifacts = auth.artifacts;
     this.auth.strategy = auth.strategy;
 
     this.auth._verified = Date.now();
 
-    this._listener._sockets.auth(this);
+    for (const subscribers of new Set(Object.values(this._subscriptions))) {     // Paths on the same route share one subscribers object
+        subscribers.reauth(this, previous);
+    }
 };
 
 

test/auth.js

@@ -582,6 +582,196 @@ describe('authentication', () => {
             await server.stop();
         });
 
+        it('does not count reauthentication towards the connections limit', async () => {
+
+            const server = Hapi.server();
+
+            server.auth.scheme('custom', internals.implementation);
+            server.auth.strategy('default', 'custom');
+            server.auth.default('default');
+
+            await server.register({ plugin: Nes, options: { auth: { index: true, maxConnectionsPerUser: 2 } } });
+            await server.start();
+
+            const client = new Nes.Client(getUri(server.info));
+            await client.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+
+            await client.reauthenticate({ headers: { authorization: 'Custom john' } });
+            // a reauth is not a new connection
+            await expect(client.reauthenticate({ headers: { authorization: 'Custom john' } })).to.not.reject();
+
+            client.disconnect();
+            await server.stop();
+        });
+
+        it('moves the socket to the new user when reauthenticating as a different user', async () => {
+
+            const server = Hapi.server();
+
+            server.auth.scheme('custom', internals.implementation);
+            server.auth.strategy('default', 'custom');
+            server.auth.default('default');
+
+            await server.register({ plugin: Nes, options: { auth: { index: true, maxConnectionsPerUser: 1 } } });
+            await server.start();
+
+            const client = new Nes.Client(getUri(server.info));
+            const updates = [];
+            client.onUpdate = (update) => updates.push(update);
+            await client.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+            await client.reauthenticate({ headers: { authorization: 'Custom ed' } });
+
+            server.broadcast('for john', { user: 'john' });
+            server.broadcast('for ed', { user: 'ed' });
+            await Hoek.wait(50);
+            expect(updates).to.equal(['for ed']);
+
+            // john's only slot is free again while the switched socket is still open
+            const client2 = new Nes.Client(getUri(server.info));
+            await client2.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+
+            client.disconnect();
+            client2.disconnect();
+            await Hoek.wait(50);
+            expect(server.plugins.nes._listener._sockets._byUser).to.equal({});
+            await server.stop();
+        });
+
+        it('rejects reauthentication as a different user who is at the connections limit', async () => {
+
+            const server = Hapi.server();
+
+            server.auth.scheme('custom', internals.implementation);
+            server.auth.strategy('default', 'custom');
+            server.auth.default('default');
+
+            await server.register({ plugin: Nes, options: { auth: { index: true, maxConnectionsPerUser: 1 } } });
+            await server.start();
+
+            const edClient = new Nes.Client(getUri(server.info));
+            const edUpdates = [];
+            edClient.onUpdate = (update) => edUpdates.push(update);
+            await edClient.connect({ reconnect: false, auth: { headers: { authorization: 'Custom ed' } } });
+
+            const client = new Nes.Client(getUri(server.info));
+            await client.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+
+            const team = new Teamwork.Team();
+            client.onDisconnect = () => team.attend();
+
+            await expect(client.reauthenticate({ headers: { authorization: 'Custom ed' } })).to.reject('Too many connections for the authenticated user');
+            await team.work;
+            await Hoek.wait(50);
+
+            // The rejected socket stayed indexed under john until it closed, so john's slot is free again
+            const client2 = new Nes.Client(getUri(server.info));
+            await client2.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+
+            server.broadcast('for ed', { user: 'ed' });
+            await Hoek.wait(50);
+            expect(edUpdates).to.equal(['for ed']);
+
+            edClient.disconnect();
+            client2.disconnect();
+            await Hoek.wait(50);
+            expect(server.plugins.nes._listener._sockets._byUser).to.equal({});
+            await server.stop();
+        });
+
+        it('does not index a socket whose reauthentication completes after it closed', async () => {
+
+            const server = Hapi.server();
+
+            const authEntered = new Teamwork.Team();
+            const authGate = new Teamwork.Team();
+            const unsubEntered = new Teamwork.Team();
+            const unsubGate = new Teamwork.Team();
+
+            server.auth.scheme('custom', () => ({
+                authenticate: async (request, h) => {
+
+                    const user = request.headers.authorization.split(' ')[1];
+                    if (user === 'ed') {
+                        authEntered.attend();
+                        await authGate.work;
+                    }
+
+                    return h.authenticated({ credentials: { user } });
+                }
+            }));
+
+            server.auth.strategy('default', 'custom');
+            server.auth.default('default');
+
+            await server.register({ plugin: Nes, options: { auth: { index: true } } });
+            server.subscription('/item', {
+                auth: { index: true },
+                onUnsubscribe: async () => {
+
+                    unsubEntered.attend();
+                    await unsubGate.work;
+                }
+            });
+
+            await server.start();
+
+            const client = new Nes.Client(getUri(server.info));
+            await client.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+            await client.subscribe('/item', Hoek.ignore);
+
+            const listener = server.plugins.nes._listener;
+            const subscribers = Object.values(listener._sockets._items)[0]._subscriptions['/item'];
+
+            // Reauth is stuck in the auth scheme while the connection closes
+            const reauth = client.reauthenticate({ headers: { authorization: 'Custom ed' } }).catch(Hoek.ignore);
+            await authEntered.work;
+            client.disconnect();
+
+            // The close handler already removed the socket and is now awaiting onUnsubscribe
+            await unsubEntered.work;
+            authGate.attend();
+            await reauth;
+            await Hoek.wait(50);
+
+            unsubGate.attend();
+            await Hoek.wait(50);
+
+            expect(listener._sockets._byUser).to.equal({});
+            expect(subscribers._byUser).to.equal({});
+            await server.stop();
+        });
+
+        it('updates the user index when reauthenticating to and from credentials without a user', async () => {
+
+            const server = Hapi.server();
+
+            server.auth.scheme('custom', internals.implementation);
+            server.auth.strategy('default', 'custom');
+            server.auth.default('default');
+
+            await server.register({ plugin: Nes, options: { auth: { index: true } } });
+            await server.start();
+
+            const client = new Nes.Client(getUri(server.info));
+            const updates = [];
+            client.onUpdate = (update) => updates.push(update);
+            await client.connect({ reconnect: false, auth: { headers: { authorization: 'Custom john' } } });
+
+            await client.reauthenticate({ headers: { authorization: 'Custom app' } });
+            server.broadcast('for john', { user: 'john' });
+            await Hoek.wait(50);
+            expect(updates).to.equal([]);
+            expect(server.plugins.nes._listener._sockets._byUser).to.equal({});
+
+            await client.reauthenticate({ headers: { authorization: 'Custom john' } });
+            server.broadcast('for john again', { user: 'john' });
+            await Hoek.wait(50);
+            expect(updates).to.equal(['for john again']);
+
+            client.disconnect();
+            await server.stop();
+        });
+
         it('protects an endpoint with prefix', async () => {
 
             const server = Hapi.server();