From a30d08313db181317f3040ea5fa28a2d19c7f649 Mon Sep 17 00:00:00 2001
From: jason5ng32 <jason5ng32@gmail.com>
Date: Fri, 26 Apr 2024 19:07:30 +0800
Subject: [PATCH] Security updates

---
 .env.example |  4 ++-
 server.js    | 71 +++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/.env.example b/.env.example
index bef4cd00..005a0021 100644
--- a/.env.example
+++ b/.env.example
@@ -6,4 +6,6 @@ KEYCDN_USER_AGENT=""
 IPCHECKING_API_KEY=""
 CLOUDFLARE_API=""
 VITE_RECAPTCHA_SITE_KEY=""
-RECAPTCHA_SECRET_KEY=""
\ No newline at end of file
+RECAPTCHA_SECRET_KEY=""
+BLACKLIST_LOG_FILE_PATH=""
+RATE_LIMIT=""
\ No newline at end of file
diff --git a/server.js b/server.js
index 8568cf01..f0ab688e 100644
--- a/server.js
+++ b/server.js
@@ -1,6 +1,7 @@
 import dotenv from 'dotenv';
 import express from 'express';
 import path from 'path';
+import fs from 'fs';
 import { fileURLToPath } from 'url';
 import mapHandler from './api/map.js';
 import ipinfoHandler from './api/ipinfo.js';
@@ -18,13 +19,75 @@ dotenv.config();
 
 const app = express();
 const port = process.env.PORT || 11966;
+const blackListIPLogFilePath = process.env.BLACKLIST_LOG_FILE_PATH || '';
+const rateLimitSet = process.env.RATE_LIMIT || '60';
 
-app.set('trust proxy', true);
+app.set('trust proxy', 1);
+
+// 获取客户端 IP 的辅助函数
+function getClientIp(req) {
+    const cfIp = req.headers['cf-connecting-ip']; // Cloudflare IP
+    const forwardedIps = req.headers['x-forwarded-for'] ? req.headers['x-forwarded-for'].split(',')[0] : null;
+    const cfIpV6 = req.headers['cf-connecting-ipv6'];
+    return cfIp || forwardedIps || cfIpV6 || req.ip;
+}
+
+// 记录限流触发的 IP 地址
+function formatDate(timestamp) {
+    return new Date(timestamp).toLocaleString('en-US', { timeZone: 'Asia/Shanghai' });
+}
+
+function logLimitedIP(ip) {
+    const logPath = path.join(__dirname, blackListIPLogFilePath);
+
+    fs.readFile(logPath, 'utf8', (err, data) => {
+        if (err && err.code !== 'ENOENT') {
+            console.error('Error reading the log file:', err);
+            return;
+        }
+
+        const now = Date.now();
+        let newCount = 1;
+        let logExists = false;
+        let updatedData = '';
+
+        if (data) {
+            const lines = data.split('\n');
+            updatedData = lines.map(line => {
+                const [currentIp, count, timestamp] = line.split(',');
+                if (currentIp === ip) {
+                    newCount = parseInt(count, 10) + 1;
+                    logExists = true;
+                    return `${ip},${newCount},${timestamp}`;  // Update count but keep the original timestamp
+                }
+                return line;
+            }).join('\n');
+        }
+
+        if (!logExists) {
+            const newLine = `${ip},${newCount},${formatDate(now)}`;
+            updatedData += (updatedData ? '\n' : '') + newLine;
+        }
+
+        fs.writeFile(logPath, updatedData, 'utf8', err => {
+            if (err) {
+                console.error('Failed to write to log file:', err);
+            }
+        });
+    });
+}
 
 const apiLimiter = rateLimit({
-    windowMs: 60 * 60 * 1000, // 60 分钟的窗口
-    max: 60,
-    message: 'Too many requests'
+    windowMs: 60 * 60 * 1000,
+    max: parseInt(rateLimitSet, 10),
+    message: 'Too Many Requests',
+    handler: (req, res, next) => {
+        const ip = getClientIp(req);
+        if (req.rateLimit.current === req.rateLimit.limit + 1 && blackListIPLogFilePath) {
+            logLimitedIP(ip);
+        }
+        res.status(429).json({ message: 'Too Many Requests' });
+    }
 });
 
 app.use('/api', apiLimiter);