1
import path from 'node:path';
2
import fs from 'node:fs';
3
import process from 'node:process';
4
import dns from 'node:dns';
5
import Handlebars from 'handlebars';
6
import ipMatching from 'ip-matching';
7
import isDocker from 'is-docker';
8

9
import { filterValidIpPatterns, getIpFromRequest, getRealOrForwardedIp } from '../express-common.js';
10
import { color, getConfigValue, safeReadFileSync } from '../util.js';
11

12
const whitelistPath = path.join(process.cwd(), './whitelist.txt');
13
const enableForwardedWhitelist = !!getConfigValue('enableForwardedWhitelist', false, 'boolean');
14
const whitelistDockerHosts = !!getConfigValue('whitelistDockerHosts', true, 'boolean');
15
/** @type {string[]} */
16
let whitelist = getConfigValue('whitelist', []);
17

18
if (fs.existsSync(whitelistPath)) {
19
    console.warn(color.yellow('whitelist.txt is deprecated and will be removed in a future release.'));
20
    console.warn(color.yellow('Please migrate its contents to the whitelist field in config.yaml. See the documentation for more details.'));
21
    try {
22
        let whitelistTxt = fs.readFileSync(whitelistPath, 'utf-8');
23
        whitelist = whitelistTxt.split('\n').filter(ip => ip).map(ip => ip.trim());
24
    } catch (e) {
25
        // Ignore errors that may occur when reading the whitelist (e.g. permissions)
26
    }
27
}
28

29
whitelist = filterValidIpPatterns(whitelist, (entry, message) => `${color.red('Warning')}: Ignoring invalid whitelist entry ${color.yellow(entry)} - ${message}`);
30

31
/**
32
 * Resolves the IP addresses of Docker hostnames and adds them to the whitelist.
33
 * @returns {Promise<void>} Promise that resolves when the Docker hostnames are resolved
34
 */
35
async function addDockerHostsToWhitelist() {
36
    if (!whitelistDockerHosts || !isDocker()) {
37
        return;
38
    }
39

40
    const whitelistHosts = ['host.docker.internal', 'gateway.docker.internal'];
41

42
    for (const entry of whitelistHosts) {
43
        try {
44
            const result = await dns.promises.lookup(entry);
45
            console.info(`Resolved whitelist hostname ${color.green(entry)} to IPv${result.family} address ${color.green(result.address)}`);
46
            whitelist.push(result.address);
47
        } catch (e) {
48
            console.warn(`Failed to resolve whitelist hostname ${color.red(entry)}: ${e.message}`);
49
        }
50
    }
51
}
52

53
/**
54
 * Returns a middleware function that checks if the client IP is in the whitelist.
55
 * @returns {Promise<import('express').RequestHandler>} Promise that resolves to the middleware function
56
 */
57
export default async function getWhitelistMiddleware() {
58
    const forbiddenWebpage = Handlebars.compile(
59
        safeReadFileSync(path.join(globalThis.DATA_ROOT, '_errors', 'forbidden-by-whitelist.html')) ?? '',
60
    );
61

62
    const noLogPaths = [
63
        '/favicon.ico',
64
    ];
65

66
    await addDockerHostsToWhitelist();
67

68
    return function (req, res, next) {
69
        const clientIp = getIpFromRequest(req);
70
        const forwardedIp = enableForwardedWhitelist && getRealOrForwardedIp(req);
71
        const userAgent = req.headers['user-agent'];
72

73
        /**
74
         * Checks if an IP address matches any entry in the whitelist.
75
         * @param {string[]} whitelist - The list of whitelisted IPs/CIDRs
76
         * @param {string} ip - The IP address to check
77
         * @returns {boolean} True if the IP matches any whitelist entry
78
         */
79
        function isIPInWhitelist(whitelist, ip) {
80
            return whitelist.some(x => ipMatching.matches(ip, ipMatching.getMatch(x)));
81
        }
82

83
        //clientIp = req.connection.remoteAddress.split(':').pop();
84
        if (!isIPInWhitelist(whitelist, clientIp)
85
            || (forwardedIp && !isIPInWhitelist(whitelist, forwardedIp))
86
        ) {
87
            // Log the connection attempt with real IP address
88
            const ipDetails = forwardedIp
89
                ? `${clientIp} (forwarded from ${forwardedIp})`
90
                : clientIp;
91

92
            if (!noLogPaths.includes(req.path)) {
93
                console.warn(
94
                    color.red(
95
                        `Blocked connection from ${ipDetails}; User Agent: ${userAgent}\n\tTo allow this connection, add its IP address to the whitelist or disable whitelist mode by editing config.yaml in the root directory of your SillyTavern installation.\n`,
96
                    ),
97
                );
98
            }
99

100
            return res.status(403).send(forbiddenWebpage({ ipDetails }));
101
        }
102
        next();
103
    };
104
}