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 privaterooms
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 aroom
in this case?