Add inference resource type

server/db/pg/schema/schema.ts

@@ -99,6 +99,7 @@ export const sites = pgTable(
         name: varchar("name").notNull(),
         pubKey: varchar("pubKey"),
         subnet: varchar("subnet"),
+        exitNodeSubnet: text("exitNodeSubnet"), // this is the subnet when connecting to an exit node
         megabytesIn: real("bytesIn").default(0),
         megabytesOut: real("bytesOut").default(0),
         lastBandwidthUpdate: varchar("lastBandwidthUpdate"),
@@ -194,7 +195,10 @@ export const resources = pgTable(
         postAuthPath: text("postAuthPath"),
         health: varchar("health").default("unknown"), // "healthy", "unhealthy", "unknown"
         wildcard: boolean("wildcard").notNull().default(false),
-        mode: text("mode").default("http").notNull(), // rdp, ssh, http, vnc
+        mode: text("mode")
+            .default("http")
+            .$type<"rdp" | "ssh" | "http" | "vnc" | "inference">()
+            .notNull(),
         pamMode: varchar("pamMode", { length: 32 })
             .$type<"passthrough" | "push">()
             .default("passthrough"),
@@ -428,7 +432,7 @@ export const siteResources = pgTable(
         name: varchar("name").notNull(),
         ssl: boolean("ssl").notNull().default(false),
         mode: varchar("mode")
-            .$type<"host" | "cidr" | "http" | "ssh">()
+            .$type<"host" | "cidr" | "http" | "ssh" | "inference">()
             .notNull(), // "host" | "cidr" | "http"
         scheme: varchar("scheme").$type<"http" | "https">(), // only for when we are doing https or http mode
         proxyPort: integer("proxyPort"), // only for port mode
@@ -1181,7 +1185,7 @@ export const clients = pgTable(
         olmId: text("olmId"), // to lock it to a specific olm optionally
         name: varchar("name").notNull(),
         pubKey: varchar("pubKey"),
-        subnet: varchar("subnet").notNull(),
+        exitNodeSubnet: varchar("exitNodeSubnet").notNull(),
         megabytesIn: real("bytesIn"),
         megabytesOut: real("bytesOut"),
         lastBandwidthUpdate: varchar("lastBandwidthUpdate"),

server/db/sqlite/schema/schema.ts

@@ -107,7 +107,7 @@ export const sites = sqliteTable("sites", {
     }),
     name: text("name").notNull(),
     pubKey: text("pubKey"),
-    subnet: text("subnet"),
+    exitNodeSubnet: text("exitNodeSubnet"),
     megabytesIn: integer("bytesIn").default(0),
     megabytesOut: integer("bytesOut").default(0),
     lastBandwidthUpdate: text("lastBandwidthUpdate"),
@@ -203,7 +203,10 @@ export const resources = sqliteTable("resources", {
     postAuthPath: text("postAuthPath"),
     health: text("health").default("unknown"), // "healthy", "unhealthy", "unknown"
     wildcard: integer("wildcard", { mode: "boolean" }).notNull().default(false),
-    mode: text("mode").default("http").notNull(), // rdp, ssh, http, vnc
+    mode: text("mode")
+        .default("http")
+        .$type<"rdp" | "ssh" | "http" | "vnc" | "inference">()
+        .notNull(), // rdp, ssh, http, vnc, inference
     pamMode: text("pamMode")
         .$type<"passthrough" | "push">()
         .default("passthrough"),
@@ -425,7 +428,9 @@ export const siteResources = sqliteTable("siteResources", {
     niceId: text("niceId").notNull(),
     name: text("name").notNull(),
     ssl: integer("ssl", { mode: "boolean" }).notNull().default(false),
-    mode: text("mode").$type<"host" | "cidr" | "http" | "ssh">().notNull(), // "host" | "cidr" | "http"
+    mode: text("mode")
+        .$type<"host" | "cidr" | "http" | "ssh" | "inference">()
+        .notNull(), // "host" | "cidr" | "http"
     scheme: text("scheme").$type<"http" | "https">(), // only for when we are doing https or http mode
     proxyPort: integer("proxyPort"), // only for port mode
     destinationPort: integer("destinationPort"), // only for port mode
@@ -599,6 +604,7 @@ export const clients = sqliteTable("clients", {
     pubKey: text("pubKey"),
     olmId: text("olmId"), // to lock it to a specific olm optionally
     subnet: text("subnet").notNull(),
+    exitNodeSubnet: text("exitNodeSubnet"), // this is the subnet when connecting to an exit node
     megabytesIn: integer("bytesIn"),
     megabytesOut: integer("bytesOut"),
     lastBandwidthUpdate: text("lastBandwidthUpdate"),

server/lib/calculateUserClientsForOrgs.ts

@@ -339,19 +339,6 @@ export async function calculateUserClientsForOrgs(
                 continue;
             }
 
-            // Get exit nodes for this org
-            const exitNodesList = await getExitNodes(orgId);
-
-            if (exitNodesList.length === 0) {
-                logger.warn(
-                    `Skipping org ${orgId} for OLM ${olm.olmId} (user ${userId}): no exit nodes found`
-                );
-                continue;
-            }
-
-            const randomExitNode =
-                exitNodesList[Math.floor(Math.random() * exitNodesList.length)];
-
             // Get next available subnet
             const { value: newSubnet, release: releaseSubnetLock } =
                 await getNextAvailableClientSubnet(orgId, trx);
@@ -370,7 +357,6 @@ export async function calculateUserClientsForOrgs(
             const newClientData: InferInsertModel<typeof clients> = {
                 userId,
                 orgId: userOrg.orgId,
-                exitNodeId: randomExitNode.exitNodeId,
                 name: olm.name || "User Client",
                 subnet: updatedSubnet,
                 olmId: olm.olmId,

server/lib/exitNodes/calculateExitNodeWeight.ts

@@ -0,0 +1,40 @@
+import { db, sites, clients } from "@server/db";
+import { and, eq, count } from "drizzle-orm";
+
+// (MAX_CONNECTIONS - current_connections) / MAX_CONNECTIONS)
+// higher = more desirable
+// like saying, this node has x% of its capacity left
+export async function calculateExitNodeWeight(
+    exitNodeId: number,
+    maxConnections: number | null | undefined
+): Promise<number | null> {
+    if (maxConnections === null || maxConnections === undefined) {
+        return 1;
+    }
+
+    const [[siteConnections], [clientConnections]] = await Promise.all([
+        db
+            .select({ count: count() })
+            .from(sites)
+            .where(
+                and(eq(sites.exitNodeId, exitNodeId), eq(sites.online, true))
+            ),
+        db
+            .select({ count: count() })
+            .from(clients)
+            .where(
+                and(
+                    eq(clients.exitNodeId, exitNodeId),
+                    eq(clients.online, true)
+                )
+            )
+    ]);
+
+    const currentConnections = siteConnections.count + clientConnections.count;
+
+    if (currentConnections >= maxConnections) {
+        return null;
+    }
+
+    return (maxConnections - currentConnections) / maxConnections;
+}

server/lib/exitNodes/exitNodes.ts

@@ -1,6 +1,5 @@
 import { db, exitNodes, Transaction } from "@server/db";
 import logger from "@server/logger";
-import { ExitNodePingResult } from "@server/routers/newt";
 import { eq } from "drizzle-orm";
 
 export async function verifyExitNodeOrgAccess(
@@ -52,6 +51,16 @@ export async function listExitNodes(
     return allExitNodes;
 }
 
+export type ExitNodePingResult = {
+    exitNodeId: number;
+    latencyMs: number;
+    weight: number;
+    error?: string;
+    exitNodeName: string;
+    endpoint: string;
+    wasPreviouslyConnected: boolean;
+};
+
 export function selectBestExitNode(
     pingResults: ExitNodePingResult[]
 ): ExitNodePingResult | null {

server/lib/exitNodes/getUniqueSubnetForExitNode.ts

@@ -0,0 +1,41 @@
+import { db, ExitNode, Transaction, sites, clients } from "@server/db";
+import { eq } from "drizzle-orm";
+import config from "@server/lib/config";
+import { findNextAvailableCidr } from "@server/lib/ip";
+import { lockManager } from "#dynamic/lib/lock";
+
+export async function getUniqueSubnetForExitNode(
+    exitNode: ExitNode,
+    trx: Transaction | typeof db = db
+): Promise<string | null> {
+    const lockKey = `subnet-allocation:${exitNode.exitNodeId}`;
+
+    return await lockManager.withLock(
+        lockKey,
+        async () => {
+            const [sitesQuery, clientsQuery] = await Promise.all([
+                trx
+                    .select({ subnet: sites.exitNodeSubnet })
+                    .from(sites)
+                    .where(eq(sites.exitNodeId, exitNode.exitNodeId)),
+                trx
+                    .select({ subnet: clients.exitNodeSubnet })
+                    .from(clients)
+                    .where(eq(clients.exitNodeId, exitNode.exitNodeId))
+            ]);
+
+            const blockSize = config.getRawConfig().gerbil.site_block_size;
+            const subnets = [...sitesQuery, ...clientsQuery]
+                .map((row) => row.subnet)
+                .filter(
+                    (subnet): subnet is string =>
+                        !!subnet &&
+                        /^(\d{1,3}\.){3}\d{1,3}\/\d{1,2}$/.test(subnet)
+                );
+            subnets.push(exitNode.address.replace(/\/\d+$/, `/${blockSize}`));
+
+            return findNextAvailableCidr(subnets, blockSize, exitNode.address);
+        },
+        5000 // 5 second lock TTL - subnet allocation should be quick
+    );
+}

server/lib/exitNodes/index.ts

@@ -2,3 +2,5 @@ export * from "./exitNodes";
 export * from "./exitNodeComms";
 export * from "./subnet";
 export * from "./getCurrentExitNodeId";
+export * from "./calculateExitNodeWeight";
+export * from "./getUniqueSubnetForExitNode";

server/lib/rebuildClientAssociations.ts

@@ -966,7 +966,7 @@ export async function updateClientSiteDestinations(
         .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
 
     for (const site of sitesData) {
-        if (!site.sites.subnet) {
+        if (!site.sites.exitNodeSubnet) {
             logger.debug(`Site ${site.sites.siteId} has no subnet, skipping`);
             continue;
         }
@@ -1002,7 +1002,7 @@ export async function updateClientSiteDestinations(
                 sourcePort: parsedEndpoint.port,
                 destinations: [
                     {
-                        destinationIP: site.sites.subnet.split("/")[0],
+                        destinationIP: site.sites.exitNodeSubnet.split("/")[0],
                         destinationPort: site.sites.listenPort || 1 // this satisfies gerbil for now but should be reevaluated
                     }
                 ]
@@ -1010,7 +1010,7 @@ export async function updateClientSiteDestinations(
         } else {
             // add to the existing destinations
             destinations.destinations.push({
-                destinationIP: site.sites.subnet.split("/")[0],
+                destinationIP: site.sites.exitNodeSubnet.split("/")[0],
                 destinationPort: site.sites.listenPort || 1 // this satisfies gerbil for now but should be reevaluated
             });
         }

server/lib/traefik/getTraefikConfig.ts

@@ -87,7 +87,7 @@ export async function getTraefikConfig(
             siteId: sites.siteId,
             siteType: sites.type,
             siteOnline: sites.online,
-            subnet: sites.subnet,
+            subnet: sites.exitNodeSubnet,
             exitNodeId: sites.exitNodeId,
             // Domain cert resolver fields
             domainCertResolver: domains.certResolver,

server/private/lib/exitNodes/exitNodes.ts

@@ -25,7 +25,6 @@ import {
     Transaction
 } from "@server/db";
 import logger from "@server/logger";
-import { ExitNodePingResult } from "@server/routers/newt";
 import { eq, and, or, ne, isNull, inArray } from "drizzle-orm";
 import axios from "axios";
 import config from "../config";
@@ -330,6 +329,16 @@ export async function listExitNodes(
     return exitNodesList;
 }
 
+export type ExitNodePingResult = {
+    exitNodeId: number;
+    latencyMs: number;
+    weight: number;
+    error?: string;
+    exitNodeName: string;
+    endpoint: string;
+    wasPreviouslyConnected: boolean;
+};
+
 /**
  * Selects the most suitable exit node from a list of ping results.
  *

server/private/lib/traefik/getTraefikConfig.ts

@@ -134,7 +134,7 @@ export async function getTraefikConfig(
             siteId: sites.siteId,
             siteType: sites.type,
             siteOnline: sites.online,
-            subnet: sites.subnet,
+            subnet: sites.exitNodeSubnet,
             exitNodeId: sites.exitNodeId,
             // Namespace
             domainNamespaceId: domainNamespaces.domainNamespaceId,

server/private/routers/re-key/reGenerateSiteSecret.ts

@@ -178,7 +178,7 @@ export async function reGenerateSiteSecret(
                     );
                 }
 
-                if (site.exitNodeId && site.subnet) {
+                if (site.exitNodeId && site.exitNodeSubnet) {
                     await deletePeer(site.exitNodeId, site.pubKey!); // the old pubkey
                     await addPeer(site.exitNodeId, {
                         publicKey: pubKey,

server/routers/client/createClient.ts

@@ -255,20 +255,6 @@ export async function createClient(
 
         let newClient: Client | null = null;
         await db.transaction(async (trx) => {
-            // TODO: more intelligent way to pick the exit node
-            const exitNodesList = await listExitNodes(orgId);
-            const randomExitNode =
-                exitNodesList[Math.floor(Math.random() * exitNodesList.length)];
-
-            if (!randomExitNode) {
-                return next(
-                    createHttpError(
-                        HttpCode.NOT_FOUND,
-                        `No exit nodes available. ${build == "saas" ? "Please contact support." : "You need to install gerbil to use the clients."}`
-                    )
-                );
-            }
-
             const [adminRole] = await trx
                 .select()
                 .from(roles)
@@ -287,7 +273,6 @@ export async function createClient(
                 .insert(clients)
                 .values({
                     niceId,
-                    exitNodeId: randomExitNode.exitNodeId,
                     orgId,
                     name,
                     subnet: updatedSubnet,

server/routers/client/createUserClient.ts

@@ -222,11 +222,6 @@ export async function createUserClient(
 
         let newClient: Client | null = null;
         await db.transaction(async (trx) => {
-            // TODO: more intelligent way to pick the exit node
-            const exitNodesList = await listExitNodes(orgId);
-            const randomExitNode =
-                exitNodesList[Math.floor(Math.random() * exitNodesList.length)];
-
             const [adminRole] = await trx
                 .select()
                 .from(roles)
@@ -244,7 +239,6 @@ export async function createUserClient(
             [newClient] = await trx
                 .insert(clients)
                 .values({
-                    exitNodeId: randomExitNode.exitNodeId,
                     orgId,
                     niceId,
                     name,

server/routers/gerbil/getAllRelays.ts

@@ -100,7 +100,7 @@ export async function generateRelayMappings(exitNode: ExitNode) {
     // Filter to sites with the required fields up front so the rest of the
     // function can safely treat endpoint/subnet/listenPort as defined.
     const validSites = sitesRes.filter(
-        (s) => s.endpoint && s.subnet && s.listenPort
+        (s) => s.endpoint && s.exitNodeSubnet && s.listenPort
     );
 
     if (validSites.length === 0) {
@@ -136,7 +136,7 @@ export async function generateRelayMappings(exitNode: ExitNode) {
         if (
             peer.orgId == null ||
             !peer.endpoint ||
-            !peer.subnet ||
+            !peer.exitNodeSubnet ||
             !peer.listenPort
         ) {
             continue;
@@ -183,7 +183,7 @@ export async function generateRelayMappings(exitNode: ExitNode) {
     // Process each site using the pre-fetched data.
     for (const site of validSites) {
         const siteDestination: PeerDestination = {
-            destinationIP: site.subnet!.split("/")[0],
+            destinationIP: site.exitNodeSubnet!.split("/")[0],
             destinationPort: site.listenPort! || 1 // this satisfies gerbil for now but should be reevaluated
         };
 
@@ -207,7 +207,7 @@ export async function generateRelayMappings(exitNode: ExitNode) {
                         continue;
                     }
                     addDestination(site.endpoint!, {
-                        destinationIP: peer.subnet!.split("/")[0],
+                        destinationIP: peer.exitNodeSubnet!.split("/")[0],
                         destinationPort: peer.listenPort! || 1 // this satisfies gerbil for now but should be reevaluated
                     });
                 }

server/routers/gerbil/getConfig.ts

@@ -89,7 +89,7 @@ export async function generateGerbilConfig(exitNode: ExitNode) {
             and(
                 eq(sites.exitNodeId, exitNode.exitNodeId),
                 isNotNull(sites.pubKey),
-                isNotNull(sites.subnet)
+                isNotNull(sites.exitNodeSubnet)
             )
         );
 
@@ -103,7 +103,7 @@ export async function generateGerbilConfig(exitNode: ExitNode) {
             } else if (site.type === "newt") {
                 return {
                     publicKey: site.pubKey,
-                    allowedIps: [site.subnet!]
+                    allowedIps: [site.exitNodeSubnet!]
                 };
             }
             return {

server/routers/gerbil/updateHolePunch.ts

@@ -186,7 +186,7 @@ export async function updateAndGenerateEndpointDestinations(
             .select({
                 siteId: sites.siteId,
                 newtId: newts.newtId,
-                subnet: sites.subnet,
+                subnet: sites.exitNodeSubnet,
                 listenPort: sites.listenPort,
                 publicKey: sites.publicKey,
                 endpoint: clientSitesAssociationsCache.endpoint,

server/routers/newt/handleNewtExitNodesRequestMessage.ts

@@ -2,14 +2,17 @@ import { db, sites } from "@server/db";
 import { MessageHandler } from "@server/routers/ws";
 import { exitNodes, Newt } from "@server/db";
 import logger from "@server/logger";
-import { ne, eq, or, and, count } from "drizzle-orm";
+import { eq } from "drizzle-orm";
 import { listExitNodes } from "#dynamic/lib/exitNodes";
+import { calculateExitNodeWeight } from "@server/lib/exitNodes";
 
-export const handleNewtPingRequestMessage: MessageHandler = async (context) => {
+export const handleNewtExitNodesRequestMessage: MessageHandler = async (
+    context
+) => {
     const { message, client, sendToClient } = context;
     const newt = client as Newt;
 
-    logger.info("Handling ping request newt message!");
+    logger.info("Handling exit nodes request newt message!");
 
     if (!newt) {
         logger.warn("Newt not found");
@@ -54,32 +57,13 @@ export const handleNewtPingRequestMessage: MessageHandler = async (context) => {
 
     const exitNodesPayload = await Promise.all(
         exitNodesList.map(async (node) => {
-            // (MAX_CONNECTIONS - current_connections) / MAX_CONNECTIONS)
-            // higher = more desirable
-            // like saying, this node has x% of its capacity left
+            const weight = await calculateExitNodeWeight(
+                node.exitNodeId,
+                node.maxConnections
+            );
 
-            let weight = 1;
-            const maxConnections = node.maxConnections;
-            if (maxConnections !== null && maxConnections !== undefined) {
-                const [currentConnections] = await db
-                    .select({
-                        count: count()
-                    })
-                    .from(sites)
-                    .where(
-                        and(
-                            eq(sites.exitNodeId, node.exitNodeId),
-                            eq(sites.online, true)
-                        )
-                    );
-
-                if (currentConnections.count >= maxConnections) {
-                    return null;
-                }
-
-                weight =
-                    (maxConnections - currentConnections.count) /
-                    maxConnections;
+            if (weight === null) {
+                return null;
             }
 
             return {

server/routers/newt/handleNewtGetConfigMessage.ts

@@ -95,16 +95,16 @@ export const handleNewtGetConfigMessage: MessageHandler = async (context) => {
             .limit(1);
         if (
             exitNode.reachableAt &&
-            existingSite.subnet &&
+            existingSite.exitNodeSubnet &&
             existingSite.listenPort
         ) {
             const payload = {
                 oldDestination: {
-                    destinationIP: existingSite.subnet?.split("/")[0],
+                    destinationIP: existingSite.exitNodeSubnet?.split("/")[0],
                     destinationPort: existingSite.listenPort || 1 // this satisfies gerbil for now but should be reevaluated
                 },
                 newDestination: {
-                    destinationIP: site.subnet?.split("/")[0],
+                    destinationIP: site.exitNodeSubnet?.split("/")[0],
                     destinationPort: site.listenPort || 1 // this satisfies gerbil for now but should be reevaluated
                 }
             };
@@ -132,7 +132,10 @@ export const handleNewtGetConfigMessage: MessageHandler = async (context) => {
         ({ targets: dedupedTargets, certs } = dedupeCertsForTargets(targets));
     }
 
-    const targetsToSend = await convertTargetsIfNecessary(newt.newtId, dedupedTargets); // for backward compatibility with old newt versions that don't support the new target format
+    const targetsToSend = await convertTargetsIfNecessary(
+        newt.newtId,
+        dedupedTargets
+    ); // for backward compatibility with old newt versions that don't support the new target format
 
     return {
         message: {

server/routers/newt/handleNewtRegisterMessage.ts

@@ -1,30 +1,20 @@
-import { db, ExitNode, newts, remoteExitNodes, Transaction } from "@server/db";
+import { db, newts, remoteExitNodes } from "@server/db";
 import { MessageHandler } from "@server/routers/ws";
 import { exitNodes, Newt, sites } from "@server/db";
 import { eq } from "drizzle-orm";
 import { addPeer, deletePeer } from "../gerbil/peers";
 import logger from "@server/logger";
 import config from "@server/lib/config";
-import { findNextAvailableCidr } from "@server/lib/ip";
 import {
+    ExitNodePingResult,
     selectBestExitNode,
     verifyExitNodeOrgAccess
 } from "#dynamic/lib/exitNodes";
+import { getUniqueSubnetForExitNode } from "@server/lib/exitNodes";
 import { fetchContainers } from "./dockerSocket";
-import { lockManager } from "#dynamic/lib/lock";
 import { buildTargetConfigurationForNewtClient } from "./buildConfiguration";
 import { canCompress } from "@server/lib/clientVersionChecks";
 
-export type ExitNodePingResult = {
-    exitNodeId: number;
-    latencyMs: number;
-    weight: number;
-    error?: string;
-    exitNodeName: string;
-    endpoint: string;
-    wasPreviouslyConnected: boolean;
-};
-
 export const handleNewtRegisterMessage: MessageHandler = async (context) => {
     const { message, client, sendToClient } = context;
     const newt = client as Newt;
@@ -94,9 +84,12 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
         fetchContainers(newt.newtId);
     }
 
-    let siteSubnet = oldSite.subnet;
+    let siteSubnet = oldSite.exitNodeSubnet;
     let exitNodeIdToQuery = oldSite.exitNodeId;
-    if (exitNodeId && (oldSite.exitNodeId !== exitNodeId || !oldSite.subnet)) {
+    if (
+        exitNodeId &&
+        (oldSite.exitNodeId !== exitNodeId || !oldSite.exitNodeSubnet)
+    ) {
         // This effectively moves the exit node to the new one
         exitNodeIdToQuery = exitNodeId; // Use the provided exitNodeId if it differs from the site's exitNodeId
 
@@ -115,7 +108,7 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
             return;
         }
 
-        const newSubnet = await getUniqueSubnetForSite(exitNode);
+        const newSubnet = await getUniqueSubnetForExitNode(exitNode);
 
         if (!newSubnet) {
             logger.error(
@@ -131,7 +124,7 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
             .set({
                 pubKey: publicKey,
                 exitNodeId: exitNodeId,
-                subnet: newSubnet
+                exitNodeSubnet: newSubnet
             })
             .where(eq(sites.siteId, siteId))
             .returning();
@@ -250,40 +243,3 @@ export const handleNewtRegisterMessage: MessageHandler = async (context) => {
         excludeSender: false // Include sender in broadcast
     };
 };
-
-async function getUniqueSubnetForSite(
-    exitNode: ExitNode,
-    trx: Transaction | typeof db = db
-): Promise<string | null> {
-    const lockKey = `subnet-allocation:${exitNode.exitNodeId}`;
-
-    return await lockManager.withLock(
-        lockKey,
-        async () => {
-            const sitesQuery = await trx
-                .select({
-                    subnet: sites.subnet
-                })
-                .from(sites)
-                .where(eq(sites.exitNodeId, exitNode.exitNodeId));
-
-            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(
-                subnets,
-                blockSize,
-                exitNode.address
-            );
-            return newSubnet;
-        },
-        5000 // 5 second lock TTL - subnet allocation should be quick
-    );
-}

server/routers/newt/index.ts

@@ -5,7 +5,7 @@ export * from "./handleNewtRegisterMessage";
 export * from "./handleReceiveBandwidthMessage";
 export * from "./handleNewtGetConfigMessage";
 export * from "./handleSocketMessages";
-export * from "./handleNewtPingRequestMessage";
+export * from "./handleNewtExitNodesRequestMessage";
 export * from "./handleApplyBlueprintMessage";
 export * from "./handleNewtPingMessage";
 export * from "./handleNewtDisconnectingMessage";

server/routers/olm/buildConfiguration.ts

@@ -167,7 +167,7 @@ export async function buildSiteConfigurationForOlmClient(
             peerOps.push(deletePeer(site.siteId, client.pubKey!));
         }
 
-        if (!site.subnet) {
+        if (!site.exitNodeSubnet) {
             logger.debug(`Site ${site.siteId} has no subnet, skipping`);
             continue;
         }

server/routers/olm/handleOlmExitNodesRequestMessage.ts

@@ -0,0 +1,93 @@
+import { db, clients } from "@server/db";
+import { MessageHandler } from "@server/routers/ws";
+import { exitNodes, Olm } from "@server/db";
+import logger from "@server/logger";
+import { eq } from "drizzle-orm";
+import { listExitNodes } from "#dynamic/lib/exitNodes";
+import { calculateExitNodeWeight } from "@server/lib/exitNodes";
+
+export const handleOlmExitNodesRequestMessage: MessageHandler = async (
+    context
+) => {
+    const { message, client: olmClient, sendToClient } = context;
+    const olm = olmClient as Olm;
+
+    logger.info("Handling exit nodes request olm message!");
+
+    if (!olm) {
+        logger.warn("olm not found");
+        return;
+    }
+
+    // Get the olm's orgId through the client relationship
+    if (!olm.clientId) {
+        logger.warn("olm clientId not found");
+        return;
+    }
+
+    const [client] = await db
+        .select({ orgId: clients.orgId })
+        .from(clients)
+        .where(eq(clients.clientId, olm.clientId))
+        .limit(1);
+
+    if (!client || !client.orgId) {
+        logger.warn("client not found");
+        return;
+    }
+
+    const { noCloud, chainId } = message.data;
+
+    const exitNodesList = await listExitNodes(
+        client.orgId,
+        true,
+        noCloud || false,
+        olm.clientId
+    ); // filter for only the online ones
+
+    let lastExitNodeId = null;
+    if (olm.clientId) {
+        const [lastExitNode] = await db
+            .select()
+            .from(clients)
+            .where(eq(clients.clientId, olm.clientId))
+            .limit(1);
+        lastExitNodeId = lastExitNode?.exitNodeId || null;
+    }
+
+    const exitNodesPayload = await Promise.all(
+        exitNodesList.map(async (node) => {
+            const weight = await calculateExitNodeWeight(
+                node.exitNodeId,
+                node.maxConnections
+            );
+
+            if (weight === null) {
+                return null;
+            }
+
+            return {
+                exitNodeId: node.exitNodeId,
+                exitNodeName: node.name,
+                endpoint: node.endpoint,
+                weight,
+                wasPreviouslyConnected: node.exitNodeId === lastExitNodeId
+            };
+        })
+    );
+
+    // filter out null values
+    const filteredExitNodes = exitNodesPayload.filter((node) => node !== null);
+
+    return {
+        message: {
+            type: "olm/ping/exitNodes",
+            data: {
+                exitNodes: filteredExitNodes,
+                chainId: chainId
+            }
+        },
+        broadcast: false, // Send to all clients
+        excludeSender: false // Include sender in broadcast
+    };
+};

server/routers/olm/handleOlmRegisterMessage.ts

@@ -1,4 +1,4 @@
-import { db, orgs, primaryDb } from "@server/db";
+import { db, ExitNode, orgs, primaryDb } from "@server/db";
 import { MessageHandler } from "@server/routers/ws";
 import {
     clients,
@@ -22,6 +22,12 @@ import { canCompress } from "@server/lib/clientVersionChecks";
 import config from "@server/lib/config";
 import cache from "#dynamic/lib/cache"; // not using regional here because we need this in the register message handler before we know where the client is
 import { waitForClientRebuildIdle } from "@server/lib/rebuildClientAssociations";
+import {
+    ExitNodePingResult,
+    selectBestExitNode,
+    verifyExitNodeOrgAccess
+} from "#dynamic/lib/exitNodes";
+import { getUniqueSubnetForExitNode } from "@server/lib/exitNodes";
 
 const HOLEPUNCH_STALE_CHAIN_THRESHOLD = 18;
 const HOLEPUNCH_STALE_CHAIN_TTL_SECONDS = 1800;
@@ -49,11 +55,20 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         olmAgent,
         orgId,
         userToken,
+        pingResults,
         fingerprint,
         postures,
+        backwardsCompatible,
         chainId
     } = message.data;
 
+    if (backwardsCompatible) {
+        logger.debug(
+            "[handleOlmRegisterMessage] Backwards compatible mode detected - not sending connect message and waiting for ping response."
+        );
+        return;
+    }
+
     if (!olm.clientId) {
         logger.warn("[handleOlmRegisterMessage] Olm client ID not found");
         sendOlmError(OlmErrorCodes.CLIENT_ID_NOT_FOUND, olm.olmId);
@@ -284,7 +299,63 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         return;
     }
 
-    if (client.pubKey !== publicKey || client.archived) {
+    let exitNodeId: number | undefined;
+    if (pingResults) {
+        const bestPingResult = selectBestExitNode(
+            pingResults as ExitNodePingResult[]
+        );
+        if (!bestPingResult) {
+            logger.warn("No suitable exit node found based on ping results");
+        }
+        exitNodeId = bestPingResult?.exitNodeId;
+    }
+
+    let clientSubnet = client.exitNodeSubnet;
+    let exitNode: ExitNode | null = null;
+    if (
+        exitNodeId &&
+        (client.exitNodeId !== exitNodeId || !client.exitNodeSubnet)
+    ) {
+        const { exitNode: exitNodeResult, hasAccess } =
+            await verifyExitNodeOrgAccess(exitNodeId, client.orgId);
+
+        exitNode = exitNodeResult;
+
+        if (!exitNode) {
+            logger.warn("[handleOlmRegisterMessage] Exit node not found", {
+                orgId: client.orgId,
+                clientId: client.clientId
+            });
+            return;
+        }
+
+        if (!hasAccess) {
+            logger.warn(
+                "[handleOlmRegisterMessage] Not authorized to use this exit node",
+                { orgId: client.orgId, clientId: client.clientId }
+            );
+            return;
+        }
+
+        const newSubnet = await getUniqueSubnetForExitNode(exitNode);
+
+        if (!newSubnet) {
+            logger.error(
+                `[handleOlmRegisterMessage] No available subnets found for exit node id ${exitNodeId} and client id ${client.clientId}`,
+                { orgId: client.orgId, clientId: client.clientId }
+            );
+            return;
+        }
+
+        clientSubnet = newSubnet;
+    }
+
+    if (
+        client.pubKey !== publicKey ||
+        client.archived ||
+        client.exitNodeId !== exitNodeId ||
+        client.exitNodeSubnet !== clientSubnet
+    ) {
         logger.info(
             "[handleOlmRegisterMessage] Public key mismatch. Updating public key and clearing session info...",
             { orgId: client.orgId, clientId: client.clientId }
@@ -294,7 +365,9 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             .update(clients)
             .set({
                 pubKey: publicKey,
-                archived: false
+                archived: false,
+                exitNodeId: exitNodeId, // this can be undefined if no exit node was selected, which is fine just means we cant talk to the node or connect to it
+                exitNodeSubnet: clientSubnet
             })
             .where(eq(clients.clientId, client.clientId));
 
@@ -394,6 +467,18 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
                 sites: siteConfigurations,
                 tunnelIP: client.subnet,
                 utilitySubnet: org.utilitySubnet,
+                exitNode:
+                    exitNode && client.exitNodeSubnet
+                        ? {
+                              endpoint: `${exitNode.endpoint}:${exitNode.listenPort}`,
+                              relayPort:
+                                  config.getRawConfig().gerbil
+                                      .clients_start_port,
+                              publicKey: exitNode.publicKey,
+                              serverIP: exitNode.address.split("/")[0],
+                              tunnelIP: client.exitNodeSubnet.split("/")[0]
+                          }
+                        : undefined,
                 chainId: chainId
             }
         },

server/routers/olm/index.ts

@@ -15,3 +15,4 @@ export * from "./handleOlmServerInitAddPeerHandshake";
 export * from "./offlineChecker";
 export * from "./handleOlmUnLocalMessage";
 export * from "./handleOlmLocalMessage";
+export * from "./handleOlmExitNodesRequestMessage";

server/routers/site/createSite.ts

@@ -311,13 +311,13 @@ export async function createSite(
                 // lets also make sure there is no overlap with other sites on the exit node
                 const sitesQuery = await db
                     .select({
-                        subnet: sites.subnet
+                        subnet: sites.exitNodeSubnet
                     })
                     .from(sites)
                     .where(
                         and(
                             eq(sites.exitNodeId, exitNodeId),
-                            eq(sites.subnet, subnet)
+                            eq(sites.exitNodeSubnet, subnet)
                         )
                     );
 
@@ -427,7 +427,7 @@ export async function createSite(
                             exitNodeId,
                             name,
                             niceId: updatedNiceId!,
-                            subnet,
+                            exitNodeSubnet: subnet,
                             type,
                             pubKey: pubKey || null,
                             status: "approved"
@@ -444,7 +444,7 @@ export async function createSite(
                             type,
                             dockerSocketEnabled: false,
                             online: true,
-                            subnet: "0.0.0.0/32",
+                            exitNodeSubnet: "0.0.0.0/32",
                             status: "approved"
                         })
                         .returning();

server/routers/site/listSites.ts

@@ -125,7 +125,7 @@ function querySitesBase() {
             niceId: sites.niceId,
             name: sites.name,
             pubKey: sites.pubKey,
-            subnet: sites.subnet,
+            subnet: sites.exitNodeSubnet,
             megabytesIn: sites.megabytesIn,
             megabytesOut: sites.megabytesOut,
             orgName: orgs.name,

server/routers/site/pickSiteDefaults.ts

@@ -43,7 +43,6 @@ const PickSiteDefaultsResponseDataSchema = z.object({
     clientAddress: z.string().optional()
 });
 
-
 registry.registerPath({
     method: "get",
     path: "/org/{orgId}/pick-site-defaults",
@@ -60,7 +59,9 @@ registry.registerPath({
             description: "Successful response",
             content: {
                 "application/json": {
-                    schema: createApiResponseSchema(PickSiteDefaultsResponseDataSchema)
+                    schema: createApiResponseSchema(
+                        PickSiteDefaultsResponseDataSchema
+                    )
                 }
             }
         }
@@ -108,7 +109,7 @@ export async function pickSiteDefaults(
         // list all of the sites on that exit node
         const sitesQuery = await db
             .select({
-                subnet: sites.subnet
+                subnet: sites.exitNodeSubnet
             })
             .from(sites)
             .where(eq(sites.exitNodeId, randomExitNode.exitNodeId));

server/routers/target/createTarget.ts

@@ -263,7 +263,7 @@ export async function createTarget(
                 // make sure the target is within the site subnet
                 if (
                     site.type == "wireguard" &&
-                    !isIpInCidr(targetData.ip, site.subnet!)
+                    !isIpInCidr(targetData.ip, site.exitNodeSubnet!)
                 ) {
                     return next(
                         createHttpError(

server/routers/ws/messageHandlers.ts

@@ -5,7 +5,7 @@ import {
     handleNewtGetConfigMessage,
     handleDockerStatusMessage,
     handleDockerContainersMessage,
-    handleNewtPingRequestMessage,
+    handleNewtExitNodesRequestMessage,
     handleApplyBlueprintMessage,
     handleNewtPingMessage,
     startNewtOfflineChecker,
@@ -22,7 +22,8 @@ import {
     handleOlmDisconnectingMessage,
     handleOlmServerInitAddPeerHandshake,
     handleOlmLocalMessage,
-    handleOlmUnLocalMessage
+    handleOlmUnLocalMessage,
+    handleOlmExitNodesRequestMessage
 } from "../olm";
 import { handleHealthcheckStatusMessage } from "../target";
 import { handleRoundTripMessage } from "./handleRoundTripMessage";
@@ -37,6 +38,7 @@ export const messageHandlers: Record<string, MessageHandler> = {
     "olm/wg/local": handleOlmLocalMessage,
     "olm/wg/unlocal": handleOlmUnLocalMessage,
     "olm/ping": handleOlmPingMessage,
+    "olm/ping/request": handleOlmExitNodesRequestMessage,
     "olm/disconnecting": handleOlmDisconnectingMessage,
     "newt/disconnecting": handleNewtDisconnectingMessage,
     "newt/ping": handleNewtPingMessage,
@@ -45,7 +47,7 @@ export const messageHandlers: Record<string, MessageHandler> = {
     "newt/receive-bandwidth": handleReceiveBandwidthMessage,
     "newt/socket/status": handleDockerStatusMessage,
     "newt/socket/containers": handleDockerContainersMessage,
-    "newt/ping/request": handleNewtPingRequestMessage,
+    "newt/ping/request": handleNewtExitNodesRequestMessage,
     "newt/blueprint/apply": handleApplyBlueprintMessage,
     "newt/healthcheck/status": handleHealthcheckStatusMessage,
     "ws/round-trip/complete": handleRoundTripMessage

src/app/[orgId]/settings/sites/[niceId]/credentials/page.tsx

@@ -72,8 +72,7 @@ export default function CredentialsPage() {
     const { data: latestVersions } = useQuery(
         productUpdatesQueries.latestVersion(true)
     );
-    const newtVersion =
-        latestVersions?.data?.newt?.latestVersion ?? "latest";
+    const newtVersion = latestVersions?.data?.newt?.latestVersion ?? "latest";
 
     // Fetch site defaults for wireguard sites to show in obfuscated config
     useEffect(() => {
@@ -354,7 +353,7 @@ export default function CredentialsPage() {
                                             text={generateObfuscatedWireGuardConfig(
                                                 {
                                                     subnet:
-                                                        site?.subnet ||
+                                                        site?.exitNodeSubnet ||
                                                         siteDefaults?.subnet ||
                                                         null,
                                                     address:

src/app/auth/login/page.tsx

@@ -91,7 +91,6 @@ export default async function Page(props: {
     let loginIdps: LoginFormIDP[] = [];
     let lastUsedIdpForSmartLogin: (LoginFormIDP & { orgId?: string }) | null =
         null;
-
     if (!useSmartLogin) {
         // Load IdPs for DashboardLoginForm (OSS or org-only IdP mode)
         if (build === "oss" || env.app.identityProviderMode !== "org") {
@@ -118,12 +117,12 @@ export default async function Page(props: {
                     `/idp/${persistedData.idpId}`
                 );
 
-                const res = idpRes.data.data;
+                const idp = idpRes.data.data.idp;
 
                 lastUsedIdpForSmartLogin = {
-                    idpId: res.idp.idpId,
-                    name: res.idp.name,
-                    variant: res.idpOidcConfig?.variant ?? res.idp.type,
+                    idpId: idp.idpId,
+                    name: idp.name,
+                    variant: idp.type,
                     orgId: persistedData.orgId,
                     lastUsed: true
                 };

src/components/IdpLoginButtons.tsx

@@ -114,52 +114,59 @@ export default function IdpLoginButtons({
 
             <div className="space-y-4">
                 {params.get("gotoapp") ? (
-                    <Button
-                        type="button"
-                        className="w-full"
-                        onClick={() => {
-                            goToApp();
-                        }}
-                    >
-                        {t("continueToApplication")}
-                    </Button>
+                    <>
+                        <Button
+                            type="button"
+                            className="w-full"
+                            onClick={() => {
+                                goToApp();
+                            }}
+                        >
+                            {t("continueToApplication")}
+                        </Button>
+                    </>
                 ) : (
-                    idps.map((idp) => {
-                        const effectiveType =
-                            idp.variant || idp.name.toLowerCase();
+                    <>
+                        {idps.map((idp) => {
+                            const effectiveType =
+                                idp.variant || idp.name.toLowerCase();
 
-                        return (
-                            <div className="w-full relative" key={idp.idpId}>
-                                <Button
+                            return (
+                                <div
+                                    className="w-full relative"
                                     key={idp.idpId}
-                                    type="button"
-                                    variant="outline"
-                                    className="w-full inline-flex items-center space-x-2  after:absolute after:inset-0 after:z-10"
-                                    onClick={() => {
-                                        startTransition(() =>
-                                            loginWithIdp(idp.idpId)
-                                        );
-                                    }}
-                                    disabled={loading}
-                                    loading={loading}
                                 >
-                                    <IdpTypeIcon
-                                        type={effectiveType}
-                                        size={16}
-                                    />
-                                    <span>{idp.name}</span>
-                                </Button>
+                                    <Button
+                                        key={idp.idpId}
+                                        type="button"
+                                        variant="outline"
+                                        className="w-full inline-flex items-center space-x-2  after:absolute after:inset-0 after:z-10"
+                                        onClick={() => {
+                                            startTransition(() =>
+                                                loginWithIdp(idp.idpId)
+                                            );
+                                        }}
+                                        disabled={loading}
+                                        loading={loading}
+                                    >
+                                        <IdpTypeIcon
+                                            type={effectiveType}
+                                            size={16}
+                                        />
+                                        <span>{idp.name}</span>
+                                    </Button>
 
-                                {idp.lastUsed && (
-                                    <div className="absolute inset-0">
-                                        <span className="absolute top-0 right-0 text-xs bg-primary text-primary-foreground rounded-bl-sm rounded-tr-sm px-2 py-0.5">
-                                            {t("idpLastUsed")}
-                                        </span>
-                                    </div>
-                                )}
-                            </div>
-                        );
-                    })
+                                    {idp.lastUsed && (
+                                        <div className="absolute inset-0">
+                                            <span className="absolute top-0 right-0 text-xs bg-primary text-primary-foreground rounded-bl-sm rounded-tr-sm px-2 py-0.5">
+                                                {t("idpLastUsed")}
+                                            </span>
+                                        </div>
+                                    )}
+                                </div>
+                            );
+                        })}
+                    </>
                 )}
             </div>
         </div>

src/components/IdpTypeIcon.tsx

@@ -23,8 +23,6 @@ export default function IdpTypeIcon({
 }: Props) {
     const effectiveType = (variant || type || "").toLowerCase();
 
-    console.log(`[IdpTypeIcon]`, { effectiveType, variant, type });
-
     let src: string | null = null;
     let defaultAlt = "";
 

src/components/OrgIdpTable.tsx

@@ -1,6 +1,17 @@
 "use client";
 
-import ConfirmDeleteDialog from "@app/components/ConfirmDeleteDialog";
+import { ColumnDef } from "@tanstack/react-table";
+import { ExtendedColumnDef } from "@app/components/ui/data-table";
+import { IdpDataTable } from "@app/components/OrgIdpDataTable";
+import { Button } from "@app/components/ui/button";
+import {
+    Command,
+    CommandEmpty,
+    CommandGroup,
+    CommandInput,
+    CommandItem,
+    CommandList
+} from "@app/components/ui/command";
 import {
     Credenza,
     CredenzaBody,
@@ -11,42 +22,37 @@ import {
     CredenzaHeader,
     CredenzaTitle
 } from "@app/components/Credenza";
-import { isIdpGlobalModeBannerVisible } from "@app/components/IdpGlobalModeBanner";
-import IdpTypeBadge from "@app/components/IdpTypeBadge";
-import IdpTypeIcon from "@app/components/IdpTypeIcon";
-import { IdpDataTable } from "@app/components/OrgIdpDataTable";
-import { Badge } from "@app/components/ui/badge";
-import { Button } from "@app/components/ui/button";
 import {
-    Command,
-    CommandEmpty,
-    CommandGroup,
-    CommandInput,
-    CommandItem,
-    CommandList
-} from "@app/components/ui/command";
-import { ExtendedColumnDef } from "@app/components/ui/data-table";
+    ArrowRight,
+    ArrowUpDown,
+    MoreHorizontal
+} from "lucide-react";
+import { useMemo, useState } from "react";
+import ConfirmDeleteDialog from "@app/components/ConfirmDeleteDialog";
+import { toast } from "@app/hooks/useToast";
+import { formatAxiosError } from "@app/lib/api";
+import { createApiClient } from "@app/lib/api";
+import { useEnvContext } from "@app/hooks/useEnvContext";
+import { useUserContext } from "@app/hooks/useUserContext";
+import { useRouter } from "next/navigation";
 import {
     DropdownMenu,
     DropdownMenuContent,
     DropdownMenuItem,
     DropdownMenuTrigger
 } from "@app/components/ui/dropdown-menu";
-import { useEnvContext } from "@app/hooks/useEnvContext";
-import { usePaidStatus } from "@app/hooks/usePaidStatus";
-import { toast } from "@app/hooks/useToast";
-import { useUserContext } from "@app/hooks/useUserContext";
-import { createApiClient, formatAxiosError } from "@app/lib/api";
-import { cn } from "@app/lib/cn";
-import { tierMatrix } from "@server/lib/billing/tierMatrix";
-import type { ListUserAdminOrgIdpsResponse } from "@server/routers/orgIdp/types";
-import { useQuery } from "@tanstack/react-query";
-import { ArrowRight, ArrowUpDown, MoreHorizontal } from "lucide-react";
-import { useTranslations } from "next-intl";
 import Link from "next/link";
-import { useRouter } from "next/navigation";
-import { useMemo, useState } from "react";
+import { useTranslations } from "next-intl";
+import IdpTypeBadge from "@app/components/IdpTypeBadge";
+import IdpTypeIcon from "@app/components/IdpTypeIcon";
+import { useQuery } from "@tanstack/react-query";
 import { useDebounce } from "use-debounce";
+import type { ListUserAdminOrgIdpsResponse } from "@server/routers/orgIdp/types";
+import { cn } from "@app/lib/cn";
+import { Badge } from "@app/components/ui/badge";
+import { usePaidStatus } from "@app/hooks/usePaidStatus";
+import { tierMatrix } from "@server/lib/billing/tierMatrix";
+import { isIdpGlobalModeBannerVisible } from "@app/components/IdpGlobalModeBanner";
 
 export type IdpRow = {
     idpId: number;
@@ -477,17 +483,15 @@ export default function IdpTable({ idps, orgId }: Props) {
                                                     {group.name}
                                                 </div>
                                                 <div className="mt-1 flex flex-wrap gap-1">
-                                                    {group.sources.map(
-                                                        (src) => (
-                                                            <Badge
-                                                                key={src.orgId}
-                                                                variant="secondary"
-                                                                className="max-w-full truncate font-normal"
-                                                            >
-                                                                {src.orgName}
-                                                            </Badge>
-                                                        )
-                                                    )}
+                                                    {group.sources.map((src) => (
+                                                        <Badge
+                                                            key={src.orgId}
+                                                            variant="secondary"
+                                                            className="max-w-full truncate font-normal"
+                                                        >
+                                                            {src.orgName}
+                                                        </Badge>
+                                                    ))}
                                                 </div>
                                             </div>
                                         </CommandItem>

src/components/UptimeBar.tsx

@@ -134,7 +134,9 @@ export default function UptimeBar({
 
     if (!data) return null;
 
-    const allNoData = data.days.every((d) => d.status === "no_data");
+    const allNoData = data.days.every(
+        (d) => d.status === "no_data" || d.status === "unknown"
+    );
 
     return (
         <div className={cn("space-y-3", className)}>

src/components/UptimeMiniBar.tsx

@@ -124,7 +124,9 @@ export function UptimeMiniBar({
 
     if (!data) return null;
 
-    const allNoData = data.days.every((d) => d.status === "no_data");
+    const allNoData = data.days.every(
+        (d) => d.status === "no_data" || d.status === "unknown"
+    );
 
     return (
         <div className="flex items-center gap-2">