Move whitelist middleware to its own module

server.js

@@ -29,8 +29,6 @@ const net = require('net');
 const dns = require('dns');
 const DeviceDetector = require('device-detector-js');
 const fetch = require('node-fetch').default;
-const ipaddr = require('ipaddr.js');
-const ipMatching = require('ip-matching');
 
 // image processing related library imports
 const jimp = require('jimp');
@@ -41,6 +39,7 @@ util.inspect.defaultOptions.maxStringLength = null;
 
 // local library imports
 const basicAuthMiddleware = require('./src/middleware/basicAuthMiddleware');
+const whitelistMiddleware = require('./src/middleware/whitelist');
 const { jsonParser, urlencodedParser } = require('./src/express-common.js');
 const contentManager = require('./src/endpoints/content-manager');
 const { migrateSecrets } = require('./src/endpoints/secrets');
@@ -108,19 +107,6 @@ app.use(responseTime());
 
 const server_port = process.env.SILLY_TAVERN_PORT || getConfigValue('port', 8000);
 
-const whitelistPath = path.join(process.cwd(), './whitelist.txt');
-let whitelist = getConfigValue('whitelist', []);
-
-if (fs.existsSync(whitelistPath)) {
-    try {
-        let whitelistTxt = fs.readFileSync(whitelistPath, 'utf-8');
-        whitelist = whitelistTxt.split('\n').filter(ip => ip).map(ip => ip.trim());
-    } catch (e) {
-        // Ignore errors that may occur when reading the whitelist (e.g. permissions)
-    }
-}
-
-const whitelistMode = getConfigValue('whitelistMode', true);
 const autorun = (getConfigValue('autorun', false) || cliArguments.autorun) && !cliArguments.ssl;
 const enableExtensions = getConfigValue('enableExtensions', true);
 const listen = getConfigValue('listen', false);
@@ -138,48 +124,7 @@ app.use(CORS);
 
 if (listen && getConfigValue('basicAuthMode', false)) app.use(basicAuthMiddleware);
 
-// IP Whitelist //
+app.use(whitelistMiddleware);
-let knownIPs = new Set();
-
-function getIpFromRequest(req) {
-    let clientIp = req.connection.remoteAddress;
-    let ip = ipaddr.parse(clientIp);
-    // Check if the IP address is IPv4-mapped IPv6 address
-    if (ip.kind() === 'ipv6' && ip instanceof ipaddr.IPv6 && ip.isIPv4MappedAddress()) {
-        const ipv4 = ip.toIPv4Address().toString();
-        clientIp = ipv4;
-    } else {
-        clientIp = ip;
-        clientIp = clientIp.toString();
-    }
-    return clientIp;
-}
-
-app.use(function (req, res, next) {
-    const clientIp = getIpFromRequest(req);
-
-    if (listen && !knownIPs.has(clientIp)) {
-        const userAgent = req.headers['user-agent'];
-        console.log(color.yellow(`New connection from ${clientIp}; User Agent: ${userAgent}\n`));
-        knownIPs.add(clientIp);
-
-        // Write access log
-        const timestamp = new Date().toISOString();
-        const log = `${timestamp} ${clientIp} ${userAgent}\n`;
-        fs.appendFile('access.log', log, (err) => {
-            if (err) {
-                console.error('Failed to write access log:', err);
-            }
-        });
-    }
-
-    //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.');
-    }
-    next();
-});
 
 // CSRF Protection //
 if (!cliArguments.disableCsrf) {

src/middleware/whitelist.js

@@ -0,0 +1,66 @@
+const path = require('path');
+const fs = require('fs');
+const ipaddr = require('ipaddr.js');
+const ipMatching = require('ip-matching');
+
+const { color, getConfigValue } = require('../util');
+
+const whitelistPath = path.join(process.cwd(), './whitelist.txt');
+let whitelist = getConfigValue('whitelist', []);
+const listen = getConfigValue('listen', false);
+
+if (fs.existsSync(whitelistPath)) {
+    try {
+        let whitelistTxt = fs.readFileSync(whitelistPath, 'utf-8');
+        whitelist = whitelistTxt.split('\n').filter(ip => ip).map(ip => ip.trim());
+    } catch (e) {
+        // Ignore errors that may occur when reading the whitelist (e.g. permissions)
+    }
+}
+
+const whitelistMode = getConfigValue('whitelistMode', true);
+
+// IP Whitelist //
+let knownIPs = new Set();
+
+function getIpFromRequest(req) {
+    let clientIp = req.connection.remoteAddress;
+    let ip = ipaddr.parse(clientIp);
+    // Check if the IP address is IPv4-mapped IPv6 address
+    if (ip.kind() === 'ipv6' && ip instanceof ipaddr.IPv6 && ip.isIPv4MappedAddress()) {
+        const ipv4 = ip.toIPv4Address().toString();
+        clientIp = ipv4;
+    } else {
+        clientIp = ip;
+        clientIp = clientIp.toString();
+    }
+    return clientIp;
+}
+
+const whitelistMiddleware = function (req, res, next) {
+    const clientIp = getIpFromRequest(req);
+
+    if (listen && !knownIPs.has(clientIp)) {
+        const userAgent = req.headers['user-agent'];
+        console.log(color.yellow(`New connection from ${clientIp}; User Agent: ${userAgent}\n`));
+        knownIPs.add(clientIp);
+
+        // Write access log
+        const timestamp = new Date().toISOString();
+        const log = `${timestamp} ${clientIp} ${userAgent}\n`;
+        fs.appendFile('access.log', log, (err) => {
+            if (err) {
+                console.error('Failed to write access log:', err);
+            }
+        });
+    }
+
+    //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.');
+    }
+    next();
+};
+
+module.exports = whitelistMiddleware;