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('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.');
+ 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('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.');
}
next();
};