Private desktop and mobile WebSocket connections with Node.js

Question

Background info

To learn a new branch of my dev career, I've started a new project in my spare time (I've never used WebSockets before and I also want to learn more about common authentication flows). I want to run (deploy to Google's App Engine) this simple WebSocket server alongside my hybrid flutter firebase app. One of the features of the flutter app is that the authenticated users (Google Firebase Authentication for now) can stream real-time face detection data (facial expressions, head tilt angles, etc.) from their mobile phones (Firebase ML kit) to desktop (face-api.js) and vice versa. This data is then used to animate 2D characters (using Flare) which are covering and replacing the user's real face. I've chosen ws over socket.io because I'm not aware of any out of the box support for flutter and I think for my relatively simple use case the ws is sufficient.

Key features

• an authenticated user is able to connect to the web socket through mobile or web app
• after a successful connection, the server creates a private room with the user's latest connected mobile and desktop device
• client messages are broadcasted to latest opened mobile and web app only
• an overwritten connection gets notified (user opens a new tab/app on another device)

Implementation

Authenticated connection

I've started with the suggested pattern for client authentication inside the ws library and I've added my own auth logic.

// ...

server.on('upgrade', async function upgrade(request, socket, head) {
  const webIdToken = request.headers['sec-websocket-protocol'];
  const mobileIdToken = request.headers['clientid'];
  const userConnectionType = webIdToken ? 'web' : 'mobile';
  const idToken = webIdToken || mobileIdToken;

  if (!idToken) {
    socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
    socket.destroy();
    return;
  }

  try {
    const authenticatedUser = await admin.auth().verifyIdToken(idToken);
    
    wss.handleUpgrade(request, socket, head, function done(ws) {
      wss.emit('connection', ws, request, { authenticatedUser, userConnectionType });
    });
  } catch (error) {
    console.log(error);
    socket.write('HTTP/1.1 403 Invalid Token\r\n\r\n');
    socket.destroy();
    return;
  }
});

wss.on('connection', function connection(ws, request, client) {
  // ...
});

// ...

For the web app, I'm using the Sec-WebSocket-Protocol header (explained here) to send the JWT of the signed-in Firebase user. For mobile, I'm simply adding the token to the headers. This way I also know if the user is trying to connect from the web or the mobile app. I'm using firebase-admin to verify the token and then emit the 'connection' event with the decoded token (authenticatedUser) and the user's connection type (userConnectionType).

Creating rooms

// ...

const privateRooms = new Map();

// ...

wss.on('connection', function connection(ws, request, { authenticatedUser, userConnectionType }) {
  const userRoom = privateRooms.get(authenticatedUser.email);

  if (!userRoom) {
    const emptyRoom = { desktop: null, mobile: null };

    emptyRoom[userConnectionType] = ws;
    privateRooms.set(authenticatedUser.email, emptyRoom);
  } else {
    const previousWs = userRoom[userConnectionType];

    if (previousWs) {
      previousWs.send('This connection has been overwritten with a new one');
    }

    userRoom[userConnectionType] = ws;
    privateRooms.set(authenticatedUser.email, userRoom);
  }
});

To store active rooms I'm using a JS Map with the user's email being the key of the map entry. On the user's first connection I'm creating a new room, which is a simple object containing the latest mobile and desktop WebSockets: userRoom = { mobile: ws, desktop: ws}. If the user's room already exists I simply overwrite the previous ws with a new connection and notify the previous ws connection.

Clearing rooms

ws.on('close', function close() {
  const currentUserRoom = privateRooms.get(authenticatedUser.email);
  const isInRoom = (currentUserRoom[userConnectionType] === ws);

  if (!isInRoom) {
    return;
  }

  if (isInRoom && (currentUserRoom.desktop == null || currentUserRoom.mobile == null)) {
    privateRooms.delete(authenticatedUser.email);
    return;
  }

  if (isInRoom) {
    currentUserRoom[userConnectionType] = null;
    privateRooms.set(authenticatedUser.email, currentUserRoom);
  }
});

After the user closes a connection, I'm checking if the ws is in the room and deleting the room from the map if there's only one remaining connection before delete. If both desktop and mobile are connected I simply set the ws to null again.

Broadcasting private messages

ws.on('message', function incoming(clientMessage) {
  const currentUserRoom = privateRooms.get(authenticatedUser.email);

  wss.clients.forEach(function each(client) {
    const isClientInRoom = (currentUserRoom.desktop === client || currentUserRoom.mobile === client);

    if (isClientInRoom && client.readyState == WebSocket.OPEN) {
      client.send(`
        Desktop connected: ${currentUserRoom.desktop != null} 
        Mobile connected: ${currentUserRoom.mobile != null} 
        Room: ${authenticatedUser.email} 
        Message: ${clientMessage}
      `);
    }
  });
});

To broadcast a private message I cycle through every client and send a message to ones that are found in the connected user's room.

Full code

'use strict';

const http = require('http');
const WebSocket = require('ws');
const admin = require('firebase-admin');
const adminCredentials = require('./credentials.json');

const server = http.createServer();
const port = process.env.PORT || 3000;
const privateRooms = new Map();
const wss = new WebSocket.Server({ noServer: true });

admin.initializeApp({
  credential: admin.credential.cert({
    projectId: adminCredentials['project_id'],
    clientEmail: adminCredentials['client_email'],
    privateKey: adminCredentials['private_key']
  })
});

server.on('upgrade', async function upgrade(request, socket, head) {
  const webIdToken = request.headers['sec-websocket-protocol'];
  const mobileIdToken = request.headers['clientid'];
  const userConnectionType = webIdToken ? 'desktop' : 'mobile';
  const idToken = webIdToken || mobileIdToken;

  if (!idToken) {
    socket.write('HTTP/1.1 400 Bad Request\r\n\r\n');
    socket.destroy();
    return;
  }

  try {
    const authenticatedUser = await admin.auth().verifyIdToken(idToken);

    wss.handleUpgrade(request, socket, head, function done(ws) {
      wss.emit('connection', ws, request, { authenticatedUser, userConnectionType });
    });
  } catch (error) {
    console.log(error);
    socket.write('HTTP/1.1 403 Invalid Token\r\n\r\n');
    socket.destroy();
    return;
  }
});

wss.on('connection', function connection(ws, request, { authenticatedUser, userConnectionType }) {
  const userRoom = privateRooms.get(authenticatedUser.email);

  if (!userRoom) {
    const emptyRoom = { desktop: null, mobile: null };

    emptyRoom[userConnectionType] = ws;
    privateRooms.set(authenticatedUser.email, emptyRoom);
  } else {
    const previousWs = userRoom[userConnectionType];

    if (previousWs) {
      previousWs.send('This connection has been overwritten with a new one');
    }

    userRoom[userConnectionType] = ws;
    privateRooms.set(authenticatedUser.email, userRoom);
  }

  ws.on('close', function close() {
    const currentUserRoom = privateRooms.get(authenticatedUser.email);
    const isInRoom = (currentUserRoom[userConnectionType] === ws);

    if (!isInRoom) {
      return;
    }

    if (isInRoom && (currentUserRoom.desktop == null || currentUserRoom.mobile == null)) {
      privateRooms.delete(authenticatedUser.email);
      return;
    }

    if (isInRoom) {
      currentUserRoom[userConnectionType] = null;
      privateRooms.set(authenticatedUser.email, currentUserRoom);
    }
  });

  ws.on('message', function incoming(clientMessage) {
    const currentUserRoom = privateRooms.get(authenticatedUser.email);

    wss.clients.forEach(function each(client) {
      const isClientInRoom = (currentUserRoom.desktop === client || currentUserRoom.mobile === client);

      if (isClientInRoom && client.readyState == WebSocket.OPEN) {
        client.send(`
          Desktop connected: ${currentUserRoom.desktop != null} 
          Mobile connected: ${currentUserRoom.mobile != null} 
          Room: ${authenticatedUser.email} 
          Message: ${clientMessage}
        `);
      }
    });
  });
});

server.listen(port);
console.log(`listening on ${port}`);

Questions

Any improvements to the code are appreciated. I also have some specific questions regarding the code:

• Is the JS Map usage to represent all private rooms with the user's email being the key fine in this case?
• I'm using strings ('desktop' or 'mobile') to access the room object properties, is there a more elegant way to represent a room in this case?

Answer 1

The code you made is looking pretty good. I'm no expert in the particular tools you're using, but I'll still give my two cents about the parts I understand more.

Let's start with your upgrade listener. I see that when authentication headers are missing, you're giving back a 400 invalid-request error. Would a 401 unauthorized error be more appropriate? (this status code is normally used when a request is missing authentication credentials)

The logic in ws.on('message', ...) seems a little off to me. Why are we looping over all of the active connections in order to send a message to the two people in the private room, especially when we have a reference to the private room handy? Can't we just check if currentUserRoom.desktop and currentUserRoom.mobile exists, and if they do, send a message to them directly? From what I saw in the docs, those client instances that you're looping over are the exact same instances as what you're saving in currentUserRoom.dekstop/mobile.

---

My main issue with this script is that the logic within it is not being categorized and grouped together in any way. You're not grouping logic together with functions or namespaces (i.e. objects or modules), unless you count the callbacks you're required to use. firebase-specific logic, websocket-specific logic, and the logic behind private room management is all intermixed.

In your connection event handler, I would try to move out all the "business logic" related to managing rooms into its own file. This would make the logic behind private rooms more testable, flexable (you can easily add other ways for users to join a room besides websockets), and in general, makes the code more readable (it's much easier to understand that a private room is being created when you see "createRoom()" vs a list of room creating instructions like "privateRooms.set(...)").

Here's an example of what a file of just the business logic might look like

privateRooms.js

'use strict';

exports.initPrivateRooms = function () {
  const privateRooms = new Map();
  return {
    getOrCreateRoom(roomId) {
      if (!privateRooms.has(roomId)) {
        privateRooms.set(roomId, createRoom(roomId));
      }
      return privateRooms.get(roomId);
    },
    removeRoom(roomId) {
      privateRooms.delete(roomId);
    },
  };
};

function createRoom(id) {
  const deviceConnections = {};
  return Object.freeze({
    id,
    hasDeviceType(deviceType) {
      return !!deviceConnections[deviceType];
    },
    messageMembers(msg) {
      for (const connection of Object.values(deviceConnections)) {
        connection.onMessaged(msg);
      }
    },
    join(deviceType, eventListeners = {}) {
      const { onMessaged = () => {}, onKickedOut = () => {} } = eventListeners;

      if (deviceConnections[deviceType] != null) {
        deviceConnections[deviceType].onKickedOut();
      }
      deviceConnections[deviceType] = { onMessaged, onKickedOut };
    },
    leave(deviceType) {
      delete deviceConnections[deviceType];
    },
    isEmpty() {
      return Object.values(deviceConnections).length === 0;
    },
  });
}

exports.DEVICE_TYPE = {
  desktop: 'DESKTOP',
  mobile: 'MOBILE',
};

And here's what your script might look like once it's using it (I haven't tested these changes, but this should give you an idea of how it gets used):

// ...

const { initPrivateRooms, DEVICE_TYPE } = require('./privateRoom');
const privateRooms = initPrivateRooms();

// ...

wss.on('connection', function connection(ws, request, { authenticatedUser, userConnectionType }) {
  const room = privateRooms.getOrCreateRoom(authenticatedUser.email);

  room.join(userConnectionType, {
    onKickedOut() {
      ws.send('This connection has been overwritten with a new one');
    },
    onMessaged(msg) {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(msg);
      }
    },
  });

  ws.on('close', function close() {
    room.leave(userConnectionType);
    if (room.isEmpty()) {
      privateRooms.removeRoom(room.id);
    }
  });

  ws.on('message', function incoming(clientMessage) {
    room.messageMembers(buildMessage(room, clientMessage));
  });
});

function buildMessage(room, clientMessage) {
  return `
Desktop connected: ${room.hasDeviceType(DEVICE_TYPE.desktop)} 
Mobile connected: ${room.hasDeviceType(DEVICE_TYPE.mobile)}
room: ${room.id} 
Message: ${clientMessage}
`;
}

server.listen(port);
console.log(`listening on ${port}`);

You might create some other functions to aid in other areas of your script too. For example, in your upgrade event handler you could use a custom parseHeaders() function, or maybe an authenticate() function that handle some of the specific implementation details for the event listener.

Now, to address some of your questions:

• Your usage of the javascript map was perfectly valid there, and is exactly the kind of thing it's designed to do. In the revised example I had above I decided to move the map into a factory function because I found that felt cleaner in that particular scenario.
• Using strings like 'desktop' and 'mobile' to represent different states in javascript is just fine - in fact, it's a commonly used pattern when working with some frameworks or libraries (like redux). When you're making a more well-defined public-facing API, you might prefer placing the different possible strings into an object (as I decided to do in privateRooms.js). Doing so can help self-document what possible values your functions can take, and can aid the code editor with auto-complete.