wireguard-ui/src/functions/backend/setup/setup-wireguard-host.ts

import type {
    BUN_SQLITE_WGUI_CLIENTS,
    BUN_SQLITE_WGUI_CLIENT_RULES,
    BUN_SQLITE_WGUI_HOSTS,
} from "@/db/types/db";
import grabDirNames from "@/src/utils/grab-dir-names";
import { execSync } from "node:child_process";
import { chmodSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import validateIpv4 from "@/src/utils/validate-ipv4";
import type { APIResponseObject } from "@moduletrace/bunext/types";
import BunSQLite from "@moduletrace/bun-sqlite";
import type { TableType, User } from "@/src/types";
import checkPrivateIPAvailability from "./check-private-ip-availability";
import manageWireguardHost from "./manage-wireguard-host";
import applyWireguardHostLive from "./apply-wireguard-host-live";
import grabHostDirnames from "./grab-host-dir-names";
import grabClientDirnames from "./grab-client-dir-names";
import buildHostIptablesScripts from "./build-host-iptables-scripts";
import allowSubnetFirewallRules from "../firewall/allow-subnet-firewall-rules";
import { AppData } from "@/src/data/app-data";
import { TlsData } from "@/src/data/tls-data";

const { WIREGUARD_PRIVATE_KEY_FILE_NAME, WIREGUARD_PUBLIC_KEY_FILE_NAME } =
    grabDirNames();

type Params = {
    host?: BUN_SQLITE_WGUI_HOSTS;
    wg_subnet_ip?: string;
    user: User;
    is_update_after_client_setup?: boolean;
    /**
     * Apply the new config to the running tunnel (`wg syncconf` + atomic
     * iptables reload) instead of restarting it, and skip the firewall
     * manager rules. Only for peer/rule changes — the interface address and
     * listen port must be unchanged.
     */
    apply_live?: boolean;
};

export default async function setupWireguardHost({
    host,
    wg_subnet_ip,
    user,
    is_update_after_client_setup,
    apply_live,
}: Params): Promise<APIResponseObject> {
    const {
        HOST_CLIENTS_DIR,
        HOST_CONFIG_DIR,
        HOST_CONFIG_FILE,
        HOST_ID,
        HOST_IPTABLES_DIR,
        HOST_PRIVATE_KEY_FILE,
        HOST_PUBLIC_KEY_FILE,
        INTERFACE_NAME,
        POST_DOWN_PATH,
        POST_UP_PATH,
    } = grabHostDirnames({ host });

    const host_clients_res = await BunSQLite.select<
        BUN_SQLITE_WGUI_CLIENTS,
        TableType
    >({
        table: "clients",
        query: {
            query: {
                host_id: {
                    value: String(HOST_ID),
                },
            },
        },
    });

    const clients = host_clients_res.payload || [];

    const client_rules_res = await BunSQLite.select<
        BUN_SQLITE_WGUI_CLIENT_RULES,
        TableType
    >({
        table: "client_rules",
    });

    const client_ids = new Set(
        clients
            .map((client) => client.id)
            .filter((id): id is number => Boolean(id)),
    );

    const rules_by_client_id = new Map<
        number,
        BUN_SQLITE_WGUI_CLIENT_RULES[]
    >();

    for (let i = 0; i < (client_rules_res.payload || []).length; i++) {
        const rule = client_rules_res.payload?.[i];

        if (!rule?.client_id || !client_ids.has(rule.client_id)) {
            continue;
        }

        const existing_rules = rules_by_client_id.get(rule.client_id) || [];
        existing_rules.push(rule);
        rules_by_client_id.set(rule.client_id, existing_rules);
    }

    const sibling_hosts_res = await BunSQLite.select<
        BUN_SQLITE_WGUI_HOSTS,
        TableType
    >({
        table: "hosts",
    });

    const sibling_subnets = [
        ...new Set(
            (sibling_hosts_res.payload || [])
                .filter(
                    (host_row) =>
                        host_row.id != HOST_ID &&
                        typeof host_row.wg_ip_address == "string" &&
                        host_row.wg_ip_address.trim() != "",
                )
                .map((host_row) => {
                    const parts = String(host_row.wg_ip_address)
                        .trim()
                        .split(".");
                    return `${parts.slice(0, 3).join(".")}.0/24`;
                }),
        ),
    ];

    const HOST_WG_IP = host?.wg_ip_address || wg_subnet_ip;

    if (!HOST_WG_IP) {
        return {
            success: false,
            msg: `No Host Private IP address provided`,
        };
    }

    if (!validateIpv4({ ip: HOST_WG_IP })) {
        return {
            success: false,
            msg: `Invalid Host Private IP address: ${HOST_WG_IP}`,
        };
    }

    const LISTEN_PORT = Number(
        host?.listen_port || AppData["DefaultWGListenPort"],
    );

    if (!Number.isInteger(LISTEN_PORT)) {
        return {
            success: false,
            msg: `Invalid Host listen port: ${host?.listen_port}`,
        };
    }

    const is_ip_available = is_update_after_client_setup
        ? { success: true }
        : host?.wg_ip_address
          ? { success: true }
          : await checkPrivateIPAvailability({
                ip_address: HOST_WG_IP,
            });

    if (!is_ip_available.success && !is_update_after_client_setup) {
        return {
            success: false,
            msg: `IP not available`,
        };
    }

    if (!is_update_after_client_setup) {
        let pre_sh = ``;

        pre_sh += `set -e\n`;
        pre_sh += `mkdir -p ${HOST_CONFIG_DIR}\n`;
        pre_sh += `mkdir -p ${HOST_CLIENTS_DIR}\n`;
        pre_sh += `mkdir -p ${HOST_IPTABLES_DIR}\n`;
        pre_sh += `cd ${HOST_CONFIG_DIR}\n`;
        pre_sh += `if [ ! -f ${WIREGUARD_PRIVATE_KEY_FILE_NAME} ]; then\n`;
        pre_sh += `    wg genkey | tee ${WIREGUARD_PRIVATE_KEY_FILE_NAME} | wg pubkey > ${WIREGUARD_PUBLIC_KEY_FILE_NAME}\n`;
        pre_sh += `fi\n`;

        try {
            const exec_pre_setup = execSync(pre_sh, { encoding: "utf-8" });
        } catch (error: any) {
            return {
                success: false,
                msg: error.message,
            };
        }
    }

    try {
        const HOST_PUBLIC_KEY = readFileSync(HOST_PUBLIC_KEY_FILE, "utf-8").trim();

        const HOST_PRIVATE_KEY = readFileSync(
            HOST_PRIVATE_KEY_FILE,
            "utf-8",
        ).trim();

        if (!is_update_after_client_setup) {
            const update_host = await BunSQLite.insert<
                BUN_SQLITE_WGUI_HOSTS,
                TableType
            >({
                table: "hosts",
                data: [
                    {
                        id: HOST_ID,
                        user_id: user.id,
                        wg_ip_address: HOST_WG_IP,
                        public_key: HOST_PUBLIC_KEY,
                    },
                ],
                update_on_duplicate: true,
            });
        }

        const HOST_SUBNET = `${HOST_WG_IP.split(".").slice(0, 3).join(".")}.0/24`;

        const iptables_scripts = buildHostIptablesScripts({
            host_id: Number(HOST_ID),
            interface_name: INTERFACE_NAME,
            host_subnet: HOST_SUBNET,
            sibling_subnets,
            ingress_ports: TlsData.IngressEnabled
                ? [TlsData.CaddyHttpsPort]
                : [],
            clients: clients.map((client) => ({
                id: client.id,
                wg_ip_address: client.wg_ip_address,
                allow_all_ips: client.allow_all_ips == 1,
                rules:
                    client.id && typeof client.id == "number"
                        ? rules_by_client_id.get(client.id) || []
                        : [],
            })),
        });

        if (
            !iptables_scripts.success ||
            !iptables_scripts.post_up ||
            !iptables_scripts.post_down
        ) {
            return {
                success: false,
                msg: iptables_scripts.msg || `Could not build iptables scripts`,
            };
        }

        mkdirSync(HOST_IPTABLES_DIR, { recursive: true });
        await Bun.write(POST_UP_PATH, iptables_scripts.post_up);
        await Bun.write(POST_DOWN_PATH, iptables_scripts.post_down);
        chmodSync(POST_UP_PATH, 0o755);
        chmodSync(POST_DOWN_PATH, 0o755);

        let HOST_CONFIG = ``;

        HOST_CONFIG += `[Interface]\n`;
        HOST_CONFIG += `Address = ${HOST_WG_IP}/24\n`;
        HOST_CONFIG += `ListenPort = ${LISTEN_PORT}\n`;
        HOST_CONFIG += `PrivateKey = ${HOST_PRIVATE_KEY}\n`;
        HOST_CONFIG += `PostUp = ${POST_UP_PATH}\n`;
        HOST_CONFIG += `PostDown = ${POST_DOWN_PATH}\n`;
        HOST_CONFIG += `\n`;

        if (clients[0]) {
            for (let i = 0; i < clients.length; i++) {
                const client = clients[i];
                if (!client?.id) continue;

                if (!validateIpv4({ ip: client.wg_ip_address })) {
                    return {
                        success: false,
                        msg: `Invalid Private IP address for client ${client.id}: ${client.wg_ip_address}`,
                    };
                }

                const { CLIENT_PUBLIC_KEY_FILE } = grabClientDirnames({
                    client,
                    host_id: HOST_ID,
                });

                const CLIENT_PUBLIC_KEY = readFileSync(
                    CLIENT_PUBLIC_KEY_FILE,
                    "utf-8",
                ).trim();

                HOST_CONFIG += `[Peer]\n`;
                HOST_CONFIG += `PublicKey = ${CLIENT_PUBLIC_KEY}\n`;
                HOST_CONFIG += `AllowedIPs = ${client.wg_ip_address}/32\n`;
                HOST_CONFIG += `\n`;
            }
        }

        writeFileSync(HOST_CONFIG_FILE, HOST_CONFIG, { mode: 0o600 });
        chmodSync(HOST_CONFIG_FILE, 0o600);

        if (!is_update_after_client_setup) {
            execSync(
                `ip link set ${INTERFACE_NAME} down 2>/dev/null; ip link del ${INTERFACE_NAME} 2>/dev/null; true`,
                { encoding: "utf-8" },
            );
        }

        if (apply_live) {
            return applyWireguardHostLive({ host_id: Number(HOST_ID) });
        }

        const manage_res = manageWireguardHost({
            host_id: HOST_ID,
            action: "restart",
        });

        if (!manage_res.success) {
            return {
                success: false,
                msg: `Host config written to ${HOST_CONFIG_FILE}, but the tunnel could not be restarted: ${manage_res.msg}`,
            };
        }

        const firewall_res = await allowSubnetFirewallRules({
            subnet: HOST_SUBNET,
            listen_port: LISTEN_PORT,
        });

        if (!firewall_res.success) {
            console.log(
                `Warning setting up firewall rules for host ${HOST_ID} => ${firewall_res.msg}`,
            );
        }

        return {
            success: true,
            msg: [manage_res.msg, firewall_res.msg].join("\n\n"),
        };
    } catch (error: any) {
        console.log(
            `Error setting up wireguard host ${HOST_ID} =>`,
            error.message,
        );

        return {
            success: false,
            msg: error.message,
        };
    }
}