Protecting against misbehaving clients with many connections

[toini]

Somewhat related to #489 that tries to avoid troubles with clients we have no control over: we've previously limited number of connections single user can open on a server by tracking their authorization tokens and blocking at onConnect if there already are 2 sockets for a token.

This has proven problematic as there is nothing that would recover the situation where server thinks there are two connections for a user. Im thinking if a better approach could be to track sockets by token but instead or blocking, just close the other sockets for same token when new connection is initialized.

Does this sound like a proper way to protect from excess connections? Any advise on how to implement it on uWs?

Br,
Toni

[enisdenjo]

All of this can be implemented in the onConnect and onClose hooks. A map of token with the sockets array should suffice. When a token appears, look up the map, and if the sockets array exceeds the size, simply socket.close() the last one. Straight forward.

You of course have to get the "number of allowed sockets" right, users open multiple tabs when working and can therefore establish multiple, valid, connections.

What's the thing you want to achieve, protect against excess connections from a single device, or across all devices?

[toini]

I was thinking that at the end of onConnect i would go through all existing clients and find the ones that have the same token and user agent and close them. (So i'd allow multiple connections but only for different user agents).

The token and user agent is stored in CustomExtra. But now I faced an issue with figuring out the CustomExtra from the Map<uWS.WebSocket<unknown>, Client>

type CustomExtra = Extra & {
    token?: string;
    userId?: string;
    userAgent?: string;
    messageCount?: number;
    subscribedHomes?: Set<string>;
};

type ExtendedWebSocketBehavior = uWS.WebSocketBehavior<unknown> & {
    getNumberOfClients: () => number;
    getClients: () => Map<uWS.WebSocket<unknown>, Client>;
};

const closeExcessConnectionsByTokenAndUserAgent = async (
    incomingSocket: uWS.WebSocket<unknown>,
    token: string,
    userAgent: string,
) => {
    const clients: Map<uWS.WebSocket<unknown>, Client> = getClients();
    const socketsToClose = Array.from(clients.keys()).filter((s) => {
        if (s === incomingSocket) return false; // Don't close the incoming socket

        // No it doesnt work like this as socket is a property of extra, not the other way around.....
        const socket = s as uWS.WebSocket<unknown> & CustomExtra;            
        return socket.token === token && socket.userAgent === userAgent;
    });

    for (const socket of socketsToClose) {
        // Close socket
    }
};

``

[toini]

What's the thing you want to achieve, protect against excess connections from a single device, or across all devices?

Im protecting against single user. They may run a device that can open very many connections. These are not browsers but home automation devices.

[toini]

A map of token with the sockets array should suffice

It would've maybe been nicer if i could enumerate the excess sockets to close just by using the single Map we already have. Adding another map means it needs properly cleaned up in all scenarios.

[toini]

Actually, since we are closing authenticated sockets, instead of token i can use userid. I.e. there would need to be a mapping userId + userAgent => socket.

[toini]

I added a class for keeping track of the connections:

import type uWS from 'uWebSockets.js';
import { logger } from '../logger.js';

// Note that currently we only allow one socket per userId, userIp and userAgent
// Thus we could just keep reference to a single socket instead of a set
export class SocketPruner {
    #map = new Map<string, Set<uWS.WebSocket<unknown>>>();

    static key(userId: string, userIp: string, userAgent: string) {
        return `${userId}|${userIp}|${userAgent}`;
    }

    #add(key: string, socket: uWS.WebSocket<unknown>) {
        let set = this.#map.get(key);
        if (!set) {
            set = new Set();
            this.#map.set(key, set);
        }
        set.add(socket);
    }

    #prune(key: string, closeFn: (socket: uWS.WebSocket<unknown>) => void) {
        const set = this.#map.get(key);
        if (!set) return;

        for (const socket of set) {
            closeFn(socket);
        }

        this.#map.delete(key); // Just delete the whole set
        logger.info('Sockets pruned', { key, count: set.size });
    }

    remove(key: string, socket: uWS.WebSocket<unknown>) {
        const set = this.#map.get(key);
        if (!set) return;
        set.delete(socket);
        if (set.size === 0) this.#map.delete(key);
    }

    pruneAdd(key: string, socket: uWS.WebSocket<unknown>, closeFn: (socket: uWS.WebSocket<unknown>) => void) {
        this.#prune(key, closeFn);
        this.#add(key, socket);
    }
}

Then

// onConnect
const key = SocketPruner.key(claims.sub, userIp, userAgent);
socketPruner.pruneAdd(key, context.extra.socket, (s) => s.end(CloseCode.TooManyInitialisationRequests, 'duplicate connection'));

// onClose
socketPruner.remove(key, context.extra.socket);

and if the sockets array exceeds the size, simply socket.close() the last one. Straight forward.

Btw if i call close() or end() on a socket, then all this does not get invoked and the socket is not removed from the clients list?

Also, im now keeping the sockets in memory with my set. I wonder if this may lead to memory leaks?

[toini]

How could i check if the socket is still open?

[enisdenjo]

You can use the onClose hook and remove the socket from your map. Alternatively you can just wrap the socket.end in a try/catch.

[toini]

I'm not sure i understand. Can I call onClose from onConnect when i need to close the old stale socket that the client apparently has discarded as they are opening a new connection? The socket.end() is currently in a try/catch.

[enisdenjo]

I meant this

// onClose
socketPruner.remove(key, context.extra.socket);

but you've already done it.

Then it's a race condition, using a try/catch is a good bet.

[toini]

There was an issue with retrieving the user ip during onClose (uWebsockets doesnt allow that for disconnected clients) and the key we used to clean up the socket didnt match. I switched to using userIp stored in context extra during onConnect. Let's see if that helps.