Updates

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

@@ -1,5 +1,6 @@
 import type {
     BUN_SQLITE_WGUI_CLIENTS,
+    BUN_SQLITE_WGUI_CLIENT_RULES,
     BUN_SQLITE_WGUI_HOSTS,
     BUN_SQLITE_WGUI_VARIABLES,
 } from "@/db/types/db";
@@ -12,8 +13,8 @@ import type { TableType, User } from "@/src/types";
 import checkPrivateIPAvailability from "./check-private-ip-availability";
 import manageWireguardHost from "./manage-wireguard-host";
 import grabHostDirnames from "./grab-host-dir-names";
-import path from "node:path";
 import grabClientDirnames from "./grab-client-dir-names";
+import buildHostIptablesScripts from "./build-host-iptables-scripts";
 
 const { WIREGUARD_PRIVATE_KEY_FILE_NAME, WIREGUARD_PUBLIC_KEY_FILE_NAME } =
     grabDirNames();
@@ -69,6 +70,36 @@ export default async function setupWireguardHost({
 
     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 TARGET_INTERFACE = await grabHostNetworkInterface();
     const HOST_WG_IP =
         host?.wg_ip_address ||
@@ -172,37 +203,37 @@ export default async function setupWireguardHost({
 
         sh += `cd ${HOST_CONFIG_DIR}\n`;
 
-        sh += `cat > ${POST_UP_PATH} << EOF\n`;
-        sh += `#!/bin/bash\n\n`;
-        sh += `# Allow WireGuard traffic to/from the server itself\n`;
-        sh += `iptables -I INPUT 1 -i ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `iptables -I OUTPUT 1 -o ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `\n`;
-        sh += `# Allow WireGuard traffic to be forwarded (insert above Docker rules)\n`;
-        sh += `iptables -I FORWARD 1 -i ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `iptables -I FORWARD 1 -o ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `\n`;
-        sh += `iptables -t nat -A POSTROUTING -o ${TARGET_INTERFACE} -j MASQUERADE\n`;
-        sh += `EOF\n`;
-        sh += `chmod +x ${POST_UP_PATH}\n`;
+        const iptables_scripts = buildHostIptablesScripts({
+            host_id: Number(HOST_ID),
+            interface_name: INTERFACE_NAME,
+            target_interface: TARGET_INTERFACE || "eth0",
+            clients: clients.map((client) => ({
+                id: client.id,
+                wg_ip_address: client.wg_ip_address,
+                rules:
+                    client.id && typeof client.id == "number"
+                        ? rules_by_client_id.get(client.id) || []
+                        : [],
+            })),
+        });
 
-        sh += `\n`;
+        if (
+            !iptables_scripts.success ||
+            !iptables_scripts.post_up ||
+            !iptables_scripts.post_down
+        ) {
+            return {
+                success: false,
+                msg: iptables_scripts.msg || `Could not build iptables scripts`,
+            };
+        }
 
-        sh += `cat > ${POST_DOWN_PATH} << EOF\n`;
-        sh += `#!/bin/bash\n\n`;
-        sh += `# Remove WireGuard INPUT/OUTPUT rules\n`;
-        sh += `iptables -D INPUT -i ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `iptables -D OUTPUT -o ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `\n`;
-        sh += `# Remove FORWARD rules\n`;
-        sh += `iptables -D FORWARD -i ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `iptables -D FORWARD -o ${INTERFACE_NAME} -j ACCEPT\n`;
-        sh += `\n`;
-        sh += `iptables -t nat -D POSTROUTING -o ${TARGET_INTERFACE} -j MASQUERADE\n`;
-        sh += `EOF\n`;
-        sh += `chmod +x ${POST_DOWN_PATH}\n`;
-
-        sh += `\n`;
+        execSync(`mkdir -p ${HOST_IPTABLES_DIR}`, { encoding: "utf-8" });
+        await Bun.write(POST_UP_PATH, iptables_scripts.post_up);
+        await Bun.write(POST_DOWN_PATH, iptables_scripts.post_down);
+        execSync(`chmod +x ${POST_UP_PATH} ${POST_DOWN_PATH}`, {
+            encoding: "utf-8",
+        });
 
         sh += `cat > ${INTERFACE_NAME}.conf << EOF\n`;
         sh += `[Interface]\n`;
@@ -218,14 +249,21 @@ export default async function setupWireguardHost({
                 const client = clients[i];
                 if (!client?.id) continue;
 
-                const { CLIENT_PUBLIC_KEY } = grabClientDirnames({
+                const { CLIENT_PUBLIC_KEY_FILE } = grabClientDirnames({
                     client,
                     host_id: HOST_ID,
                 });
 
+                const CLIENT_PUBLIC_KEY = execSync(
+                    `cat ${CLIENT_PUBLIC_KEY_FILE}`,
+                    {
+                        encoding: "utf-8",
+                    },
+                ).trim();
+
                 sh += `[Peer]\n`;
                 sh += `PublicKey = ${CLIENT_PUBLIC_KEY}\n`;
-                sh += `AllowedIPs = ${client.allowed_ips}\n`;
+                sh += `AllowedIPs = ${client.wg_ip_address}/32\n`;
                 sh += `\n`;
             }
         }