Add ws

server/index.ts

@@ -12,6 +12,7 @@ import {
 } from "@server/middlewares";
 import internal from "@server/routers/internal";
 import { authenticated, unauthenticated } from "@server/routers/external";
+import { router as wsRouter, handleWSUpgrade } from '@server/routers/ws';
 import cookieParser from "cookie-parser";
 import { User } from "@server/db/schema";
 import { ensureActions } from "./db/ensureActions";
@@ -50,22 +51,25 @@ app.prepare().then(() => {
     externalServer.use(logIncomingMiddleware);
     externalServer.use(prefix, unauthenticated);
     externalServer.use(prefix, authenticated);
-
+    externalServer.use(`${prefix}/ws`, wsRouter);
+    
     externalServer.use(notFoundMiddleware);
-
+    
     // We are using NEXT from here on
     externalServer.all("*", (req: Request, res: Response) => {
         const parsedUrl = parse(req.url!, true);
         handle(req, res, parsedUrl);
     });
-
+    
-    externalServer.listen(externalPort, (err?: any) => {
+    const httpServer = externalServer.listen(externalPort, (err?: any) => {
         if (err) throw err;
         logger.info(
             `Main server is running on http://localhost:${externalPort}`
         );
     });
 
+    handleWSUpgrade(httpServer);
+
     externalServer.use(errorHandlerMiddleware);
 
     // Internal server

server/routers/ws.ts

@@ -0,0 +1,144 @@
+import { Router, Request, Response } from 'express';
+import { Server as HttpServer } from 'http';
+import { WebSocket, WebSocketServer } from 'ws';
+import { IncomingMessage } from 'http';
+import { Socket } from 'net';
+
+// Custom interfaces
+interface WebSocketRequest extends IncomingMessage {
+  token?: string;
+}
+
+interface AuthenticatedWebSocket extends WebSocket {
+  userId?: string;
+  isAlive?: boolean;
+}
+
+interface TokenPayload {
+  userId: string;
+  // Add other token payload properties as needed
+}
+
+const router: Router = Router();
+const wss: WebSocketServer = new WebSocketServer({ noServer: true });
+
+// Token verification middleware
+const verifyToken = async (token: string): Promise<TokenPayload | null> => {
+  try {
+    // This is where you'd implement your token verification logic
+    // For example, verify JWT, check against database, etc.
+    // Return the token payload if valid, null if invalid
+    return { userId: 'dummy-user-id' }; // Placeholder return
+  } catch (error) {
+    console.error('Token verification failed:', error);
+    return null;
+  }
+};
+
+// Handle WebSocket upgrade requests
+router.get('/ws', (req: Request, res: Response) => {
+  // WebSocket upgrade will be handled by the server
+  res.status(200).send('WebSocket endpoint');
+});
+
+// Set up WebSocket server handling
+const handleWSUpgrade = (server: HttpServer): void => {
+  server.on('upgrade', async (request: WebSocketRequest, socket: Socket, head: Buffer) => {
+    try {
+      // Extract token from query parameters or headers
+      const token = request.url?.includes('?')
+        ? new URLSearchParams(request.url.split('?')[1]).get('token') || ''
+        : request.headers['sec-websocket-protocol'];
+
+      if (!token) {
+        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+        socket.destroy();
+        return;
+      }
+
+      // Verify the token
+      const tokenPayload = await verifyToken(token);
+      if (!tokenPayload) {
+        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+        socket.destroy();
+        return;
+      }
+
+      // Store token payload data in the request for later use
+      request.token = token;
+
+      wss.handleUpgrade(request, socket, head, (ws: AuthenticatedWebSocket) => {
+        // Attach user data to the WebSocket instance
+        ws.userId = tokenPayload.userId;
+        ws.isAlive = true;
+        wss.emit('connection', ws, request);
+      });
+    } catch (error) {
+      console.error('Upgrade error:', error);
+      socket.write('HTTP/1.1 500 Internal Server Error\r\n\r\n');
+      socket.destroy();
+    }
+  });
+};
+
+// WebSocket message interface
+interface WSMessage {
+  type: string;
+  data: any;
+}
+
+// WebSocket connection handler
+wss.on('connection', (ws: AuthenticatedWebSocket, request: WebSocketRequest) => {
+  console.log(`Client connected - User ID: ${ws.userId}`);
+
+  // Set up ping-pong for connection health check
+  const pingInterval = setInterval(() => {
+    if (ws.isAlive === false) {
+      clearInterval(pingInterval);
+      return ws.terminate();
+    }
+    ws.isAlive = false;
+    ws.ping();
+  }, 30000);
+
+  // Handle pong response
+  ws.on('pong', () => {
+    ws.isAlive = true;
+  });
+
+  // Set up message handler
+  ws.on('message', (data) => {
+    try {
+      const message: WSMessage = JSON.parse(data.toString());
+      console.log('Received:', message);
+
+      // Echo the message back
+      ws.send(JSON.stringify({
+        type: 'echo',
+        data: message
+      }));
+    } catch (error) {
+      console.error('Message parsing error:', error);
+      ws.send(JSON.stringify({
+        type: 'error',
+        data: 'Invalid message format'
+      }));
+    }
+  });
+
+  // Handle client disconnect
+  ws.on('close', () => {
+    clearInterval(pingInterval);
+    console.log(`Client disconnected - User ID: ${ws.userId}`);
+  });
+
+  // Handle errors
+  ws.on('error', (error: Error) => {
+    console.error('WebSocket error:', error);
+  });
+});
+
+export {
+  router,
+  handleWSUpgrade
+};

server/routers/ws/ws.ts

@@ -1,68 +0,0 @@
-const express = require('express');
-const router = express.Router();
-const WebSocket = require('ws');
-const wss = new WebSocket.Server({ noServer: true });
-
-// Token verification middleware
-const verifyToken = (token) => {
-  // This is where you'd implement your token verification logic
-  // For example, verify JWT, check against database, etc.
-  // Return true if token is valid, false otherwise
-  return true; // Placeholder return
-};
-
-// Handle WebSocket upgrade requests
-router.get('/ws', (req, res) => {
-  // WebSocket upgrade will be handled by the server
-  res.status(200).send('WebSocket endpoint');
-});
-
-// Set up WebSocket server handling
-const handleWSUpgrade = (server) => {
-  server.on('upgrade', (request, socket, head) => {
-    // Extract token from query parameters or headers
-    const token = request.url.includes('?') 
-      ? new URLSearchParams(request.url.split('?')[1]).get('token')
-      : request.headers['sec-websocket-protocol'];
-
-    // Verify the token
-    if (!token || !verifyToken(token)) {
-      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
-      socket.destroy();
-      return;
-    }
-
-    wss.handleUpgrade(request, socket, head, (ws) => {
-      wss.emit('connection', ws, request);
-    });
-  });
-};
-
-// WebSocket connection handler
-wss.on('connection', (ws, request) => {
-  console.log('Client connected');
-
-  // Set up message handler
-  ws.on('message', (message) => {
-    console.log('Received:', message.toString());
-    
-    // Echo the message back
-    ws.send(`Server received: ${message}`);
-  });
-
-  // Handle client disconnect
-  ws.on('close', () => {
-    console.log('Client disconnected');
-  });
-
-  // Handle errors
-  ws.on('error', (error) => {
-    console.error('WebSocket error:', error);
-  });
-});
-
-// Export both the router and the upgrade handler
-module.exports = {
-  router,
-  handleWSUpgrade
-};