Nodejs ws module: Heartbeat in TypeScript

Question

I'm currently working on getting a WebSocket server running on NodeJs with TypeScript. As a WebSocket server implementation, I am using ws. Along with that I use the @types/ws package for the typings. I want the server to periodically send out a ping package to every connected client. The readme file of the ws module already provides an example implementation in JavaScript which I would like to use aswell. You can find the example implementation here:

The problem I am facing is that the sample implementation adds the "isAlive" attribute to a socket object which I obviously can not do in TypeScript since the "isAlive" attribute does not exist for the socket class.

My question: What is the best way to add this attribute to the socket class? Since I pass these socket objects around as parameters, I want to avoid adding some sort of extra import to all my files that somehow work with a socket.

I want to avoid typecasting the socket object to any for this.

I try to implement a wrapper class around ws, so that I can just instanciate my WebSocket class and use it throughout my application.

My class looks like this

import * as ws from "ws";


export default class WebSocketServer {
  private wss: ws.Server;
  private onSubscription: (socket: ws, deviceId: string, topicFilter: string) => void;

  constructor(port: number, onSubscription: (socket: ws, deviceId: string, topicFilter: string) => void) {
    this.wss = new ws.Server({ port });
    this.onSubscription = onSubscription;
    const conhandler = this.handleConnection.bind(this);
    this.wss.on("connection", conhandler);
  }

  public static init(onSubscription: (socket: ws, deviceId: string, topicFilter: string) => void): WebSocketServer {
    const port: number = Number(process.env.WSS_PORT);
    if (!port) {
      throw new Error("Unable to create the WebSocket server: Environment variable \"WSS_PORT\" is missing");
    }

    const server: WebSocketServer = new WebSocketServer(port, onSubscription);

    console.info(`Websocket Server listening on port ${port}`);
    return server;
  }


  private handleConnection(socket: ws): void {
    const messageHandler: (socket: ws, data: ws.Data) => void = this.handleMessage.bind(this);
    socket.on("message", data => messageHandler(socket, data));

    /* This is where I try to add the isAlive attribute to a socket object */
    /* but since ws does not have an attribute isAlive, TypeScript complains */
    socket.on("pong", () => socket.isAlive = true);
  }


  private handleMessage(socket: ws, data: ws.Data): void {
    /* Do stuff with the message */
  }


  public quit() {
    this.wss.close();
  }
}

Edit

As mentioned by a commenter, i tried declaration merging. I created a new file called "Augment.ts" and added the following code:

import * as ws from "ws";
declare module "ws" {
  interface WebSocket {
    isAlive: boolean;
  }
}

In the file that contains my WebSocketServer class, i added the import for the Augment.ts file like so: import "./Augment". The result is another error message (but somehow stating the same?): Property 'isAlive' does not exist on type 'import("c:/path/to/project/node_modules/@types/ws/index.d.ts")'.ts(2339)

Answer 1

Using this answer to a similar question, I found how to add a property to a WebSocket in TypeScript.

Method 1: Single TypeScript file

declare interface Socket extends WebSocket { // you can name Socket whatever you want
    isAlive: boolean;
}

Place this in the same file as your WebSocket above your new Server/new WebSocketServer.

Example:

import { Server, WebSocket } from "ws";

declare interface Socket extends WebSocket { // you can name Socket whatever you want
    isAlive: boolean;
}

const wss = new Server({ port: 8080 });

wss.on("connection", (ws:Socket) => { // ws is defined as a Socket here
    ws.isAlive = true;
    ws.on("pong", () => { ws.isAlive = true; });
});

// ect

Method 2: Multiple TypeScript files

Place this code block inside a new TypeScript file:

import { WebSocket } from "ws";

export interface Socket extends WebSocket {
    isAlive: boolean;
}

Note the export instead of declare.

Inside the files you want to import the definition, import Socket:

import { Socket } from "./example";

Then, inside your wss.on, use ws:Socket instead of ws/ws:WebSocket

Example:

import { Server } from "ws";
import { Socket } from "./example";

const wss = new Server({ port: 8080 });

wss.on("connection", (ws:Socket) => { // ws is defined as a Socket here
    ws.isAlive = true;
    ws.on("pong", () => { ws.isAlive = true; });
});

// ect

Using forEach

To make forEach work, use this:

(wss.clients as Set<Socket>).forEach((ws:Socket) => {
    // your code here
});

Happy coding!

Answer 2

I got it to work by creating a Custom Socket interface that extends Web Socket.

interface CustomSocket extends WebSocket {
    isAlive: boolean
}
export default class WebSocket {
    private heartbeat: ReturnType<typeof setInterval>
    private wss: WebSocket.Server


    constructor(port: number) {
        this.wss = new WebSocket.Server({port});

        this.wss.on('connection', this.handleConnection)
    }

    private handleConnection(socket: CustomSocket) {
        socket.on('pong', () => socket.isAlive = true)
    }
}