[OA][Infrastructure] - Init WS Server

README.md

@@ -35,11 +35,13 @@ cd sarge/
 
 ```sh
 pnpm install
+cd src/ws && pnpm install
 ```
 
 3. Setup the .env file
 
 ```
+# in the project's root directory
 mv .env.example .env
 ```
 

src/ws/Dockerfile

@@ -4,4 +4,4 @@ RUN corepack enable pnpm
 COPY package.json index.js pnpm-lock.yaml* ./
 RUN pnpm install
 
-CMD ["node", "index.js"]
+CMD ["node", "index.js"]

src/ws/index.js

@@ -1,9 +1,82 @@
 import { WebSocketServer } from 'ws';
+import { jwtVerify } from 'jose';
 
 const ws = new WebSocketServer({ port: 8080 });
+const secret = new TextEncoder().encode(process.env.JWT_SECRET);
 
-ws.on('connection', (ws) => {
-    ws.on('message', (data) => {
-        ws.send(data.toString('utf-8'));
+console.log('Sarge WS server listening on 8080');
+
+// <userId, socket>
+const clients = new Map();
+
+const HEARTBEAT_INTERVAL = 5000;
+
+function getTokenFromRequest(request) {
+    return new URL(request.url ?? '/', 'ws://x').searchParams.get('token');
+}
+
+async function verifyToken(token) {
+    const { payload } = await jwtVerify(token, secret);
+    return payload;
+}
+
+function startHeartbeat(socket) {
+    return setInterval(() => {
+        if (!socket.isAlive) {
+            // TODO: add an offline timer before the socket gets terminated unique to each user
+            // Once done, send an API request back to server to end the OA
+
+            // terminate() jumps to socket.on('close')
+            socket.terminate();
+            return;
+        }
+        socket.isAlive = false;
+        socket.ping();
+    }, HEARTBEAT_INTERVAL);
+}
+
+ws.on('connection', async (socket, request) => {
+    const token = getTokenFromRequest(request);
+    try {
+        const payload = await verifyToken(token ?? '');
+        socket.userId = payload.userId;
+        socket.isAlive = true;
+        clients.set(payload.userId, socket);
+    } catch {
+        socket.close(1008, 'Unauthorized');
+        return;
+    }
+
+    socket.heartbeat = startHeartbeat(socket);
+
+    socket.on('message', (data) => {
+        socket.send(data.toString('utf-8'));
     });
+
+    socket.on('pong', () => {
+        socket.isAlive = true;
+    });
+
+    socket.on('close', (code, reason) => {
+        if (!socket.userId) return;
+        clearInterval(socket.heartbeat);
+        clients.delete(socket.userId);
+        console.log(
+            `Client ${socket.userId} disconnected | code: ${code}, reason: ${reason.toString()}`
+        );
+    });
+
+    socket.on('error', (err) => {
+        console.error(`Socket error for ${socket.userId}:`, err.message);
+    });
+});
+
+ws.on('error', (err) => {
+    console.error('Socket error: ', err);
+});
+
+ws.on('close', () => {
+    for (const socket of clients.values()) {
+        socket.close(1001, 'Shutting down server');
+    }
 });

src/ws/package.json

@@ -5,12 +5,17 @@
     "main": "index.js",
     "type": "module",
     "scripts": {
-        "test": "echo \"Error: no test specified\" && exit 1"
+        "start": "node index.js",
+        "dev": "nodemon index.js"
     },
     "keywords": [],
     "author": "",
     "license": "ISC",
     "dependencies": {
+        "jose": "^6.2.2",
         "ws": "^8.20.0"
+    },
+    "devDependencies": {
+        "nodemon": "^3.1.14"
     }
 }