25866259812

Committed 14 May 2026 02:38PM UTC coverage: 68.383% (+0.02%) from 68.367%

Build # 25866259812

Build Type

Pull #775

github

Committed by

web-flow

Commit Message

Merge 1d9353c4b into aee30e868

Pull Request Pull Request #775: Emit a `user:join` on reconnection if we don't have a lost connection

Coverage Stats

685 of 1169 branches covered (58.6%)

Branch coverage included in aggregate %.

4 of 4 new or added lines in 1 file covered. (100.0%)

8 existing lines in 1 file now uncovered.

2025 of 2794 relevant lines covered (72.48%)

173.93 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

71.19

/src/SocketServer.js

import { promisify } from 'node:util';
import lodash from 'lodash';
import { WebSocketServer } from 'ws';
import { Ajv } from 'ajv';
import { stdSerializers } from 'pino';
import { ulid, encodeTime } from 'ulid';
import { socketVote } from './controllers/booth.js';
import { disconnectUser } from './controllers/users.js';
import AuthRegistry from './AuthRegistry.js';
import GuestConnection from './sockets/GuestConnection.js';
import AuthedConnection from './sockets/AuthedConnection.js';
import LostConnection from './sockets/LostConnection.js';
import { serializeUser } from './utils/serialize.js';
import { jsonb } from './utils/sqlite.js';
import { subMinutes } from './utils/date.js';
import { ChatMutedError } from './errors/index.js';

const { isEmpty } = lodash;

export const KEY_ACTIVE_SESSIONS = 'users';

const PING_INTERVAL = 10_000;
const GUEST_COUNT_INTERVAL = 2_000;

/**
 * @typedef {import('./schema.js').User} User
 */

/**
 * @typedef {GuestConnection | AuthedConnection | LostConnection} Connection
 */

/**
 * @typedef {object} ClientActionParameters
 * @prop {string} sendChat
 * @prop {-1 | 1} vote
 * @prop {undefined} logout
 */

/**
 * @typedef {{
 *   [Name in keyof ClientActionParameters]: (
 *     user: User,
 *     parameter: ClientActionParameters[Name],
 *     connection: AuthedConnection
 *   ) => void
 * }} ClientActions
 */

const ajv = new Ajv({
  coerceTypes: false,
  ownProperties: true,
  removeAdditional: true,
  useDefaults: false,
});

/**
 * @template {object} T
 * @param {T} object
 * @param {PropertyKey} property
 * @returns {property is keyof T}
 */
function has(object, property) {
  return Object.prototype.hasOwnProperty.call(object, property);
}

class SocketServer {
  /**
   * @param {import('./Uwave.js').Boot} uw
   * @param {{ secret: Buffer|string, sessionStore: import('express-session').Store }} options
   */
  static async plugin(uw, options) {
    uw.socketServer = new SocketServer(uw, {
      secret: options.secret,
      sessionStore: options.sessionStore,
      server: uw.server,
    });

    uw.after(async () => {
      try {
        await uw.socketServer.initLostConnections();
      } catch (err) {
        // No need to prevent startup for this.
        // We do need to clear the `users` list because the lost connection handlers
        // will not do so.
        uw.socketServer.#logger.warn({ err }, 'could not initialise lost connections');
        await uw.keyv.delete(KEY_ACTIVE_SESSIONS);
      }
    });

    uw.onClose(async () => {
      await uw.socketServer.destroy();
    });
  }

  #uw;

  #logger;

  #wss;

  #closing = false;

  /** @type {Connection[]} */
  #connections = [];

  #pinger;

  /** Update online guests count and broadcast an update if necessary. */
  #guestCountInterval;

  #guestCountDirty = true;

  /**
   * Handlers for commands that come in from clients.
   *
   * @type {ClientActions}
   */
  #clientActions;

  /**
   * @type {{
   *   [K in keyof ClientActionParameters]:
   *     import('ajv').ValidateFunction<ClientActionParameters[K]>
   * }}
   */
  #clientActionSchemas;

  /**
   * Handlers for commands that come in from the server side.
   *
   * @type {import('./redisMessages.js').ServerActions}
   */
  #serverActions;

  #unsubscribe;

  /**
   * Create a socket server.
   *
   * @param {import('./Uwave.js').Boot} uw üWave Core instance.
   * @param {object} options Socket server options.
   * @param {number} [options.timeout] Time in seconds to wait for disconnected
   *     users to reconnect before removing them.
   * @param {Buffer|string} options.secret
   * @param {import('express-session').Store} options.sessionStore
   * @param {import('http').Server | import('https').Server} [options.server]
   * @param {number} [options.port]
   * @private
   */
  constructor(uw, options) {
    if (!uw) {
      throw new TypeError(`The "uw" argument must be of type UwaveServer. Received ${typeof uw}`);
    }

    this.#uw = uw;
    this.#logger = uw.logger.child({ ns: 'uwave:sockets' }, {
      serializers: {
        req: stdSerializers.req,
      },
    });

    this.options = {
      /** @type {(_socket: import('ws').WebSocket | undefined, err: Error) => void} */
      onError: (_socket, err) => {
        throw err;
      },
      timeout: 30,
      ...options,
    };

    // TODO put this behind a symbol, it's just public for tests
    this.authRegistry = new AuthRegistry(uw.db);

    this.#wss = new WebSocketServer({
      server: options.server,
      port: options.server ? undefined : options.port,
    });

    this.#unsubscribe = uw.events.onAny((command, data) => {
      this.#onServerMessage(command, data);
    });

    this.#wss.on('error', (error) => {
      this.onError(error);
    });
    this.#wss.on('connection', (socket, request) => {
      this.onSocketConnected(socket, request);
    });

    this.#pinger = setInterval(() => {
      this.ping();
    }, PING_INTERVAL);

    this.#guestCountInterval = setInterval(() => {
      if (!this.#guestCountDirty) {
        return;
      }

      this.#recountGuests();
    }, GUEST_COUNT_INTERVAL);

    this.#clientActions = {
      sendChat: (user, message) => {
        this.#logger.trace({ user, message }, 'sendChat');
        this.#uw.chat.send(user, message).catch((err) => {
          if (err instanceof ChatMutedError) {
            return;
          }

          this.#logger.warn({ err }, 'could not send chat message');
          // TODO: should this return the error to the client?
        });
      },
      vote: (user, direction) => {
        socketVote(this.#uw, user.id, direction);
      },
      logout: (user, _, connection) => {
        this.replace(connection, this.createGuestConnection(connection.socket));
        if (!this.connection(user)) {
          disconnectUser(this.#uw, user.id);
        }
      },
    };

    this.#clientActionSchemas = {
      sendChat: ajv.compile({
        type: 'string',
      }),
      vote: ajv.compile({
        type: 'integer',
        enum: [-1, 1],
      }),
      logout: ajv.compile(true),
    };

    this.#serverActions = {
      /**
       * Broadcast the next track.
       */
      'advance:complete': (next) => {
        if (next) {
          this.broadcast('advance', {
            historyID: next.historyID,
            userID: next.userID,
            media: next.media,
            playedAt: new Date(next.playedAt).getTime(),
          });
        } else {
          this.broadcast('advance', null);
        }
      },
      /**
       * Broadcast a skip notification.
       */
      'booth:skip': ({ moderatorID, userID, reason }) => {
        this.broadcast('skip', { moderatorID, userID, reason });
      },
      /**
       * Broadcast a chat message.
       */
      'chat:message': (message) => {
        this.broadcast('chatMessage', message);
      },
      /**
       * Delete chat messages. The delete filter can have an _id property to
       * delete a specific message, a userID property to delete messages by a
       * user, or be empty to delete all messages.
       */
      'chat:delete': ({ moderatorID, filter }) => {
        if ('id' in filter) {
          this.broadcast('chatDeleteByID', {
            moderatorID,
            _id: filter.id,
          });
        } else if ('userID' in filter) {
          this.broadcast('chatDeleteByUser', {
            moderatorID,
            userID: filter.userID,
          });
        } else if (isEmpty(filter)) {
          this.broadcast('chatDelete', { moderatorID });
        }
      },
      /**
       * Broadcast that a user was muted in chat.
       */
      'chat:mute': ({ moderatorID, userID, duration }) => {
        this.broadcast('chatMute', {
          userID,
          moderatorID,
          expiresAt: Date.now() + duration,
        });
      },
      /**
       * Broadcast that a user was unmuted in chat.
       */
      'chat:unmute': ({ moderatorID, userID }) => {
        this.broadcast('chatUnmute', { userID, moderatorID });
      },
      /**
       * Broadcast a vote for the current track.
       */
      'booth:vote': ({ userID, direction }) => {
        this.broadcast('vote', {
          _id: userID,
          value: direction,
        });
      },
      /**
       * Broadcast a favorite for the current track.
       */
      'booth:favorite': ({ userID }) => {
        this.broadcast('favorite', { userID });
      },
      /**
       * Cycle a single user's playlist.
       */
      'playlist:cycle': ({ userID, playlistID }) => {
        this.sendTo(userID, 'playlistCycle', { playlistID });
      },
      /**
       * Broadcast that a user joined the waitlist.
       */
      'waitlist:join': ({ userID, waitlist }) => {
        this.broadcast('waitlistJoin', { userID, waitlist });
      },
      /**
       * Broadcast that a user left the waitlist.
       */
      'waitlist:leave': ({ userID, waitlist }) => {
        this.broadcast('waitlistLeave', { userID, waitlist });
      },
      /**
       * Broadcast that a user was added to the waitlist.
       */
      'waitlist:add': ({
        userID, moderatorID, position, waitlist,
      }) => {
        this.broadcast('waitlistAdd', {
          userID, moderatorID, position, waitlist,
        });
      },
      /**
       * Broadcast that a user was removed from the waitlist.
       */
      'waitlist:remove': ({ userID, moderatorID, waitlist }) => {
        this.broadcast('waitlistRemove', { userID, moderatorID, waitlist });
      },
      /**
       * Broadcast that a user was moved in the waitlist.
       */
      'waitlist:move': ({
        userID, moderatorID, position, waitlist,
      }) => {
        this.broadcast('waitlistMove', {
          userID, moderatorID, position, waitlist,
        });
      },
      /**
       * Broadcast a waitlist update.
       */
      'waitlist:update': (waitlist) => {
        this.broadcast('waitlistUpdate', waitlist);
      },
      /**
       * Broadcast that the waitlist was cleared.
       */
      'waitlist:clear': ({ moderatorID }) => {
        this.broadcast('waitlistClear', { moderatorID });
      },
      /**
       * Broadcast that the waitlist was locked.
       */
      'waitlist:lock': ({ moderatorID, locked }) => {
        this.broadcast('waitlistLock', { moderatorID, locked });
      },

      'acl:allow': ({ userID, roles }) => {
        this.broadcast('acl:allow', { userID, roles });
      },
      'acl:disallow': ({ userID, roles }) => {
        this.broadcast('acl:disallow', { userID, roles });
      },

      'user:update': ({ userID, moderatorID, new: update }) => {
        // TODO Remove this remnant of the old roles system
        if ('role' in update) {
          this.broadcast('roleChange', {
            moderatorID,
            userID,
            role: update.role,
          });
        }
        if ('username' in update) {
          this.broadcast('nameChange', {
            moderatorID,
            userID,
            username: update.username,
          });
        }
      },
      'user:join': async ({ userID }) => {
        const { users, keyv } = this.#uw;
        const user = await users.getUser(userID);
        if (user) {
          // TODO this should not be the socket server code's responsibility
          const userIDs = /** @type {import('./schema.js').UserID[] | null} */ (
            await keyv.get(KEY_ACTIVE_SESSIONS)
          ) ?? [];
          userIDs.push(user.id);
          await keyv.set(KEY_ACTIVE_SESSIONS, userIDs);
          this.broadcast('join', serializeUser(user));
        }
      },
      /**
       * Broadcast that a user left the server.
       */
      'user:leave': ({ userID }) => {
        this.broadcast('leave', userID);
      },
      /**
       * Broadcast a ban event.
       */
      'user:ban': ({
        moderatorID, userID, permanent = false, duration, expiresAt,
      }) => {
        this.broadcast('ban', {
          moderatorID, userID, permanent, duration, expiresAt,
        });

        this.#connections.forEach((connection) => {
          if (connection instanceof AuthedConnection && connection.user.id === userID) {
            connection.ban();
          } else if (connection instanceof LostConnection && connection.user.id === userID) {
            connection.close();
          }
        });
      },
      /**
       * Broadcast an unban event.
       */
      'user:unban': ({ moderatorID, userID }) => {
        this.broadcast('unban', { moderatorID, userID });
      },
      /**
       * Force-close a connection.
       */
      'http-api:socket:close': (userID) => {
        this.#connections.forEach((connection) => {
          if ('user' in connection && connection.user.id === userID) {
            connection.close();
          }
        });
      },

      'emotes:reload': () => {
        this.broadcast('reloadEmotes', null);
      },
    };
  }

  /**
   * Create `LostConnection`s for every user that's known to be online, but that
   * is not currently connected to the socket server.
   *
   * @private
   */
  async initLostConnections() {
    const { db, keyv } = this.#uw;
    const userIDs = /** @type {import('./schema.js').UserID[] | null} */ (
      await keyv.get(KEY_ACTIVE_SESSIONS)
    ) ?? [];
    const disconnectedIDs = userIDs.filter((userID) => !this.connection(userID));

    if (disconnectedIDs.length === 0) {
      return;
    }

    const disconnectedUsers = await db.selectFrom('users')
      .where('id', 'in', disconnectedIDs)
      .selectAll()
      .execute();
    disconnectedUsers.forEach((_user) => {
      // TODO (commented out but it already didn't really work)
      void _user;
      // this.add(this.createLostConnection(user, 'TODO: Actual session ID!!', null));
    });
  }

  /**
   * @param {import('ws').WebSocket} socket
   * @param {import('http').IncomingMessage} request
   * @private
   */
  onSocketConnected(socket, request) {
    this.#logger.info({ req: request }, 'new connection');

    socket.on('error', (error) => {
      this.onSocketError(socket, error);
    });
    this.add(this.createGuestConnection(socket));
  }

  /**
   * @param {import('ws').WebSocket} socket
   * @param {Error} error
   * @private
   */
  onSocketError(socket, error) {
    this.#logger.warn({ err: error }, 'socket error');

    this.options.onError(socket, error);
  }

  /**
   * @param {Error} error
   * @private
   */
  onError(error) {
    this.#logger.error({ err: error }, 'server error');

    this.options.onError(undefined, error);
  }

  /**
   * Get a LostConnection for a user, if one exists.
   *
   * @param {string} sessionID
   * @private
   */
  getLostConnection(sessionID) {
    return this.#connections.find((connection) => (
      connection instanceof LostConnection && connection.sessionID === sessionID
    ));
  }

  /**
   * Create a connection instance for an unauthenticated user.
   *
   * @param {import('ws').WebSocket} socket
   * @private
   */
  createGuestConnection(socket) {
    const connection = new GuestConnection(
      this.#uw,
      this.options.sessionStore,
      socket,
      { authRegistry: this.authRegistry },
    );
    connection.on('close', () => {
      this.remove(connection);
    });
    connection.on('authenticate', async ({ user, sessionID, lastEventID }) => {
      let isReconnect = await connection.isReconnect(sessionID);
      this.#logger.info({ userId: user.id, isReconnect, lastEventID }, 'authenticated socket');
      if (isReconnect) {
        const previousConnection = this.getLostConnection(sessionID);
        if (previousConnection) {
          this.remove(previousConnection);
        } else {
          // If there's actually no lost connection, we might've derived the reconnection flag from stale state?
          // XXX(@goto-bus-stop): validate this
          isReconnect = false;
        }
      }

      this.replace(connection, this.createAuthedConnection(socket, user, sessionID, lastEventID));

      if (!isReconnect) {
        this.#uw.publish('user:join', { userID: user.id });
      }
    });
    return connection;
  }

  /**
   * Create a connection instance for an authenticated user.
   *
   * @param {import('ws').WebSocket} socket
   * @param {User} user
   * @param {string} sessionID
   * @param {string|null} lastEventID
   * @returns {AuthedConnection}
   * @private
   */
  createAuthedConnection(socket, user, sessionID, lastEventID) {
    const connection = new AuthedConnection(
      this.#uw,
      this.options.sessionStore,
      socket,
      user,
      sessionID,
      lastEventID,
    );
    connection.on('close', ({ banned, lastEventID }) => {
      if (banned) {
        this.#logger.info({ userId: user.id }, 'removing connection after ban');
        disconnectUser(this.#uw, user.id);
      } else if (!this.#closing) {
        this.#logger.info({ userId: user.id }, 'lost connection');
        this.add(this.createLostConnection(user, sessionID, lastEventID));
      }
      this.remove(connection);
    });
    connection.on(
      'command',
      ({ command, data }) => {
        this.#logger.trace({ userId: user.id, command, data }, 'command');
        if (has(this.#clientActions, command)) {
          // Ignore incorrect input
          const validate = this.#clientActionSchemas[command];
          if (validate && !validate(data)) {
            return;
          }

          const action = this.#clientActions[command];
          // @ts-expect-error TS2345 `data` is validated
          action(user, data, connection);
        }
      },
    );
    return connection;
  }

  /**
   * Create a connection instance for a user who disconnected.
   *
   * @param {User} user
   * @param {string} sessionID
   * @param {string|null} lastEventID
   * @returns {LostConnection}
   * @private
   */
  createLostConnection(user, sessionID, lastEventID) {
    const connection = new LostConnection(
      this.#uw,
      this.options.sessionStore,
      user,
      sessionID,
      lastEventID,
      this.options.timeout,
    );
    connection.on('close', () => {
      this.#logger.info({ userId: user.id }, 'user left');
      this.remove(connection);
      // Only register that the user left if they didn't have another connection
      // still open.
      if (!this.connection(user)) {
        disconnectUser(this.#uw, user.id);
      }
    });
    return connection;
  }

  /**
   * Add a connection.
   *
   * @param {Connection} connection
   * @private
   */
  add(connection) {
    const userID = 'user' in connection ? connection.user.id : null;
    const sessionID = 'sessionID' in connection ? connection.sessionID : null;
    this.#logger.trace({ type: connection.constructor.name, userID, sessionID }, 'add connection');

    this.#connections.push(connection);
    this.#guestCountDirty = true;
  }

  /**
   * Remove a connection.
   *
   * @param {Connection} connection
   * @private
   */
  remove(connection) {
    const userID = 'user' in connection ? connection.user.id : null;
    const sessionID = 'sessionID' in connection ? connection.sessionID : null;
    this.#logger.trace({ type: connection.constructor.name, userID, sessionID }, 'remove connection');

    const i = this.#connections.indexOf(connection);
    this.#connections.splice(i, 1);

    connection.removed();
    this.#guestCountDirty = true;
  }

  /**
   * Replace a connection instance with another connection instance. Useful when
   * a connection changes "type", like GuestConnection → AuthedConnection.
   *
   * @param {Connection} oldConnection
   * @param {Connection} newConnection
   * @private
   */
  replace(oldConnection, newConnection) {
    this.remove(oldConnection);
    this.add(newConnection);
  }

  /**
   * Handle command messages coming in from elsewhere in the app.
   * Some commands are intended to broadcast immediately to all connected
   * clients, but others require special action.
   *
   * @template {keyof import('./redisMessages.js').ServerActionParameters} K
   * @param {K} command
   * @param {import('./redisMessages.js').ServerActionParameters[K]} data
   */
  #onServerMessage(command, data) {
    this.#logger.trace({ channel: command, command, data }, 'server message');

    const action = this.#serverActions[command];
    if (action !== undefined) {
      action(data);
    }
  }

  /**
   * Stop the socket server.
   *
   * @returns {Promise<void>}
   */
  async destroy() {
    this.#closing = true;

    this.#unsubscribe();
    clearInterval(this.#pinger);
    clearInterval(this.#guestCountInterval);

    for (const connection of this.#connections) {
      connection.close();
    }

    const closeWsServer = promisify(this.#wss.close.bind(this.#wss));
    await closeWsServer();
  }

  /**
   * Get the connection instance for a specific user.
   *
   * @param {User|import('./schema.js').UserID} user The user.
   * @returns {Connection|undefined}
   */
  connection(user) {
    const userID = typeof user === 'object' ? user.id : user;
    return this.#connections.find((connection) => 'user' in connection && connection.user.id === userID);
  }

  ping() {
    this.#connections.forEach((connection) => {
      connection.ping();
    });

    this.#cleanupMessageQueue().catch((err) => {
      this.#logger.error({ err }, 'failed to clean up socket message queue');
    });
  }

  async #cleanupMessageQueue() {
    const oldestID = encodeTime(subMinutes(new Date(), 10).getTime());

    await this.#uw.db.deleteFrom('socketMessageQueue')
      .where('id', '<', oldestID)
      .execute();
  }

  /**
   * Broadcast a command to all connected clients.
   *
   * @param {string} command Command name.
   * @param {import('type-fest').JsonValue} data Command data.
   * @param {import('./schema.js').UserID | null} targetUserID
   */
  #recordMessage(command, data, targetUserID = null) {
    const id = ulid();

    this.#uw.db.insertInto('socketMessageQueue')
      .values({
        id,
        command,
        data: jsonb(data),
        targetUserID,
      })
      .execute();

    return id;
  }

  /**
   * Broadcast a command to all connected clients.
   *
   * @param {string} command Command name.
   * @param {import('type-fest').JsonValue} data Command data.
   */
  broadcast(command, data) {
    const id = this.#recordMessage(command, data);

    this.#logger.trace({
      id,
      command,
      data,
      to: this.#connections.map((connection) => (
        'user' in connection ? connection.user.id : null
      )),
    }, 'broadcast');

    this.#connections.forEach((connection) => {
      connection.send(id, command, data);
    });
  }

  /**
   * Send a command to a single user.
   *
   * @param {User|import('./schema.js').UserID} user User or user ID to send the command to.
   * @param {string} command Command name.
   * @param {import('type-fest').JsonValue} data Command data.
   */
  sendTo(user, command, data) {
    const userID = typeof user === 'object' ? user.id : user;
    const id = this.#recordMessage(command, data, userID);

    this.#connections.forEach((connection) => {
      if ('user' in connection && connection.user.id === userID) {
        connection.send(id, command, data);
      }
    });
  }

  #lastGuestCount = 0;

  /** The number of unauthenticated connections. */
  get guestCount() {
    return this.#connections.reduce((acc, connection) => {
      if (connection instanceof GuestConnection) {
        return acc + 1;
      }
      return acc;
    }, 0);
  }

  #recountGuests() {
    const guests = this.guestCount;
    if (guests !== this.#lastGuestCount) {
      this.#lastGuestCount = guests;
      this.broadcast('guests', guests);
    }
  }
}

export default SocketServer;

u-wave / core / 25866259812

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous