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(); };