Improve latency on register messages for newt

This commit is contained in:
Owen
2026-10-22 16:11:18 -04:00
parent 5ff8ecdf78
commit 16a9d99b20
4 changed files with 220 additions and 144 deletions
+4 -11
View File
@@ -492,8 +492,8 @@ export async function getNextAvailableOrgSubnet(): Promise<{
}
}
const remoteSubnetCidrSchema = z.union([z.cidrv4(), z.cidrv6()]);
const remoteSubnetIpSchema = z.union([z.ipv4(), z.ipv6()]);
const cidrSchema = z.union([z.cidrv4(), z.cidrv6()]);
const ipSchema = z.union([z.ipv4(), z.ipv6()]);
export function generateRemoteSubnets(
allSiteResources: SiteResource[]
@@ -506,16 +506,12 @@ export function generateRemoteSubnets(
// we purposely filter out the gateway resource here because we add it manually on the client
if (sr.mode === "cidr") {
// check if its a valid CIDR using zod
const parseResult = remoteSubnetCidrSchema.safeParse(
sr.destination
);
const parseResult = cidrSchema.safeParse(sr.destination);
return parseResult.success;
}
if (sr.mode === "host" || sr.mode === "ssh") {
// check if its a valid IP using zod
const parseResult = remoteSubnetIpSchema.safeParse(
sr.destination
);
const parseResult = ipSchema.safeParse(sr.destination);
return parseResult.success;
}
return false;
@@ -603,7 +599,6 @@ export function generateSubnetProxyTargets(
if (siteResource.mode == "host" || siteResource.mode == "ssh") {
let destination = siteResource.destination;
// check if this is a valid ip
const ipSchema = z.union([z.ipv4(), z.ipv6()]);
if (ipSchema.safeParse(destination).success) {
destination = `${destination}/32`;
@@ -790,7 +785,6 @@ export async function generateSubnetProxyTargetV2(
if (siteResource.mode == "host" || siteResource.mode == "ssh") {
let destination = siteResource.destination;
// check if this is a valid ip
const ipSchema = z.union([z.ipv4(), z.ipv6()]);
if (ipSchema.safeParse(destination).success) {
destination = `${destination}/32`;
@@ -825,7 +819,6 @@ export async function generateSubnetProxyTargetV2(
} else if (siteResource.mode == "http") {
let destination = siteResource.destination;
// check if this is a valid ip
const ipSchema = z.union([z.ipv4(), z.ipv6()]);
if (ipSchema.safeParse(destination).success) {
destination = `${destination}/32`;
}
+151 -125
View File
@@ -5,6 +5,7 @@ import {
db,
ExitNode,
networks,
olms,
remoteExitNodeResources,
resources,
Site,
@@ -14,7 +15,7 @@ import {
targets
} from "@server/db";
import logger from "@server/logger";
import { initPeerAddHandshake, updatePeer } from "../olm/peers";
import { initPeerAddHandshakeBatch, updatePeersBatch } from "../olm/peers";
import { eq, and, inArray, or, isNotNull, sql } from "drizzle-orm";
import config from "@server/lib/config";
import { decrypt } from "@server/lib/crypto";
@@ -35,8 +36,10 @@ export async function buildClientConfigurationForNewtClient(
) {
const siteId = site.siteId;
const sendPeers = !!(site.publicKey && site.endpoint && exitNode);
// Get all clients connected to this site
const clientsRes = await db
const clientsQuery = db
.select()
.from(clients)
.innerJoin(
@@ -45,114 +48,25 @@ export async function buildClientConfigurationForNewtClient(
)
.where(eq(clientSitesAssociationsCache.siteId, siteId));
let peers: Array<{
publicKey: string;
allowedIps: string[];
endpoint?: string;
}> = [];
if (site.publicKey && site.endpoint && exitNode) {
// Prepare peers data for the response
peers = await Promise.all(
clientsRes
.filter((client) => {
if (!client.clients.pubKey) {
logger.debug(
`Client ${client.clients.clientId} has no public key, skipping`
);
return false;
}
if (!client.clients.subnet) {
logger.debug(
`Client ${client.clients.clientId} has no subnet, skipping`
);
return false;
}
return true;
})
.map(async (client) => {
// Add or update this peer on the olm if it is connected
// const allSiteResources = await db // only get the site resources that this client has access to
// .select()
// .from(siteResources)
// .innerJoin(
// clientSiteResourcesAssociationsCache,
// eq(
// siteResources.siteResourceId,
// clientSiteResourcesAssociationsCache.siteResourceId
// )
// )
// .where(
// and(
// eq(siteResources.siteId, site.siteId),
// eq(
// clientSiteResourcesAssociationsCache.clientId,
// client.clients.clientId
// )
// )
// );
if (!client.clientSitesAssociationsCache.isJitMode) {
// if we are adding sites through jit then dont add the site to the olm
// update the peer info on the olm
// if the peer has not been added yet this will be a no-op
await updatePeer(client.clients.clientId, {
siteId: site.siteId,
endpoint: site.endpoint!,
relayEndpoint: formatEndpoint(
exitNode.endpoint,
exitNode.relayPort
),
publicKey: site.publicKey!,
serverIP: site.address,
serverPort: site.listenPort
// remoteSubnets: generateRemoteSubnets(
// allSiteResources.map(
// ({ siteResources }) => siteResources
// )
// ),
// aliases: generateAliasConfig(
// allSiteResources.map(
// ({ siteResources }) => siteResources
// )
// )
});
// also trigger the peer add handshake in case the peer was not already added to the olm and we need to hole punch
// if it has already been added this will be a no-op
await initPeerAddHandshake(
// this will kick off the add peer process for the client
client.clients.clientId,
{
siteId,
exitNode: {
publicKey: exitNode.publicKey,
relayPort: exitNode.relayPort,
endpoint: exitNode.endpoint
}
}
);
}
return {
publicKey: client.clients.pubKey!,
allowedIps: [
`${client.clients.subnet.split("/")[0]}/32`
], // we want to only allow from that client
endpoint: client.clientSitesAssociationsCache.isRelayed
? ""
: client.clientSitesAssociationsCache.endpoint! // if its relayed it should be localhost
};
})
);
}
// Filter out any null values from peers that didn't have an olm
const validPeers = peers.filter((peer) => peer !== null);
// Batch-fetch the olm for every client on this site so the peer updates
// can be sent as one batch instead of looking up the olm per client
const olmRowsQuery = sendPeers
? db
.select({
clientId: olms.clientId,
olmId: olms.olmId,
version: olms.version
})
.from(olms)
.innerJoin(
clientSitesAssociationsCache,
eq(olms.clientId, clientSitesAssociationsCache.clientId)
)
.where(eq(clientSitesAssociationsCache.siteId, siteId))
: Promise.resolve([]);
// Get all enabled site resources for this site by joining through siteNetworks and networks
const allSiteResources = await db
const allSiteResourcesQuery = db
.select()
.from(siteResources)
.innerJoin(networks, eq(siteResources.networkId, networks.networkId))
@@ -165,11 +79,110 @@ export async function buildClientConfigurationForNewtClient(
)
.then((rows) => rows.map((r) => r.siteResources));
const [clientsRes, olmRows, allSiteResources] = await Promise.all([
clientsQuery,
olmRowsQuery,
allSiteResourcesQuery
]);
const peers: Array<{
publicKey: string;
allowedIps: string[];
endpoint?: string;
}> = [];
if (sendPeers) {
const olmByClientId = new Map<
number,
{ olmId: string; version: string | null }
>();
for (const row of olmRows) {
if (row.clientId !== null && !olmByClientId.has(row.clientId)) {
olmByClientId.set(row.clientId, {
olmId: row.olmId,
version: row.version
});
}
}
const olmPeerUpdates: Parameters<typeof updatePeersBatch>[0] = [];
const olmPeerAddHandshakes: Parameters<
typeof initPeerAddHandshakeBatch
>[0] = [];
// Prepare peers data for the response
for (const client of clientsRes) {
if (!client.clients.pubKey) {
logger.debug(
`Client ${client.clients.clientId} has no public key, skipping`
);
continue;
}
if (!client.clients.subnet) {
logger.debug(
`Client ${client.clients.clientId} has no subnet, skipping`
);
continue;
}
// Add or update this peer on the olm if it is connected
const olm = olmByClientId.get(client.clients.clientId);
if (olm && !client.clientSitesAssociationsCache.isJitMode) {
// if we are adding sites through jit then dont add the site to the olm
// update the peer info on the olm
// if the peer has not been added yet this will be a no-op
olmPeerUpdates.push({
clientId: client.clients.clientId,
peer: {
siteId: site.siteId,
endpoint: site.endpoint!,
relayEndpoint: formatEndpoint(
exitNode!.endpoint,
exitNode!.relayPort
),
publicKey: site.publicKey!,
serverIP: site.address,
serverPort: site.listenPort
},
olmId: olm.olmId,
version: olm.version
});
// also trigger the peer add handshake in case the peer was not already added to the olm and we need to hole punch
// if it has already been added this will be a no-op
olmPeerAddHandshakes.push({
clientId: client.clients.clientId,
peer: {
siteId,
exitNode: {
publicKey: exitNode!.publicKey,
relayPort: exitNode!.relayPort,
endpoint: exitNode!.endpoint
}
},
olmId: olm.olmId
});
}
peers.push({
publicKey: client.clients.pubKey,
allowedIps: [`${client.clients.subnet.split("/")[0]}/32`], // we want to only allow from that client
endpoint: client.clientSitesAssociationsCache.isRelayed
? ""
: client.clientSitesAssociationsCache.endpoint! // if its relayed it should be localhost
});
}
// send the updates before the handshakes, matching the per-client order
await updatePeersBatch(olmPeerUpdates);
await initPeerAddHandshakeBatch(olmPeerAddHandshakes);
}
const targetsToSend: SubnetProxyTargetV2[] = [];
if (allSiteResources.length === 0) {
return {
peers: validPeers,
peers,
targets: targetsToSend,
clientIdentities: [] as ClientIdentity[]
};
@@ -256,7 +269,7 @@ export async function buildClientConfigurationForNewtClient(
);
return {
peers: validPeers,
peers,
targets: targetsToSend,
clientIdentities
};
@@ -268,7 +281,7 @@ export async function buildTargetConfigurationForNewtClient(
remoteExitNodeId?: string
) {
// Get enabled HTTP/TCP/UDP targets for resources and AI providers
const allTargets = await db
const allTargetsQuery = db
.select({
resourceId: targets.resourceId,
targetId: targets.targetId,
@@ -292,7 +305,7 @@ export async function buildTargetConfigurationForNewtClient(
)
);
const allHealthChecks = await db
const allHealthChecksQuery = db
.select({
targetHealthCheckId: targetHealthCheck.targetHealthCheckId,
hcEnabled: targetHealthCheck.hcEnabled,
@@ -316,7 +329,7 @@ export async function buildTargetConfigurationForNewtClient(
.where(eq(targetHealthCheck.siteId, siteId));
// Get all enabled targets with their resource mode information
const allBrowserGatewayTargets = await db
const allBrowserGatewayTargetsQuery = db
.select({
resourceId: targets.resourceId,
targetId: targets.targetId,
@@ -337,6 +350,27 @@ export async function buildTargetConfigurationForNewtClient(
)
);
const remoteNodeResourcesQuery = remoteExitNodeId
? db
.select()
.from(remoteExitNodeResources)
.where(
eq(remoteExitNodeResources.remoteExitNodeId, remoteExitNodeId)
)
: Promise.resolve([]);
const [
allTargets,
allHealthChecks,
allBrowserGatewayTargets,
remoteNodeResources
] = await Promise.all([
allTargetsQuery,
allHealthChecksQuery,
allBrowserGatewayTargetsQuery,
remoteNodeResourcesQuery
]);
const { tcpTargets, udpTargets } = allTargets.reduce(
(acc, target) => {
// Filter out invalid targets
@@ -424,18 +458,10 @@ export async function buildTargetConfigurationForNewtClient(
};
});
let remoteExitNodeSubnets: string[] = [];
if (remoteExitNodeId) {
const remoteNodeResources = await db
.select()
.from(remoteExitNodeResources)
.where(
eq(remoteExitNodeResources.remoteExitNodeId, remoteExitNodeId)
);
// filter through these and provide the subnets
remoteExitNodeSubnets = remoteNodeResources.map((r) => r.destination);
}
// filter through these and provide the subnets
const remoteExitNodeSubnets: string[] = remoteNodeResources.map(
(r) => r.destination
);
return {
validHealthCheckTargets,
@@ -93,7 +93,7 @@ export const handleNewtGetConfigMessage: MessageHandler = async (context) => {
.where(eq(exitNodes.exitNodeId, site.exitNodeId))
.limit(1);
if (
exitNode.reachableAt &&
exitNode?.reachableAt &&
existingSite.exitNodeSubnet &&
existingSite.listenPort
) {
+64 -7
View File
@@ -148,6 +148,55 @@ export async function updatePeer(
logger.info(`Updated peer ${peer.publicKey} on olm ${olmId}`);
}
export async function updatePeersBatch(
updates: {
clientId: number;
peer: {
siteId: number;
publicKey: string;
endpoint: string;
relayEndpoint?: string;
serverIP?: string | null;
serverPort?: number | null;
remoteSubnets?: string[] | null;
aliases?: Alias[] | null;
};
olmId: string;
version?: string | null;
}[]
) {
if (updates.length === 0) {
return;
}
await sendToClientsBatch(
updates.map((item) => ({
clientId: item.olmId,
message: {
type: "olm/wg/peer/update",
data: {
siteId: item.peer.siteId,
publicKey: item.peer.publicKey,
endpoint: item.peer.endpoint,
relayEndpoint: item.peer.relayEndpoint,
serverIP: item.peer.serverIP,
serverPort: item.peer.serverPort,
remoteSubnets: item.peer.remoteSubnets,
aliases: item.peer.aliases
}
},
options: {
incrementConfigVersion: true,
compress: canCompress(item.version, "olm")
}
}))
).catch((error) => {
logger.warn(`Error sending batched olm peer updates:`, error);
});
logger.info(`Updated ${updates.length} peer(s) on olms (batch)`);
}
export async function initPeerAddHandshake(
clientId: number,
peer: {
@@ -330,20 +379,28 @@ export async function initPeerAddHandshakeBatch(
logger.warn(`Error sending batched olm handshakes:`, error);
});
// One update per site instead of one per handshake
const clientIdsBySiteId = new Map<number, number[]>();
for (const item of handshakes) {
const clientIds = clientIdsBySiteId.get(item.peer.siteId);
if (clientIds) {
clientIds.push(item.clientId);
} else {
clientIdsBySiteId.set(item.peer.siteId, [item.clientId]);
}
}
await Promise.all(
handshakes.map((item) =>
Array.from(clientIdsBySiteId.entries()).map(([siteId, clientIds]) =>
db
.update(clientSitesAssociationsCache)
.set({ isJitMode: false })
.where(
and(
eq(
eq(clientSitesAssociationsCache.siteId, siteId),
inArray(
clientSitesAssociationsCache.clientId,
item.clientId
),
eq(
clientSitesAssociationsCache.siteId,
item.peer.siteId
clientIds
)
)
)