This commit is contained in:
Owen
2025-10-04 18:36:44 -07:00
parent 3123f858bb
commit c2c907852d
320 changed files with 35785 additions and 2984 deletions

View File

@@ -14,7 +14,7 @@ import {
import { clients, clientSites, Newt, sites } from "@server/db";
import { eq, and, inArray } from "drizzle-orm";
import { updatePeer } from "../olm/peers";
import { sendToExitNode } from "../../lib/exitNodeComms";
import { sendToExitNode } from "@server/lib/exitNodes";
const inputSchema = z.object({
publicKey: z.string(),
@@ -66,7 +66,7 @@ export const handleGetConfigMessage: MessageHandler = async (context) => {
// we need to wait for hole punch success
if (!existingSite.endpoint) {
logger.warn(`Site ${existingSite.siteId} has no endpoint, skipping`);
logger.debug(`In newt get config: existing site ${existingSite.siteId} has no endpoint, skipping`);
return;
}
@@ -74,12 +74,12 @@ export const handleGetConfigMessage: MessageHandler = async (context) => {
// TODO: somehow we should make sure a recent hole punch has happened if this occurs (hole punch could be from the last restart if done quickly)
}
if (existingSite.lastHolePunch && now - existingSite.lastHolePunch > 6) {
logger.warn(
`Site ${existingSite.siteId} last hole punch is too old, skipping`
);
return;
}
// if (existingSite.lastHolePunch && now - existingSite.lastHolePunch > 6) {
// logger.warn(
// `Site ${existingSite.siteId} last hole punch is too old, skipping`
// );
// return;
// }
// update the endpoint and the public key
const [site] = await db
@@ -176,7 +176,7 @@ export const handleGetConfigMessage: MessageHandler = async (context) => {
if (!endpoint) {
logger.warn(
`Site ${site.siteId} has no endpoint, skipping`
`In Newt get config: Peer site ${site.siteId} has no endpoint, skipping`
);
return null;
}

View File

@@ -29,7 +29,14 @@ export const handleNewtPingRequestMessage: MessageHandler = async (context) => {
.where(eq(sites.siteId, newt.siteId))
.limit(1);
const exitNodesList = await listExitNodes(site.orgId, true); // filter for only the online ones
if (!site || !site.orgId) {
logger.warn("Site not found");
return;
}
const { noCloud } = message.data;
const exitNodesList = await listExitNodes(site.orgId, true, noCloud || false); // filter for only the online ones
let lastExitNodeId = null;
if (newt.siteId) {

View File

@@ -1,6 +1,7 @@
import { db, newts } from "@server/db";
import { db, exitNodeOrgs, newts } from "@server/db";
import { MessageHandler } from "../ws";
import { exitNodes, Newt, resources, sites, Target, targets } from "@server/db";
import { targetHealthCheck } from "@server/db";
import { eq, and, sql, inArray } from "drizzle-orm";
import { addPeer, deletePeer } from "../gerbil/peers";
import logger from "@server/logger";
@@ -9,7 +10,12 @@ import {
findNextAvailableCidr,
getNextAvailableClientSubnet
} from "@server/lib/ip";
import { selectBestExitNode, verifyExitNodeOrgAccess } from "@server/lib/exitNodes";
import { usageService } from "@server/lib/private/billing/usageService";
import { FeatureId } from "@server/lib/private/billing";
import {
selectBestExitNode,
verifyExitNodeOrgAccess
} from "@server/lib/exitNodes";
import { fetchContainers } from "./dockerSocket";
export type ExitNodePingResult = {
@@ -22,6 +28,8 @@ export type ExitNodePingResult = {
wasPreviouslyConnected: boolean;
};
let numTimesLimitExceededForId: Record<string, number> = {};
export const handleNewtRegisterMessage: MessageHandler = async (context) => {
const { message, client, sendToClient } = context;
const newt = client as Newt;
@@ -86,13 +94,52 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
fetchContainers(newt.newtId);
}
const rejectSiteUptime = await usageService.checkLimitSet(
oldSite.orgId,
false,
FeatureId.SITE_UPTIME
);
const rejectEgressDataMb = await usageService.checkLimitSet(
oldSite.orgId,
false,
FeatureId.EGRESS_DATA_MB
);
// Do we need to check the users and domains daily limits here?
// const rejectUsers = await usageService.checkLimitSet(oldSite.orgId, false, FeatureId.USERS);
// const rejectDomains = await usageService.checkLimitSet(oldSite.orgId, false, FeatureId.DOMAINS);
// if (rejectEgressDataMb || rejectSiteUptime || rejectUsers || rejectDomains) {
if (rejectEgressDataMb || rejectSiteUptime) {
logger.info(
`Usage limits exceeded for org ${oldSite.orgId}. Rejecting newt registration.`
);
// PREVENT FURTHER REGISTRATION ATTEMPTS SO WE DON'T SPAM
// Increment the limit exceeded count for this site
numTimesLimitExceededForId[newt.newtId] =
(numTimesLimitExceededForId[newt.newtId] || 0) + 1;
if (numTimesLimitExceededForId[newt.newtId] > 15) {
logger.debug(
`Newt ${newt.newtId} has exceeded usage limits 15 times. Terminating...`
);
}
return;
}
let siteSubnet = oldSite.subnet;
let exitNodeIdToQuery = oldSite.exitNodeId;
if (exitNodeId && (oldSite.exitNodeId !== exitNodeId || !oldSite.subnet)) {
// This effectively moves the exit node to the new one
exitNodeIdToQuery = exitNodeId; // Use the provided exitNodeId if it differs from the site's exitNodeId
const { exitNode, hasAccess } = await verifyExitNodeOrgAccess(exitNodeIdToQuery, oldSite.orgId);
const { exitNode, hasAccess } = await verifyExitNodeOrgAccess(
exitNodeIdToQuery,
oldSite.orgId
);
if (!exitNode) {
logger.warn("Exit node not found");
@@ -114,6 +161,10 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
const blockSize = config.getRawConfig().gerbil.site_block_size;
const subnets = sitesQuery
.map((site) => site.subnet)
.filter(
(subnet) =>
subnet && /^(\d{1,3}\.){3}\d{1,3}\/\d{1,2}$/.test(subnet)
)
.filter((subnet) => subnet !== null);
subnets.push(exitNode.address.replace(/\/\d+$/, `/${blockSize}`));
const newSubnet = findNextAvailableCidr(
@@ -122,7 +173,9 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
exitNode.address
);
if (!newSubnet) {
logger.error("No available subnets found for the new exit node");
logger.error(
`No available subnets found for the new exit node id ${exitNodeId} and site id ${siteId}`
);
return;
}
@@ -168,11 +221,25 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
return;
}
// add the peer to the exit node
await addPeer(exitNodeIdToQuery, {
publicKey: publicKey,
allowedIps: [siteSubnet]
});
try {
// add the peer to the exit node
await addPeer(exitNodeIdToQuery, {
publicKey: publicKey,
allowedIps: [siteSubnet]
});
} catch (error) {
logger.error(`Failed to add peer to exit node: ${error}`);
}
if (newtVersion && newtVersion !== newt.version) {
// update the newt version in the database
await db
.update(newts)
.set({
version: newtVersion as string
})
.where(eq(newts.newtId, newt.newtId));
}
if (newtVersion && newtVersion !== newt.version) {
// update the newt version in the database
@@ -194,10 +261,25 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
port: targets.port,
internalPort: targets.internalPort,
enabled: targets.enabled,
protocol: resources.protocol
protocol: resources.protocol,
hcEnabled: targetHealthCheck.hcEnabled,
hcPath: targetHealthCheck.hcPath,
hcScheme: targetHealthCheck.hcScheme,
hcMode: targetHealthCheck.hcMode,
hcHostname: targetHealthCheck.hcHostname,
hcPort: targetHealthCheck.hcPort,
hcInterval: targetHealthCheck.hcInterval,
hcUnhealthyInterval: targetHealthCheck.hcUnhealthyInterval,
hcTimeout: targetHealthCheck.hcTimeout,
hcHeaders: targetHealthCheck.hcHeaders,
hcMethod: targetHealthCheck.hcMethod
})
.from(targets)
.innerJoin(resources, eq(targets.resourceId, resources.resourceId))
.leftJoin(
targetHealthCheck,
eq(targets.targetId, targetHealthCheck.targetId)
)
.where(and(eq(targets.siteId, siteId), eq(targets.enabled, true)));
const { tcpTargets, udpTargets } = allTargets.reduce(
@@ -222,6 +304,59 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
{ tcpTargets: [] as string[], udpTargets: [] as string[] }
);
const healthCheckTargets = allTargets.map((target) => {
// make sure the stuff is defined
if (
!target.hcPath ||
!target.hcHostname ||
!target.hcPort ||
!target.hcInterval ||
!target.hcMethod
) {
logger.debug(
`Skipping target ${target.targetId} due to missing health check fields`
);
return null; // Skip targets with missing health check fields
}
// parse headers
const hcHeadersParse = target.hcHeaders
? JSON.parse(target.hcHeaders)
: null;
let hcHeadersSend: { [key: string]: string } = {};
if (hcHeadersParse) {
hcHeadersParse.forEach(
(header: { name: string; value: string }) => {
hcHeadersSend[header.name] = header.value;
}
);
}
return {
id: target.targetId,
hcEnabled: target.hcEnabled,
hcPath: target.hcPath,
hcScheme: target.hcScheme,
hcMode: target.hcMode,
hcHostname: target.hcHostname,
hcPort: target.hcPort,
hcInterval: target.hcInterval, // in seconds
hcUnhealthyInterval: target.hcUnhealthyInterval, // in seconds
hcTimeout: target.hcTimeout, // in seconds
hcHeaders: hcHeadersSend,
hcMethod: target.hcMethod
};
});
// Filter out any null values from health check targets
const validHealthCheckTargets = healthCheckTargets.filter(
(target) => target !== null
);
logger.debug(
`Sending health check targets to newt ${newt.newtId}: ${JSON.stringify(validHealthCheckTargets)}`
);
return {
message: {
type: "newt/wg/connect",
@@ -233,10 +368,11 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
targets: {
udp: udpTargets,
tcp: tcpTargets
}
},
healthCheckTargets: validHealthCheckTargets
}
},
broadcast: false, // Send to all clients
excludeSender: false // Include sender in broadcast
};
};
};

View File

@@ -1,10 +1,12 @@
import { Target } from "@server/db";
import { Target, TargetHealthCheck, db, targetHealthCheck } from "@server/db";
import { sendToClient } from "../ws";
import logger from "@server/logger";
import { eq, inArray } from "drizzle-orm";
export async function addTargets(
newtId: string,
targets: Target[],
healthCheckData: TargetHealthCheck[],
protocol: string,
port: number | null = null
) {
@@ -21,6 +23,62 @@ export async function addTargets(
targets: payloadTargets
}
});
// Create a map for quick lookup
const healthCheckMap = new Map<number, TargetHealthCheck>();
healthCheckData.forEach(hc => {
healthCheckMap.set(hc.targetId, hc);
});
const healthCheckTargets = targets.map((target) => {
const hc = healthCheckMap.get(target.targetId);
// If no health check data found, skip this target
if (!hc) {
logger.warn(`No health check configuration found for target ${target.targetId}`);
return null;
}
// Ensure all necessary fields are present
if (!hc.hcPath || !hc.hcHostname || !hc.hcPort || !hc.hcInterval || !hc.hcMethod) {
logger.debug(`Skipping target ${target.targetId} due to missing health check fields`);
return null; // Skip targets with missing health check fields
}
const hcHeadersParse = hc.hcHeaders ? JSON.parse(hc.hcHeaders) : null;
const hcHeadersSend: { [key: string]: string } = {};
if (hcHeadersParse) {
// transform
hcHeadersParse.forEach((header: { name: string; value: string }) => {
hcHeadersSend[header.name] = header.value;
});
}
return {
id: target.targetId,
hcEnabled: hc.hcEnabled,
hcPath: hc.hcPath,
hcScheme: hc.hcScheme,
hcMode: hc.hcMode,
hcHostname: hc.hcHostname,
hcPort: hc.hcPort,
hcInterval: hc.hcInterval, // in seconds
hcUnhealthyInterval: hc.hcUnhealthyInterval, // in seconds
hcTimeout: hc.hcTimeout, // in seconds
hcHeaders: hcHeadersSend,
hcMethod: hc.hcMethod
};
});
// Filter out any null values from health check targets
const validHealthCheckTargets = healthCheckTargets.filter((target) => target !== null);
await sendToClient(newtId, {
type: `newt/healthcheck/add`,
data: {
targets: validHealthCheckTargets
}
});
}
export async function removeTargets(
@@ -42,4 +100,15 @@ export async function removeTargets(
targets: payloadTargets
}
});
const healthCheckTargets = targets.map((target) => {
return target.targetId
});
await sendToClient(newtId, {
type: `newt/healthcheck/remove`,
data: {
ids: healthCheckTargets
}
});
}