From 792be63c5c89b6ae6ecf4808c1347866e2a87d53 Mon Sep 17 00:00:00 2001
From: Zhongyi Lu <zhongyi_lu@outlook.com>
Date: Fri, 12 Apr 2024 22:03:36 -0700
Subject: [PATCH] fix: whitelist in real-ip

---
 src/middleware/whitelist.js | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/middleware/whitelist.js b/src/middleware/whitelist.js
index 87d5ac5a5..2ebc5e48f 100644
--- a/src/middleware/whitelist.js
+++ b/src/middleware/whitelist.js
@@ -19,6 +19,22 @@ if (fs.existsSync(whitelistPath)) {
     }
 }
 
+function getForwardedIp(req) {
+    // Check if X-Real-IP is available
+    if (req.headers['x-real-ip']) {
+        return req.headers['x-real-ip'];
+    }
+
+    // Check for X-Forwarded-For and parse if available
+    if (req.headers['x-forwarded-for']) {
+        const ipList = req.headers['x-forwarded-for'].split(',').map(ip => ip.trim());
+        return ipList[0];
+    }
+
+    // If none of the headers are available, return undefined
+    return undefined;
+}
+
 function getIpFromRequest(req) {
     let clientIp = req.connection.remoteAddress;
     let ip = ipaddr.parse(clientIp);
@@ -41,6 +57,7 @@ function getIpFromRequest(req) {
 function whitelistMiddleware(listen) {
     return function (req, res, next) {
         const clientIp = getIpFromRequest(req);
+        const forwardedIp = getForwardedIp(req);
 
         if (listen && !knownIPs.has(clientIp)) {
             const userAgent = req.headers['user-agent'];
@@ -58,9 +75,13 @@ function whitelistMiddleware(listen) {
         }
 
         //clientIp = req.connection.remoteAddress.split(':').pop();
-        if (whitelistMode === true && !whitelist.some(x => ipMatching.matches(clientIp, ipMatching.getMatch(x)))) {
-            console.log(color.red('Forbidden: Connection attempt from ' + clientIp + '. If you are attempting to connect, please add your IP address in whitelist or disable whitelist mode in config.yaml in root of SillyTavern folder.\n'));
-            return res.status(403).send('<b>Forbidden</b>: Connection attempt from <b>' + clientIp + '</b>. If you are attempting to connect, please add your IP address in whitelist or disable whitelist mode in config.yaml in root of SillyTavern folder.');
+        if (whitelistMode === true && !whitelist.some(x => ipMatching.matches(clientIp, ipMatching.getMatch(x)))
+            || forwardedIp && whitelistMode === true && !whitelist.some(x => ipMatching.matches(forwardedIp, ipMatching.getMatch(x)))
+        ) {
+            // Log the connection attempt with real IP address
+            const ipDetails = forwardedIp ? `${clientIp} (forwarded from ${forwardedIp})` : clientIp;
+            console.log(color.red('Forbidden: Connection attempt from ' + ipDetails + '. If you are attempting to connect, please add your IP address in whitelist or disable whitelist mode in config.yaml in root of SillyTavern folder.\n'));
+            return res.status(403).send('<b>Forbidden</b>: Connection attempt from <b>' + ipDetails + '</b>. If you are attempting to connect, please add your IP address in whitelist or disable whitelist mode in config.yaml in root of SillyTavern folder.');
         }
         next();
     };