Merge pull request #28 from TheRedScreen64/backend-rate-limits

Backend rate limits

Author: TheRedScreen64

backend/package.json

@@ -31,6 +31,7 @@
       "cors": "^2.8.5",
       "dotenv": "^16.4.4",
       "express": "^4.18.2",
+      "express-rate-limit": "^7.2.0",
       "lucia": "^3.0.1",
       "node-cron": "^3.0.3",
       "oslo": "^1.1.1",

backend/src/index.ts

@@ -1,6 +1,7 @@
 import cookieParser from "cookie-parser";
 import * as dotenv from "dotenv";
 import express from "express";
+import rateLimit from "express-rate-limit";
 import { createServer } from "http";
 import { WebSocketServer } from "ws";
 import { lucia } from "./lib/auth.js";
@@ -38,6 +39,15 @@ const errorHandler = (err: any, req: any, res: any, next: any) => {
    }
 };
 
+const authLimiter = rateLimit({
+   windowMs: 60 * 1000,
+   max: 5,
+});
+const limiter = rateLimit({
+   windowMs: 60 * 1000,
+   max: 100,
+});
+
 const app = express();
 const server = createServer(app);
 const wss = new WebSocketServer({ server });
@@ -70,6 +80,9 @@ app.use(async (req, res, next) => {
    return next();
 });
 
+app.use(limiter);
+app.use("/auth/*", authLimiter);
+
 app.use(
    loginRouter,
    logoutRouter,

backend/src/lib/websocket.ts

@@ -13,8 +13,36 @@ const messageSchema = z.object({
 
 // TODO: Overhaul
 
+let connections = new Map<string, any>();
+
+type RateLimit = {
+   count: number;
+};
+
+const generalRateLimit = {
+   timeWindowMs: 10000,
+   count: 10,
+};
+const globalChatRateLimit: RateLimit = {
+   count: 1,
+};
+
 function initWebsocket(wss: WebSocketServer) {
    wss.on("connection", async (ws, req) => {
+      if (!req.socket.remoteAddress) {
+         ws.close();
+         return;
+      }
+      const ip = req.socket.remoteAddress;
+      let connection = connections.get(ip);
+      if (!connection) {
+         connection = {
+            count: 0,
+            lastMessageTime: Date.now(),
+         };
+         connections.set(ip, connection);
+      }
+
       let sessionId = url.parse(req.url!, true).query.session!;
       if (!sessionId || Array.isArray(sessionId)) {
          ws.close();
@@ -31,7 +59,15 @@ function initWebsocket(wss: WebSocketServer) {
 
       ws.on("message", async (reqData: string) => {
          try {
-            await validateSession(ws, sessionId);
+            const exceededRateLimit = await handleRateLimit(connection!);
+            if (exceededRateLimit) {
+               return ws.send(JSON.stringify({ error: { message: "Rate limit exceeded, slow down" } }));
+            }
+
+            const sessionValid = await validateSession(sessionId);
+            if (!sessionValid) {
+               return ws.send(JSON.stringify({ error: { message: "Invalid session" } }));
+            }
 
             let reqJson = JSON.parse(reqData);
             let parsed = messageSchema.safeParse(reqJson);
@@ -43,6 +79,10 @@ function initWebsocket(wss: WebSocketServer) {
 
             switch (type) {
                case "globalMessage": {
+                  const exceededRateLimit = await checkRateLimit(connection!, globalChatRateLimit);
+                  if (exceededRateLimit) {
+                     return ws.send(JSON.stringify({ error: { message: "Rate limit exceeded, slow down" } }));
+                  }
                   handleGlobalMessage(ws, wss, data, user);
                   break;
                }
@@ -123,19 +163,45 @@ async function handleChatMessage(ws: WebSocket, data: any, user: User) {
    }
 }
 
-async function validateSession(ws: WebSocket, sessionId: string | string[]) {
+async function validateSession(sessionId: string | string[]): Promise<boolean> {
    if (Array.isArray(sessionId)) {
-      ws.close();
-      return;
+      return false;
    }
    const { session, user } = await lucia.validateSession(sessionId);
    if (!session || !user) {
-      ws.close();
-      return;
+      return false;
+   } else {
+      return true;
+   }
+}
+
+async function handleRateLimit(connection: any): Promise<boolean> {
+   const now = Date.now();
+   const elapsedTime = now - connection.lastMessageTime;
+
+   if (elapsedTime < generalRateLimit.timeWindowMs) {
+      connection.count++;
+      if (connection.count > generalRateLimit.count) {
+         return true;
+      }
+   } else {
+      connection.count = 1;
+      connection.lastMessageTime = now;
+   }
+   return false;
+}
+
+async function checkRateLimit(connection: any, rateLimit: RateLimit): Promise<boolean> {
+   console.log(connection.count > rateLimit.count);
+   if (connection.count > rateLimit.count) {
+      return true;
+   } else {
+      return false;
    }
 }
 
 function broadcast(wss: WebSocketServer, data: any, skip: WebSocket | null) {
+   wss.clients;
    wss.clients.forEach((client) => {
       if (skip && client == skip) return;
       if (client.readyState === WebSocket.OPEN) {

backend/src/routes/user/exists.ts

@@ -6,6 +6,9 @@ import { formatPrismaError } from "../../lib/utils.js";
 export const existsRouter = express.Router();
 
 existsRouter.post("/user/exists", async (req, res, next) => {
+   if (!res.locals.user) {
+      return next({ msg: "Not authorized", status: 401 });
+   }
    if (!req.body) {
       return next({ msg: "No input provided", status: 400 });
    }