Working on handlers
This commit is contained in:
Owen
parent
e112fcba29
commit
b9de0f8e38
5 changed files with 62 additions and 171 deletions
|
|
@ -16,7 +16,7 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
|
|||
|
||||
const newt = client;
|
||||
|
||||
logger.info("Handling register message!");
|
||||
logger.info("Handling register newt message!");
|
||||
|
||||
if (!newt) {
|
||||
logger.warn("Newt not found");
|
||||
|
|
|
|||
46
server/routers/newt/peers.ts
Normal file
46
server/routers/newt/peers.ts
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
import db from '@server/db';
|
||||
import { newts, sites } from '@server/db/schema';
|
||||
import { eq } from 'drizzle-orm';
|
||||
import { sendToClient } from '../ws';
|
||||
|
||||
export async function addPeer(siteId: number, peer: {
|
||||
publicKey: string;
|
||||
allowedIps: string[];
|
||||
}) {
|
||||
|
||||
const [site] = await db.select().from(sites).where(eq(sites.siteId, siteId)).limit(1);
|
||||
if (!site) {
|
||||
throw new Error(`Exit node with ID ${siteId} not found`);
|
||||
}
|
||||
|
||||
// get the newt on the site
|
||||
const [newt] = await db.select().from(newts).where(eq(newts.siteId, siteId)).limit(1);
|
||||
if (!newt) {
|
||||
throw new Error(`Newt not found for site ${siteId}`);
|
||||
}
|
||||
|
||||
sendToClient(newt.newtId, {
|
||||
type: 'add_peer',
|
||||
data: peer
|
||||
});
|
||||
}
|
||||
|
||||
export async function deletePeer(siteId: number, publicKey: string) {
|
||||
const [site] = await db.select().from(sites).where(eq(sites.siteId, siteId)).limit(1);
|
||||
if (!site) {
|
||||
throw new Error(`Exit node with ID ${siteId} not found`);
|
||||
}
|
||||
|
||||
// get the newt on the site
|
||||
const [newt] = await db.select().from(newts).where(eq(newts.siteId, siteId)).limit(1);
|
||||
if (!newt) {
|
||||
throw new Error(`Newt not found for site ${siteId}`);
|
||||
}
|
||||
|
||||
sendToClient(newt.newtId, {
|
||||
type: 'delete_peer',
|
||||
data: {
|
||||
publicKey
|
||||
}
|
||||
});
|
||||
}
|
||||
|
|
@ -1,147 +0,0 @@
|
|||
import { z } from "zod";
|
||||
import { MessageHandler } from "../ws";
|
||||
import logger from "@server/logger";
|
||||
import { fromError } from "zod-validation-error";
|
||||
import db from "@server/db";
|
||||
import { olms, Site, sites } from "@server/db/schema";
|
||||
import { eq, isNotNull } from "drizzle-orm";
|
||||
import { findNextAvailableCidr } from "@server/lib/ip";
|
||||
import config from "@server/lib/config";
|
||||
|
||||
const inputSchema = z.object({
|
||||
publicKey: z.string(),
|
||||
endpoint: z.string(),
|
||||
listenPort: z.number()
|
||||
});
|
||||
|
||||
type Input = z.infer<typeof inputSchema>;
|
||||
|
||||
export const handleGetConfigMessage: MessageHandler = async (context) => {
|
||||
const { message, newt, sendToClient } = context;
|
||||
|
||||
logger.debug("Handling Newt get config message!");
|
||||
|
||||
if (!newt) {
|
||||
logger.warn("Newt not found");
|
||||
return;
|
||||
}
|
||||
|
||||
if (!newt.siteId) {
|
||||
logger.warn("Newt has no site!"); // TODO: Maybe we create the site here?
|
||||
return;
|
||||
}
|
||||
|
||||
const parsed = inputSchema.safeParse(message.data);
|
||||
if (!parsed.success) {
|
||||
logger.error(
|
||||
"handleGetConfigMessage: Invalid input: " +
|
||||
fromError(parsed.error).toString()
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
const { publicKey, endpoint, listenPort } = message.data as Input;
|
||||
|
||||
const siteId = newt.siteId;
|
||||
|
||||
const [siteRes] = await db
|
||||
.select()
|
||||
.from(sites)
|
||||
.where(eq(sites.siteId, siteId));
|
||||
|
||||
if (!siteRes) {
|
||||
logger.warn("handleGetConfigMessage: Site not found");
|
||||
return;
|
||||
}
|
||||
|
||||
let site: Site | undefined;
|
||||
if (!site) {
|
||||
const address = await getNextAvailableSubnet();
|
||||
|
||||
// create a new exit node
|
||||
const [updateRes] = await db
|
||||
.update(sites)
|
||||
.set({
|
||||
publicKey,
|
||||
endpoint,
|
||||
address,
|
||||
listenPort
|
||||
})
|
||||
.where(eq(sites.siteId, siteId))
|
||||
.returning();
|
||||
|
||||
site = updateRes;
|
||||
|
||||
logger.info(`Updated site ${siteId} with new WG Newt info`);
|
||||
} else {
|
||||
site = siteRes;
|
||||
}
|
||||
|
||||
if (!site) {
|
||||
logger.error("handleGetConfigMessage: Failed to update site");
|
||||
return;
|
||||
}
|
||||
|
||||
const clientsRes = await db
|
||||
.select()
|
||||
.from(olms)
|
||||
.where(eq(olms.siteId, siteId));
|
||||
|
||||
const peers = await Promise.all(
|
||||
clientsRes.map(async (client) => {
|
||||
return {
|
||||
publicKey: client.pubKey,
|
||||
allowedIps: "0.0.0.0/0"
|
||||
};
|
||||
})
|
||||
);
|
||||
|
||||
const configResponse = {
|
||||
listenPort: site.listenPort, // ?????
|
||||
// ipAddress: exitNode[0].address,
|
||||
peers
|
||||
};
|
||||
|
||||
logger.debug("Sending config: ", configResponse);
|
||||
|
||||
return {
|
||||
message: {
|
||||
type: "olm/wg/connect", // what to make the response type?
|
||||
data: {
|
||||
config: configResponse
|
||||
}
|
||||
},
|
||||
broadcast: false, // Send to all clients
|
||||
excludeSender: false // Include sender in broadcast
|
||||
};
|
||||
};
|
||||
|
||||
async function getNextAvailableSubnet(): Promise<string> {
|
||||
const existingAddresses = await db
|
||||
.select({
|
||||
address: sites.address
|
||||
})
|
||||
.from(sites)
|
||||
.where(isNotNull(sites.address));
|
||||
|
||||
const addresses = existingAddresses
|
||||
.map((a) => a.address)
|
||||
.filter((a) => a) as string[];
|
||||
|
||||
let subnet = findNextAvailableCidr(
|
||||
addresses,
|
||||
config.getRawConfig().wg_site.block_size,
|
||||
config.getRawConfig().wg_site.subnet_group
|
||||
);
|
||||
if (!subnet) {
|
||||
throw new Error("No available subnets remaining in space");
|
||||
}
|
||||
|
||||
// replace the last octet with 1
|
||||
subnet =
|
||||
subnet.split(".").slice(0, 3).join(".") +
|
||||
".1" +
|
||||
"/" +
|
||||
subnet.split("/")[1];
|
||||
return subnet;
|
||||
}
|
||||
|
|
@ -1,14 +1,11 @@
|
|||
import db from "@server/db";
|
||||
import { MessageHandler } from "../ws";
|
||||
import {
|
||||
exitNodes,
|
||||
resources,
|
||||
olms,
|
||||
sites,
|
||||
Target,
|
||||
targets
|
||||
} from "@server/db/schema";
|
||||
import { eq, and, sql } from "drizzle-orm";
|
||||
import { addPeer, deletePeer } from "../gerbil/peers";
|
||||
import { eq, } from "drizzle-orm";
|
||||
import { addPeer, deletePeer } from "../newt/peers";
|
||||
import logger from "@server/logger";
|
||||
|
||||
export const handleOlmRegisterMessage: MessageHandler = async (context) => {
|
||||
|
|
@ -16,7 +13,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
|
|||
|
||||
const olm = client;
|
||||
|
||||
logger.info("Handling register message!");
|
||||
logger.info("Handling register olm message!");
|
||||
|
||||
if (!olm) {
|
||||
logger.warn("Olm not found");
|
||||
|
|
@ -42,28 +39,22 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
|
|||
.where(eq(sites.siteId, siteId))
|
||||
.limit(1);
|
||||
|
||||
if (!site || !site.exitNodeId) {
|
||||
if (!site) {
|
||||
logger.warn("Site not found or does not have exit node");
|
||||
return;
|
||||
}
|
||||
|
||||
await db
|
||||
.update(sites)
|
||||
.update(olms)
|
||||
.set({
|
||||
pubKey: publicKey
|
||||
})
|
||||
.where(eq(sites.siteId, siteId))
|
||||
.where(eq(olms.olmId, olm.olmId))
|
||||
.returning();
|
||||
|
||||
const [exitNode] = await db
|
||||
.select()
|
||||
.from(exitNodes)
|
||||
.where(eq(exitNodes.exitNodeId, site.exitNodeId))
|
||||
.limit(1);
|
||||
|
||||
if (site.pubKey && site.pubKey !== publicKey) {
|
||||
if (olm.pubKey && olm.pubKey !== publicKey) {
|
||||
logger.info("Public key mismatch. Deleting old peer...");
|
||||
await deletePeer(site.exitNodeId, site.pubKey);
|
||||
await deletePeer(site.siteId, site.pubKey);
|
||||
}
|
||||
|
||||
if (!site.subnet) {
|
||||
|
|
@ -72,7 +63,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
|
|||
}
|
||||
|
||||
// add the peer to the exit node
|
||||
await addPeer(site.exitNodeId, {
|
||||
await addPeer(site.siteId, {
|
||||
publicKey: publicKey,
|
||||
allowedIps: [site.subnet]
|
||||
});
|
||||
|
|
@ -81,9 +72,9 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
|
|||
message: {
|
||||
type: "olm/wg/connect",
|
||||
data: {
|
||||
endpoint: `${exitNode.endpoint}:${exitNode.listenPort}`,
|
||||
publicKey: exitNode.publicKey,
|
||||
serverIP: exitNode.address.split("/")[0],
|
||||
endpoint: `${site.endpoint}:${site.listenPort}`,
|
||||
publicKey: site.publicKey,
|
||||
serverIP: site.address!.split("/")[0],
|
||||
tunnelIP: site.subnet.split("/")[0]
|
||||
}
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1 +1,2 @@
|
|||
export * from "./pickOlmDefaults";
|
||||
export * from "./pickOlmDefaults";
|
||||
export * from "./handleOlmRegisterMessage";
|
||||
Loading…
Add table
Reference in a new issue