From 20b1eb29aeafd7603f57443bf03eea44bc7fed7e Mon Sep 17 00:00:00 2001
From: Samuel Jensen <44519206+nichtsam@users.noreply.github.com>
Date: Tue, 3 Dec 2024 23:22:36 +0800
Subject: [PATCH] rate limit

---
 package.json   |  1 +
 pnpm-lock.yaml | 13 +++++++++++++
 server.js      | 30 ++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/package.json b/package.json
index 016ddcc..0da1cbf 100644
--- a/package.json
+++ b/package.json
@@ -66,6 +66,7 @@
     "drizzle-orm": "0.36.3",
     "drizzle-zod": "0.5.1",
     "express": "4.21.1",
+    "express-rate-limit": "7.4.1",
     "get-port": "7.1.0",
     "gray-matter": "4.0.3",
     "helmet": "8.0.0",
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 9f52437..1efadfc 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -146,6 +146,9 @@ importers:
       express:
         specifier: 4.21.1
         version: 4.21.1
+      express-rate-limit:
+        specifier: 7.4.1
+        version: 7.4.1(express@4.21.1)
       get-port:
         specifier: 7.1.0
         version: 7.1.0
@@ -3968,6 +3971,12 @@ packages:
     resolution: {integrity: sha512-bFi65yM+xZgk+u/KRIpekdSYkTB5W1pEf0Lt8Q8Msh7b+eQ7LXVtIB1Bkm4fvclDEL1b2CZkMhv2mOeF8tMdkA==}
     engines: {node: '>=12.0.0'}
 
+  express-rate-limit@7.4.1:
+    resolution: {integrity: sha512-KS3efpnpIDVIXopMc65EMbWbUht7qvTCdtCR2dD/IZmi9MIkopYESwyRqLgv8Pfu589+KqDqOdzJWW7AHoACeg==}
+    engines: {node: '>= 16'}
+    peerDependencies:
+      express: 4 || 5 || ^5.0.0-beta.1
+
   express@4.21.1:
     resolution: {integrity: sha512-YSFlK1Ee0/GC8QaO91tHcDxJiE/X4FbpAyQWkxAvG6AXCuR65YzK8ua6D9hvi/TzUfZMpc+BwuM1IPw8fmQBiQ==}
     engines: {node: '>= 0.10.0'}
@@ -11151,6 +11160,10 @@ snapshots:
 
   expect-type@1.1.0: {}
 
+  express-rate-limit@7.4.1(express@4.21.1):
+    dependencies:
+      express: 4.21.1
+
   express@4.21.1:
     dependencies:
       accepts: 1.3.8
diff --git a/server.js b/server.js
index 6413e42..97191df 100644
--- a/server.js
+++ b/server.js
@@ -5,6 +5,7 @@ import chalk from 'chalk'
 import closeWithGrace from 'close-with-grace'
 import compression from 'compression'
 import express from 'express'
+import { rateLimit } from 'express-rate-limit'
 import getPort, { portNumbers } from 'get-port'
 import helmet from 'helmet'
 import morgan from 'morgan'
@@ -56,6 +57,35 @@ app.use(
 	}),
 )
 
+const rateLimitConfig = {
+	skip: () => MODE !== 'production',
+	windowMs: 60 * 1000,
+	max: 1000,
+	standardHeaders: true,
+	legacyHeaders: false,
+	keyGenerator: (req) => {
+		return req.get('fly-client-ip') ?? `${req.ip}`
+	},
+}
+
+const strongestRateLimit = rateLimit({ ...rateLimitConfig, max: 10 })
+const strongRateLimit = rateLimit({ ...rateLimitConfig, max: 100 })
+const generalRateLimit = rateLimit(rateLimitConfig)
+
+app.use((req, res, next) => {
+	const criticalActions = ['/auth', '/onboarding']
+
+	if (req.method !== 'GET' && req.method !== 'HEAD') {
+		if (criticalActions.some((p) => req.path.startsWith(p))) {
+			return strongestRateLimit(req, res, next)
+		}
+
+		return strongRateLimit(req, res, next)
+	}
+
+	return generalRateLimit(req, res, next)
+})
+
 app.use((_, res, next) => {
 	res.locals.cspNonce = crypto.randomBytes(16).toString('hex')
 	next()