mirror of
https://github.com/fosrl/pangolin.git
synced 2026-10-23 13:09:05 +02:00
Improve latency on register messages for newt
This commit is contained in:
+4
-11
@@ -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`;
|
||||
}
|
||||
|
||||
@@ -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
|
||||
) {
|
||||
|
||||
@@ -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
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user