Merge pull request #3974 from fosrl/dev

Improve speed of register handlers
This commit is contained in:
Owen Schwartz
2026-10-22 16:27:52 -04:00
committed by GitHub
26 changed files with 541 additions and 260 deletions
+1
View File
@@ -1488,6 +1488,7 @@ export const olms = pgTable(
dateCreated: varchar("dateCreated").notNull(),
version: text("version"),
agent: text("agent"),
agentVersion: text("agentVersion"),
name: varchar("name"),
clientId: integer("clientId").references(() => clients.clientId, {
// we will switch this depending on the current org it wants to connect to
+1
View File
@@ -843,6 +843,7 @@ export const olms = sqliteTable(
dateCreated: text("dateCreated").notNull(),
version: text("version"),
agent: text("agent"),
agentVersion: text("agentVersion"),
name: text("name"),
clientId: integer("clientId").references(() => clients.clientId, {
// we will switch this depending on the current org it wants to connect to
+39 -3
View File
@@ -1,16 +1,20 @@
import semver from "semver";
// NEWT AND OLM
// Pass the newt/olm library version (newts.version / olms.version), not
// agentVersion: the agent may be an app embedding the library (the CLI,
// Pangolin Windows, ...) whose version is unrelated to the protocol support.
export function canCompress(
clientVersion: string | null | undefined,
type: "newt" | "olm"
agent: "newt" | "olm"
): boolean {
try {
if (!clientVersion) return false;
// check if it is a valid semver
if (!semver.valid(clientVersion)) return false;
if (type === "newt") {
if (agent === "newt") {
return semver.gte(clientVersion, "1.10.3");
} else if (type === "olm") {
} else if (agent === "olm") {
return semver.gte(clientVersion, "1.4.3");
}
return false;
@@ -23,6 +27,7 @@ export function canCompress(
// sync message's `certs` array, instead of requiring each target to carry
// its own inline `tlsCert`/`tlsKey` PEM data. Bump the version floor here to
// match whatever release first ships the newt-side support.
// NEWT
export function supportsCertReferences(
clientVersion: string | null | undefined
): boolean {
@@ -38,6 +43,7 @@ export function supportsCertReferences(
// Whether this newt client understands newt/wg/client-identities/add, used to
// set Remote-* headers on requests through its HTTP proxy. Bump the version
// floor here to match whatever release first ships the newt-side support.
// NEWT
export function supportsClientIdentities(
clientVersion: string | null | undefined
): boolean {
@@ -49,3 +55,33 @@ export function supportsClientIdentities(
return false;
}
}
// Whether this newt client understands the v2 subnet proxy target format.
// Older clients need targets converted back to v1 before they are sent.
// NEWT
export function supportsTargetsV2(
clientVersion: string | null | undefined
): boolean {
try {
if (!clientVersion) return false;
if (!semver.valid(clientVersion)) return false;
return semver.gte(clientVersion, "1.10.3");
} catch {
return false;
}
}
// Whether this newt client understands the sync message sent when its config
// version falls behind the server's.
// NEWT
export function supportsSync(
clientVersion: string | null | undefined
): boolean {
try {
if (!clientVersion) return false;
if (!semver.valid(clientVersion)) return false;
return semver.gte(clientVersion, "1.14.0");
} catch {
return false;
}
}
+3 -5
View File
@@ -492,6 +492,9 @@ export async function getNextAvailableOrgSubnet(): Promise<{
}
}
const cidrSchema = z.union([z.cidrv4(), z.cidrv6()]);
const ipSchema = z.union([z.ipv4(), z.ipv6()]);
export function generateRemoteSubnets(
allSiteResources: SiteResource[]
): string[] {
@@ -503,13 +506,11 @@ 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 cidrSchema = z.union([z.cidrv4(), z.cidrv6()]);
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 ipSchema = z.union([z.ipv4(), z.ipv6()]);
const parseResult = ipSchema.safeParse(sr.destination);
return parseResult.success;
}
@@ -598,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`;
@@ -785,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`;
@@ -820,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`;
}
+15
View File
@@ -110,6 +110,21 @@ export class LockManager {
};
}
/**
* Check whether any of the given locks currently exist
* @param lockKeys - Unique identifiers for the locks
* @returns Promise<boolean> - true if at least one lock is held
*/
async anyLockExists(lockKeys: string[]): Promise<boolean> {
for (const lockKey of lockKeys) {
this.clearExpiredLocalLock(lockKey);
if (localLocks.has(lockKey)) {
return true;
}
}
return false;
}
/**
* Extend the TTL of an existing lock, provided the token matches the
* acquisition currently holding it.
+26 -16
View File
@@ -92,6 +92,24 @@ export async function hasActiveSiteResourceRebuild(
return rebuildQueue.isQueued({ type: "site-resource", id: siteResourceId });
}
/**
* Returns true if a rebuild for any of the given site resources is currently
* active or pending in the rebuild queue. Checks all of them in a couple of
* round trips instead of one per resource.
*/
export async function hasAnyActiveSiteResourceRebuild(
siteResourceIds: number[]
): Promise<boolean> {
if (siteResourceIds.length === 0) return false;
const lockKeys = siteResourceIds.map(
(id) => `rebuild-client-associations:site-resource:${id}`
);
if (await lockManager.anyLockExists(lockKeys)) return true;
return rebuildQueue.anyQueued(
siteResourceIds.map((id) => ({ type: "site-resource", id }))
);
}
/**
* Resolves once there is no active or queued rebuild for the given site resource.
* Logs a warning and resolves early if the timeout is reached.
@@ -130,14 +148,10 @@ export async function waitForSiteRebuildIdle(
eq(siteNetworks.networkId, siteResources.networkId)
)
.where(eq(siteNetworks.siteId, siteId));
let allIdle = true;
for (const { siteResourceId } of resourceRows) {
if (await hasActiveSiteResourceRebuild(siteResourceId)) {
allIdle = false;
break;
}
}
if (allIdle) return;
const busy = await hasAnyActiveSiteResourceRebuild(
resourceRows.map((row) => row.siteResourceId)
);
if (!busy) return;
await new Promise<void>((r) =>
setTimeout(r, REBUILD_IDLE_POLL_INTERVAL_MS)
);
@@ -164,14 +178,10 @@ export async function waitForClientRebuildIdle(
})
.from(clientSiteResourcesAssociationsCache)
.where(eq(clientSiteResourcesAssociationsCache.clientId, clientId));
let allIdle = true;
for (const { siteResourceId } of resourceRows) {
if (await hasActiveSiteResourceRebuild(siteResourceId)) {
allIdle = false;
break;
}
}
if (allIdle) return;
const busy = await hasAnyActiveSiteResourceRebuild(
resourceRows.map((row) => row.siteResourceId)
);
if (!busy) return;
await new Promise<void>((r) =>
setTimeout(r, REBUILD_IDLE_POLL_INTERVAL_MS)
);
+5
View File
@@ -20,6 +20,7 @@ export interface RebuildQueueManager {
enqueue(job: RebuildJob): Promise<void>;
startProcessing(handlers: RebuildJobHandlers): void;
isQueued(job: RebuildJob): Promise<boolean>;
anyQueued(jobs: RebuildJob[]): Promise<boolean>;
}
// In-process FIFO used when there is no Redis to back a distributed queue
@@ -48,6 +49,10 @@ class InMemoryRebuildQueue implements RebuildQueueManager {
return this.queuedSet.has(dedupeKey(job));
}
async anyQueued(jobs: RebuildJob[]): Promise<boolean> {
return jobs.some((job) => this.queuedSet.has(dedupeKey(job)));
}
async enqueue(job: RebuildJob): Promise<void> {
const key = dedupeKey(job);
if (this.queuedSet.has(key)) {
+36
View File
@@ -17,6 +17,9 @@ import { v4 as uuidv4 } from "uuid";
const instanceId = uuidv4();
// Max keys passed to a single EXISTS call in anyLockExists
const LOCK_EXISTS_CHUNK_SIZE = 1000;
type LocalLockRecord = {
owner: string;
expiresAt: number;
@@ -254,6 +257,39 @@ export class LockManager {
}
}
/**
* Check whether any of the given locks currently exist, using one EXISTS
* call per chunk of keys instead of a round trip per key
* @param lockKeys - Unique identifiers for the locks
* @returns Promise<boolean> - true if at least one lock is held
*/
async anyLockExists(lockKeys: string[]): Promise<boolean> {
if (!redis || !redis.status || redis.status !== "ready") {
for (const lockKey of lockKeys) {
this.clearExpiredLocalLock(lockKey);
if (localLocks.has(lockKey)) {
return true;
}
}
return false;
}
try {
for (let i = 0; i < lockKeys.length; i += LOCK_EXISTS_CHUNK_SIZE) {
const redisKeys = lockKeys
.slice(i, i + LOCK_EXISTS_CHUNK_SIZE)
.map((lockKey) => `lock:${lockKey}`);
if ((await redis.exists(...redisKeys)) > 0) {
return true;
}
}
return false;
} catch (error) {
logger.error("Failed to check lock existence:", error);
return false;
}
}
/**
* Extend the TTL of an existing lock, provided the token matches the
* acquisition currently holding it.
+18
View File
@@ -66,6 +66,24 @@ class RedisRebuildQueue {
}
}
// Pipelines the membership checks so any number of jobs costs one round trip
async anyQueued(jobs: RebuildJob[]): Promise<boolean> {
if (jobs.length === 0) return false;
if (!redis || redis.status !== "ready") return false;
try {
const pipeline = redis.pipeline();
for (const job of jobs) {
pipeline.sismember(QUEUED_SET_KEY, `${job.type}:${job.id}`);
}
const results = await pipeline.exec();
return (results ?? []).some(
([err, member]) => !err && member === 1
);
} catch {
return false;
}
}
async enqueue(job: RebuildJob): Promise<void> {
if (!redis || redis.status !== "ready") {
logger.warn(
+12 -4
View File
@@ -792,13 +792,21 @@ const sendToClientsBatch = async (
)
);
for (const entry of entries) {
// Issue all increments at once so Redis receives them back to back on one
// connection (in entry order) instead of waiting a round trip per entry
const configVersions = await Promise.all(
entries.map((entry) =>
entry.options?.incrementConfigVersion
? incrementClientConfigVersion(entry.clientId)
: stableConfigVersionByClient.get(entry.clientId)
)
);
for (const [index, entry] of entries.entries()) {
const options = entry.options || {};
const { clientId, message } = entry;
const configVersion = options.incrementConfigVersion
? await incrementClientConfigVersion(clientId)
: stableConfigVersionByClient.get(clientId);
const configVersion = configVersions[index];
logger.debug(
`sendToClientsBatch: Message type ${message.type} queued for clientId ${clientId} (new configVersion: ${configVersion})`
+2
View File
@@ -228,6 +228,7 @@ export type GetClientResponse = NonNullable<
olmId: string | null;
agent: string | null;
olmVersion: string | null;
agentVersion: string | null;
userEmail: string | null;
userName: string | null;
userUsername: string | null;
@@ -421,6 +422,7 @@ export async function getClient(
olmId: client.olms ? client.olms.olmId : null,
agent: client.olms?.agent || null,
olmVersion: client.olms?.version || null,
agentVersion: client.olms?.agentVersion || null,
userEmail: client.user?.email ?? null,
userName: client.user?.name ?? null,
userUsername: client.user?.username ?? null,
+1
View File
@@ -187,6 +187,7 @@ function queryClientsBase() {
userEmail: users.email,
niceId: clients.niceId,
agent: olms.agent,
agentVersion: olms.agentVersion,
approvalState: clients.approvalState,
olmArchived: olms.archived,
archived: clients.archived,
+1
View File
@@ -172,6 +172,7 @@ function queryUserDevicesBase() {
idpVariant: idpOidcConfig.variant,
niceId: clients.niceId,
agent: olms.agent,
agentVersion: olms.agentVersion,
approvalState: clients.approvalState,
olmArchived: olms.archived,
archived: clients.archived,
+6 -12
View File
@@ -6,12 +6,12 @@ import {
SubnetProxyTarget,
SubnetProxyTargetV2
} from "@server/lib/ip";
import { canCompress } from "@server/lib/clientVersionChecks";
import {
canCompress,
supportsTargetsV2
} from "@server/lib/clientVersionChecks";
import logger from "@server/logger";
import { eq, inArray } from "drizzle-orm";
import semver from "semver";
const NEWT_V2_TARGETS_VERSION = ">=1.10.3";
export async function convertTargetsIfNecessary(
newtId: string,
@@ -27,10 +27,7 @@ export async function convertTargetsIfNecessary(
}
// check the semver
if (
newt.version &&
!semver.satisfies(newt.version, NEWT_V2_TARGETS_VERSION)
) {
if (newt.version && !supportsTargetsV2(newt.version)) {
logger.debug(
`addTargets Newt version ${newt.version} does not support targets v2 falling back`
);
@@ -167,10 +164,7 @@ export async function updateTargets(
}
// check the semver
if (
newt.version &&
!semver.satisfies(newt.version, NEWT_V2_TARGETS_VERSION)
) {
if (newt.version && !supportsTargetsV2(newt.version)) {
logger.debug(
`addTargets Newt version ${newt.version} does not support targets v2 falling back`
);
+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
) {
+2 -6
View File
@@ -5,10 +5,9 @@ import { Newt } from "@server/db";
import { eq } from "drizzle-orm";
import logger from "@server/logger";
import { sendNewtSyncMessage } from "./sync";
import semver from "semver";
import { recordSitePing } from "./pingAccumulator";
import { supportsSync } from "@server/lib/clientVersionChecks";
const NEWT_SUPPORTS_SYNC_VERSION = ">=1.14.0";
const PONG = {
message: {
type: "pong",
@@ -50,10 +49,7 @@ export const handleNewtPingMessage: MessageHandler = async (context) => {
// cross-region latency to the database.
recordSitePing(newt.siteId);
if (
newt.version &&
!semver.satisfies(newt.version, NEWT_SUPPORTS_SYNC_VERSION)
) {
if (newt.version && !supportsSync(newt.version)) {
// Newt does not support the sync message so not checking - stop here -
return PONG;
}
+31
View File
@@ -98,6 +98,37 @@ export async function deletePeer(
return site;
}
export async function addPeersBatch(
peers: {
siteId: number;
newtId: string;
peer: {
publicKey: string;
allowedIps: string[];
endpoint: string;
};
}[]
) {
if (peers.length === 0) {
return;
}
await sendToClientsBatch(
peers.map((peer) => ({
clientId: peer.newtId,
message: {
type: "newt/wg/peer/add",
data: peer.peer
},
options: { incrementConfigVersion: true }
}))
).catch((error) => {
logger.warn(`Error sending batched newt peer additions:`, error);
});
logger.info(`Added ${peers.length} peer(s) to newts (batch)`);
}
export async function deletePeersBatch(
peers: {
siteId: number;
+72 -56
View File
@@ -3,8 +3,8 @@ import {
clientSiteResourcesAssociationsCache,
clientSitesAssociationsCache,
db,
exitNodes,
networks,
newts,
SiteResource,
siteNetworks,
siteResources,
@@ -16,9 +16,8 @@ import {
generateRemoteSubnets
} from "@server/lib/ip";
import logger from "@server/logger";
import { and, eq, inArray } from "drizzle-orm";
import { addPeer, deletePeer } from "../newt/peers";
import { SiR } from "react-icons/si";
import { and, eq } from "drizzle-orm";
import { addPeersBatch, deletePeersBatch } from "../newt/peers";
export async function buildSiteConfigurationForOlmClient(
client: Client,
@@ -41,7 +40,7 @@ export async function buildSiteConfigurationForOlmClient(
let exitNodeAliases: string[] = [];
// Get all sites data
const sitesData = await db
const sitesDataQuery = db
.select()
.from(sites)
.innerJoin(
@@ -50,10 +49,26 @@ export async function buildSiteConfigurationForOlmClient(
)
.where(eq(clientSitesAssociationsCache.clientId, client.clientId));
// Batch-fetch the newt for every site so peer updates can be sent as one
// batch instead of looking up the newt per site. Not needed in JIT mode
// because no peers are pushed to the sites there.
const newtRowsQuery = jitMode
? Promise.resolve([])
: db
.select({ siteId: newts.siteId, newtId: newts.newtId })
.from(newts)
.innerJoin(
clientSitesAssociationsCache,
eq(newts.siteId, clientSitesAssociationsCache.siteId)
)
.where(
eq(clientSitesAssociationsCache.clientId, client.clientId)
);
// Batch-fetch every site resource this client has access to across ALL sites
// in a single query, then group by siteId in memory. This avoids issuing one
// query per site (which would be N round-trips for N sites).
const allClientSiteResources = await db
const allClientSiteResourcesQuery = db
.select({
siteResource: siteResources,
siteId: siteNetworks.siteId
@@ -78,6 +93,19 @@ export async function buildSiteConfigurationForOlmClient(
)
);
const [sitesData, newtRows, allClientSiteResources] = await Promise.all([
sitesDataQuery,
newtRowsQuery,
allClientSiteResourcesQuery
]);
const newtIdBySiteId = new Map<number, string>();
for (const row of newtRows) {
if (row.siteId != null && !newtIdBySiteId.has(row.siteId)) {
newtIdBySiteId.set(row.siteId, row.newtId);
}
}
const siteResourcesBySiteId = new Map<number, SiteResource[]>();
const siteResourcesForExitNode = [];
for (const row of allClientSiteResources) {
@@ -107,28 +135,8 @@ export async function buildSiteConfigurationForOlmClient(
};
}
// Batch-fetch exit nodes for all sites in one query (only needed in relay mode).
const exitNodesById = new Map<number, typeof exitNodes.$inferSelect>();
if (!jitMode && relay) {
const exitNodeIds = Array.from(
new Set(
sitesData
.map(({ sites: s }) => s.exitNodeId)
.filter((id): id is number => id != null)
)
);
if (exitNodeIds.length > 0) {
const nodes = await db
.select()
.from(exitNodes)
.where(inArray(exitNodes.exitNodeId, exitNodeIds));
for (const n of nodes) {
exitNodesById.set(n.exitNodeId, n);
}
}
}
const peerOps: Promise<unknown>[] = [];
const peersToDelete: Parameters<typeof deletePeersBatch>[0] = [];
const peersToAdd: Parameters<typeof addPeersBatch>[0] = [];
// Process each site
for (const {
@@ -175,12 +183,23 @@ export async function buildSiteConfigurationForOlmClient(
// continue;
// }
const newtId = newtIdBySiteId.get(site.siteId);
if (!newtId) {
logger.warn(
`Newt not found for site ${site.siteId}, skipping peer updates`
);
}
// If public key changed, delete old peer from this site
if (client.pubKey && client.pubKey != publicKey) {
logger.info(
if (newtId && client.pubKey && client.pubKey != publicKey) {
logger.debug(
`Public key mismatch. Deleting old peer from site ${site.siteId}...`
);
peerOps.push(deletePeer(site.siteId, client.pubKey!));
peersToDelete.push({
siteId: site.siteId,
publicKey: client.pubKey,
newtId
});
}
if (!site.exitNodeSubnet) {
@@ -191,37 +210,30 @@ export async function buildSiteConfigurationForOlmClient(
// Add the peer to the exit node for this site. The endpoint comes from
// the already-joined association row above, so no extra query needed.
if (association.endpoint && publicKey) {
logger.info(
`Adding peer ${publicKey} to site ${site.siteId} with endpoint ${association.endpoint}`
);
peerOps.push(
addPeer(site.siteId, {
publicKey: publicKey,
allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client
endpoint: relay ? "" : association.endpoint
})
);
if (newtId) {
logger.debug(
`Adding peer ${publicKey} to site ${site.siteId} with endpoint ${association.endpoint}`
);
peersToAdd.push({
siteId: site.siteId,
newtId,
peer: {
publicKey: publicKey,
allowedIps: [`${client.subnet.split("/")[0]}/32`], // we want to only allow from that client
endpoint: relay ? "" : association.endpoint
}
});
}
} else {
logger.warn(
`Client ${client.clientId} has no endpoint, skipping peer addition`
);
}
let relayEndpoint: string | undefined = undefined;
if (relay) {
const exitNode = exitNodesById.get(site.exitNodeId);
if (!exitNode) {
logger.warn(`Exit node not found for site ${site.siteId}`);
continue;
}
relayEndpoint = `${exitNode.endpoint}:${exitNode.relayPort}`;
}
// Add site configuration to the array
siteConfigurations.push({
siteId: site.siteId,
name: site.name,
// relayEndpoint: relayEndpoint, // this can be undefined now if not relayed // lets not do this for now because it would conflict with the hole punch testing
endpoint: site.endpoint,
localEndpoints: site.localEndpoints
? JSON.parse(site.localEndpoints)
@@ -234,10 +246,14 @@ export async function buildSiteConfigurationForOlmClient(
});
}
// Run all peer add/delete operations concurrently rather than serially per
// site, so total time is bounded by the slowest call instead of the sum.
if (peerOps.length > 0) {
Promise.allSettled(peerOps).catch((err) => {
// Send the peer removals and additions to the newts as one batch each
// without blocking the response on them. Removals go first so the old key
// is gone before the new one is added.
if (peersToDelete.length > 0 || peersToAdd.length > 0) {
(async () => {
await deletePeersBatch(peersToDelete);
await addPeersBatch(peersToAdd);
})().catch((err) => {
logger.error("Error processing peer operations: ", err);
});
}
@@ -55,6 +55,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
relay,
olmVersion,
olmAgent,
agentVersion,
orgId,
userToken,
pingResults,
@@ -92,9 +93,17 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
await handleFingerprintInsertion(olm, fingerprint, postures);
}
// Olm clients that predate agentVersion report the version of the app
// embedding olm (Pangolin Windows, the CLI, ...) in olmVersion, so it
// doubles as the agent version. Newer clients report the olm library
// version in olmVersion and the app version in agentVersion.
const reportedAgentVersion =
agentVersion === undefined ? olmVersion : agentVersion;
if (
(olmVersion && olm.version !== olmVersion) ||
(olmAgent && olm.agent !== olmAgent) ||
(reportedAgentVersion && olm.agentVersion !== reportedAgentVersion) ||
olm.archived
) {
await db
@@ -102,6 +111,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
.set({
version: olmVersion,
agent: olmAgent,
agentVersion: reportedAgentVersion,
archived: false
})
.where(eq(olms.olmId, olm.olmId));
+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
)
)
)
@@ -89,6 +89,7 @@ export default async function ClientsPage(props: ClientsPageProps) {
orgId: params.orgId,
online: client.online,
olmVersion: client.olmVersion || undefined,
agentVersion: client.agentVersion || undefined,
olmUpdateAvailable: client.olmUpdateAvailable || false,
userId: client.userId,
username: client.username,
@@ -92,6 +92,7 @@ export default async function ClientsPage(props: ClientsPageProps) {
orgId: params.orgId,
online: client.online,
olmVersion: client.olmVersion || undefined,
agentVersion: client.agentVersion || undefined,
olmUpdateAvailable: Boolean(client.olmUpdateAvailable),
userId: client.userId,
username: client.username,
+12 -11
View File
@@ -44,26 +44,29 @@ export default function SiteInfoCard({}: ClientInfoCardProps) {
const data = useQuery(productUpdatesQueries.latestVersion(true));
const latestPlatformVersions = data.data?.data;
// agentVersion was added after olmVersion, and older olm clients report
// the app version in olmVersion, so fall back to it for clients that have
// not re-registered yet.
const agentVersion = client.agentVersion ?? client.olmVersion;
let updateAvailable = false;
if (client.agent && client.olmVersion && latestPlatformVersions) {
if (client.agent && agentVersion && latestPlatformVersions) {
const agent = agentVersionMap[
client.agent
] as keyof LatestVersionResponse;
if (agent in latestPlatformVersions) {
const agentVersion = latestPlatformVersions[agent];
const latestAgentVersion = latestPlatformVersions[agent];
updateAvailable = Boolean(
semver.valid(client.olmVersion) &&
semver.lt(client.olmVersion, agentVersion.latestVersion)
semver.valid(agentVersion) &&
semver.lt(agentVersion, latestAgentVersion.latestVersion)
);
}
}
const cols =
2 +
(userDisplayName ? 1 : 0) +
(client.agent && client.olmVersion ? 1 : 0);
2 + (userDisplayName ? 1 : 0) + (client.agent && agentVersion ? 1 : 0);
return (
<Alert>
@@ -73,15 +76,13 @@ export default function SiteInfoCard({}: ClientInfoCardProps) {
<InfoSectionTitle>{t("name")}</InfoSectionTitle>
<InfoSectionContent>{client.name}</InfoSectionContent>
</InfoSection>
{client.agent && client.olmVersion ? (
{client.agent && agentVersion ? (
<InfoSection>
<InfoSectionTitle>{t("agent")}</InfoSectionTitle>
<InfoSectionContent>
<div className="flex items-center gap-2">
<span>
{client.agent +
" v" +
client.olmVersion}
{client.agent + " v" + agentVersion}
</span>
{updateAvailable && (
<InfoPopup
+15 -7
View File
@@ -56,6 +56,7 @@ export type ClientRow = {
orgId: string;
online: boolean;
olmVersion?: string;
agentVersion?: string;
olmUpdateAvailable: boolean;
userId: string | null;
username: string | null;
@@ -397,10 +398,16 @@ export default function MachineClientsTable({
"Olm CLI": "olm"
};
// agentVersion was added after olmVersion, and older olm
// clients report the app version in olmVersion, so fall
// back to it for rows that have not re-registered yet.
const displayedVersion =
originalRow.agentVersion ?? originalRow.olmVersion;
let updateAvailable = false;
if (
originalRow.olmVersion &&
displayedVersion &&
originalRow.agent &&
latestPlatformVersions
) {
@@ -409,13 +416,14 @@ export default function MachineClientsTable({
] as keyof LatestVersionResponse;
if (agent in latestPlatformVersions) {
const agentVersion = latestPlatformVersions[agent];
const latestAgentVersion =
latestPlatformVersions[agent];
updateAvailable = Boolean(
semver.valid(originalRow.olmVersion) &&
semver.valid(displayedVersion) &&
semver.lt(
originalRow.olmVersion,
agentVersion.latestVersion
displayedVersion,
latestAgentVersion.latestVersion
)
);
}
@@ -423,11 +431,11 @@ export default function MachineClientsTable({
return (
<div className="flex items-center space-x-1">
{originalRow.agent && originalRow.olmVersion ? (
{originalRow.agent && displayedVersion ? (
<Badge variant="secondary">
{originalRow.agent +
" v" +
originalRow.olmVersion}
displayedVersion}
</Badge>
) : (
"-"
+15 -7
View File
@@ -55,6 +55,7 @@ export type ClientRow = {
orgId: string;
online: boolean;
olmVersion?: string;
agentVersion?: string;
olmUpdateAvailable: boolean;
userId: string | null;
username: string | null;
@@ -578,10 +579,16 @@ export default function UserDevicesTable({
"Olm CLI": "olm"
};
// agentVersion was added after olmVersion, and older olm
// clients report the app version in olmVersion, so fall
// back to it for rows that have not re-registered yet.
const displayedVersion =
originalRow.agentVersion ?? originalRow.olmVersion;
let updateAvailable = false;
if (
originalRow.olmVersion &&
displayedVersion &&
originalRow.agent &&
latestPlatformVersions
) {
@@ -590,13 +597,14 @@ export default function UserDevicesTable({
] as keyof LatestVersionResponse;
if (agent in latestPlatformVersions) {
const agentVersion = latestPlatformVersions[agent];
const latestAgentVersion =
latestPlatformVersions[agent];
updateAvailable = Boolean(
semver.valid(originalRow.olmVersion) &&
semver.valid(displayedVersion) &&
semver.lt(
originalRow.olmVersion,
agentVersion.latestVersion
displayedVersion,
latestAgentVersion.latestVersion
)
);
}
@@ -604,11 +612,11 @@ export default function UserDevicesTable({
return (
<div className="flex items-center space-x-1">
{originalRow.agent && originalRow.olmVersion ? (
{originalRow.agent && displayedVersion ? (
<Badge variant="secondary">
{originalRow.agent +
" v" +
originalRow.olmVersion}
displayedVersion}
</Badge>
) : (
"-"