From bec303821b8b05e67c1a9ee87277e787273ba5a5 Mon Sep 17 00:00:00 2001 From: Owen Date: Fri, 21 Feb 2025 12:17:56 -0500 Subject: [PATCH 1/4] Handle types in handlers --- server/routers/newt/handleGetConfigMessage.ts | 5 +-- .../routers/newt/handleNewtRegisterMessage.ts | 4 +-- .../routers/olm/handleOlmRegisterMessage.ts | 34 +++++++++++++------ 3 files changed, 28 insertions(+), 15 deletions(-) diff --git a/server/routers/newt/handleGetConfigMessage.ts b/server/routers/newt/handleGetConfigMessage.ts index 17ac63dd..9bcde107 100644 --- a/server/routers/newt/handleGetConfigMessage.ts +++ b/server/routers/newt/handleGetConfigMessage.ts @@ -3,7 +3,7 @@ import { MessageHandler } from "../ws"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import db from "@server/db"; -import { clients, Site, sites } from "@server/db/schema"; +import { clients, Newt, Site, sites } from "@server/db/schema"; import { eq, isNotNull } from "drizzle-orm"; import { findNextAvailableCidr } from "@server/lib/ip"; import config from "@server/lib/config"; @@ -17,7 +17,8 @@ const inputSchema = z.object({ type Input = z.infer; export const handleGetConfigMessage: MessageHandler = async (context) => { - const { message, newt, sendToClient } = context; + const { message, client, sendToClient } = context; + const newt = client as Newt; logger.debug("Handling Newt get config message!"); diff --git a/server/routers/newt/handleNewtRegisterMessage.ts b/server/routers/newt/handleNewtRegisterMessage.ts index 54a62735..21ad4ba6 100644 --- a/server/routers/newt/handleNewtRegisterMessage.ts +++ b/server/routers/newt/handleNewtRegisterMessage.ts @@ -2,6 +2,7 @@ import db from "@server/db"; import { MessageHandler } from "../ws"; import { exitNodes, + Newt, resources, sites, Target, @@ -13,8 +14,7 @@ import logger from "@server/logger"; export const handleNewtRegisterMessage: MessageHandler = async (context) => { const { message, client, sendToClient } = context; - - const newt = client; + const newt = client as Newt; logger.info("Handling register newt message!"); diff --git a/server/routers/olm/handleOlmRegisterMessage.ts b/server/routers/olm/handleOlmRegisterMessage.ts index 859f756c..147de81c 100644 --- a/server/routers/olm/handleOlmRegisterMessage.ts +++ b/server/routers/olm/handleOlmRegisterMessage.ts @@ -1,6 +1,8 @@ import db from "@server/db"; import { MessageHandler } from "../ws"; import { + clients, + Olm, olms, sites, } from "@server/db/schema"; @@ -9,9 +11,8 @@ import { addPeer, deletePeer } from "../newt/peers"; import logger from "@server/logger"; export const handleOlmRegisterMessage: MessageHandler = async (context) => { - const { message, client, sendToClient } = context; - - const olm = client; + const { message, client: c, sendToClient } = context; + const olm = c as Olm; logger.info("Handling register olm message!"); @@ -20,12 +21,12 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { return; } - if (!olm.siteId) { + if (!olm.clientId) { logger.warn("Olm has no site!"); // TODO: Maybe we create the site here? return; } - const siteId = olm.siteId; + const clientId = olm.clientId; const { publicKey } = message.data; if (!publicKey) { @@ -33,28 +34,39 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => { return; } + const [client] = await db + .select() + .from(clients) + .where(eq(clients.clientId, clientId)) + .limit(1); + + if (!client || !client.siteId) { + logger.warn("Site not found or does not have exit node"); + return; + } + const [site] = await db .select() .from(sites) - .where(eq(sites.siteId, siteId)) + .where(eq(sites.siteId, client.siteId)) .limit(1); - if (!site) { + if (!client) { logger.warn("Site not found or does not have exit node"); return; } await db - .update(olms) + .update(clients) .set({ pubKey: publicKey }) - .where(eq(olms.olmId, olm.olmId)) + .where(eq(clients.clientId, olm.clientId)) .returning(); - if (olm.pubKey && olm.pubKey !== publicKey) { + if (client.pubKey && client.pubKey !== publicKey) { logger.info("Public key mismatch. Deleting old peer..."); - await deletePeer(site.siteId, site.pubKey); + await deletePeer(site.siteId, client.pubKey); } if (!site.subnet) { From ef69bf9256cedbfcb0555523215224b15beec6ae Mon Sep 17 00:00:00 2001 From: Owen Date: Fri, 21 Feb 2025 12:34:05 -0500 Subject: [PATCH 2/4] Update ws paths --- server/routers/newt/handleGetConfigMessage.ts | 4 ++-- server/routers/newt/peers.ts | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/server/routers/newt/handleGetConfigMessage.ts b/server/routers/newt/handleGetConfigMessage.ts index 9bcde107..4f03bdd8 100644 --- a/server/routers/newt/handleGetConfigMessage.ts +++ b/server/routers/newt/handleGetConfigMessage.ts @@ -92,7 +92,7 @@ export const handleGetConfigMessage: MessageHandler = async (context) => { clientsRes.map(async (client) => { return { publicKey: client.pubKey, - allowedIps: "0.0.0.0/0" + allowedIps: "0.0.0.0/0" // TODO: We should lock this down more }; }) ); @@ -107,7 +107,7 @@ export const handleGetConfigMessage: MessageHandler = async (context) => { return { message: { - type: "newt/wg/connect", // what to make the response type? + type: "newt/wg/receive-config", // what to make the response type? data: { config: configResponse } diff --git a/server/routers/newt/peers.ts b/server/routers/newt/peers.ts index ee22c052..afc3b5d6 100644 --- a/server/routers/newt/peers.ts +++ b/server/routers/newt/peers.ts @@ -20,7 +20,7 @@ export async function addPeer(siteId: number, peer: { } sendToClient(newt.newtId, { - type: 'add_peer', + type: 'newt/wg/peer/add', data: peer }); } @@ -38,7 +38,7 @@ export async function deletePeer(siteId: number, publicKey: string) { } sendToClient(newt.newtId, { - type: 'delete_peer', + type: 'newt/wg/peer/remove', data: { publicKey } From a57d32d05d04db68d63682fa2df080f6d33e2cd5 Mon Sep 17 00:00:00 2001 From: Owen Date: Fri, 21 Feb 2025 12:41:02 -0500 Subject: [PATCH 3/4] Add receive bandwidth --- server/routers/gerbil/receiveBandwidth.ts | 9 +-- server/routers/messageHandlers.ts | 3 +- .../newt/handleReceiveBandwidthMessage.ts | 68 +++++++++++++++++++ server/routers/newt/index.ts | 3 +- 4 files changed, 77 insertions(+), 6 deletions(-) create mode 100644 server/routers/newt/handleReceiveBandwidthMessage.ts diff --git a/server/routers/gerbil/receiveBandwidth.ts b/server/routers/gerbil/receiveBandwidth.ts index b2063c08..b8839577 100644 --- a/server/routers/gerbil/receiveBandwidth.ts +++ b/server/routers/gerbil/receiveBandwidth.ts @@ -30,12 +30,13 @@ export const receiveBandwidth = async ( const { publicKey, bytesIn, bytesOut } = peer; // Find the site by public key - const site = await trx.query.sites.findFirst({ - where: eq(sites.pubKey, publicKey) - }); + const [site] = await trx + .select() + .from(sites) + .where(eq(sites.pubKey, publicKey)) + .limit(1); if (!site) { - logger.warn(`Site not found for public key: ${publicKey}`); continue; } let online = site.online; diff --git a/server/routers/messageHandlers.ts b/server/routers/messageHandlers.ts index bf8f357c..f23ea0a8 100644 --- a/server/routers/messageHandlers.ts +++ b/server/routers/messageHandlers.ts @@ -1,4 +1,4 @@ -import { handleNewtRegisterMessage } from "./newt"; +import { handleNewtRegisterMessage, handleReceiveBandwidthMessage } from "./newt"; import { handleOlmRegisterMessage } from "./olm"; import { handleGetConfigMessage } from "./newt/handleGetConfigMessage"; import { MessageHandler } from "./ws"; @@ -7,4 +7,5 @@ export const messageHandlers: Record = { "newt/wg/register": handleNewtRegisterMessage, "olm/wg/register": handleOlmRegisterMessage, "newt/wg/get-config": handleGetConfigMessage, + "newt/receive-bandwidth": handleReceiveBandwidthMessage }; diff --git a/server/routers/newt/handleReceiveBandwidthMessage.ts b/server/routers/newt/handleReceiveBandwidthMessage.ts new file mode 100644 index 00000000..1e1642aa --- /dev/null +++ b/server/routers/newt/handleReceiveBandwidthMessage.ts @@ -0,0 +1,68 @@ +import db from "@server/db"; +import { MessageHandler } from "../ws"; +import { clients, Newt } from "@server/db/schema"; +import { eq } from "drizzle-orm"; +import logger from "@server/logger"; + +interface PeerBandwidth { + publicKey: string; + bytesIn: number; + bytesOut: number; +} + +export const handleReceiveBandwidthMessage: MessageHandler = async (context) => { + const { message, client, sendToClient } = context; + const newt = client as Newt; + + const bandwidthData: PeerBandwidth[] = message.data; + + if (!Array.isArray(bandwidthData)) { + throw new Error("Invalid bandwidth data"); + } + + await db.transaction(async (trx) => { + for (const peer of bandwidthData) { + const { publicKey, bytesIn, bytesOut } = peer; + + // Find the site by public key + const [client] = await trx + .select() + .from(clients) + .where(eq(clients.pubKey, publicKey)) + .limit(1); + + if (!client) { + continue; + } + let online = client.online; + + // if the bandwidth for the site is > 0 then set it to online. if it has been less than 0 (no update) for 5 minutes then set it to offline + if (bytesIn > 0 || bytesOut > 0) { + online = true; + } else if (client.lastBandwidthUpdate) { + const lastBandwidthUpdate = new Date( + client.lastBandwidthUpdate + ); + const currentTime = new Date(); + const diff = + currentTime.getTime() - lastBandwidthUpdate.getTime(); + if (diff < 300000) { + online = false; + } + } + + // Update the site's bandwidth usage + await trx + .update(clients) + .set({ + megabytesOut: (client.megabytesIn || 0) + bytesIn, + megabytesIn: (client.megabytesOut || 0) + bytesOut, + lastBandwidthUpdate: new Date().toISOString(), + online + }) + .where(eq(clients.clientId, client.clientId)); + } + }); + + logger.info("Handling register olm message!"); +}; diff --git a/server/routers/newt/index.ts b/server/routers/newt/index.ts index dcc49749..84b9a6e9 100644 --- a/server/routers/newt/index.ts +++ b/server/routers/newt/index.ts @@ -1,3 +1,4 @@ export * from "./createNewt"; export * from "./getToken"; -export * from "./handleRegisterMessage"; \ No newline at end of file +export * from "./handleNewtRegisterMessage"; +export* from "./handleReceiveBandwidthMessage"; \ No newline at end of file From 757d628bc84c729e1ca5e103a714b8fc6b44c5f9 Mon Sep 17 00:00:00 2001 From: Owen Date: Fri, 21 Feb 2025 12:52:24 -0500 Subject: [PATCH 4/4] Handle port correctly --- install/fs/config.yml | 5 ++++ server/lib/config.ts | 2 ++ server/routers/newt/handleGetConfigMessage.ts | 24 ++++++++++++++++++- 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/install/fs/config.yml b/install/fs/config.yml index 8e4411e7..cf9e6464 100644 --- a/install/fs/config.yml +++ b/install/fs/config.yml @@ -32,6 +32,11 @@ gerbil: site_block_size: 30 subnet_group: 100.89.137.0/20 +wg_site: + start_port: 51820 + block_size: 24 + subnet_group: 100.89.137.0/20 + rate_limits: global: window_minutes: 1 diff --git a/server/lib/config.ts b/server/lib/config.ts index fc1c0531..f607fe0d 100644 --- a/server/lib/config.ts +++ b/server/lib/config.ts @@ -12,6 +12,7 @@ import { } from "@server/lib/consts"; import { passwordSchema } from "@server/auth/passwordSchema"; import stoi from "./stoi"; +import { start } from "repl"; const portSchema = z.number().positive().gt(0).lte(65535); const hostnameSchema = z @@ -112,6 +113,7 @@ const configSchema = z.object({ wg_site: z.object({ block_size: z.number().positive().gt(0), subnet_group: z.string(), + start_port: portSchema }), rate_limits: z.object({ global: z.object({ diff --git a/server/routers/newt/handleGetConfigMessage.ts b/server/routers/newt/handleGetConfigMessage.ts index 4f03bdd8..6d8cb8c8 100644 --- a/server/routers/newt/handleGetConfigMessage.ts +++ b/server/routers/newt/handleGetConfigMessage.ts @@ -41,7 +41,7 @@ export const handleGetConfigMessage: MessageHandler = async (context) => { return; } - const { publicKey, endpoint, listenPort } = message.data as Input; + const { publicKey, endpoint } = message.data as Input; const siteId = newt.siteId; @@ -58,6 +58,7 @@ export const handleGetConfigMessage: MessageHandler = async (context) => { let site: Site | undefined; if (!site) { const address = await getNextAvailableSubnet(); + const listenPort = await getNextAvailablePort(); // create a new exit node const [updateRes] = await db @@ -146,3 +147,24 @@ async function getNextAvailableSubnet(): Promise { subnet.split("/")[1]; return subnet; } + +async function getNextAvailablePort(): Promise { + // Get all existing ports from exitNodes table + const existingPorts = await db.select({ + listenPort: sites.listenPort, + }).from(sites); + + // Find the first available port between 1024 and 65535 + let nextPort = config.getRawConfig().wg_site.start_port; + for (const port of existingPorts) { + if (port.listenPort && port.listenPort > nextPort) { + break; + } + nextPort++; + if (nextPort > 65535) { + throw new Error('No available ports remaining in space'); + } + } + + return nextPort; +}