Handle JIT for ssh

server/db/pg/schema/schema.ts

@@ -720,6 +720,7 @@ export const clientSitesAssociationsCache = pgTable(
             .notNull(),
         siteId: integer("siteId").notNull(),
         isRelayed: boolean("isRelayed").notNull().default(false),
+        isJitMode: boolean("isJitMode").notNull().default(false),
         endpoint: varchar("endpoint"),
         publicKey: varchar("publicKey") // this will act as the session's public key for hole punching so we can track when it changes
     }

server/db/sqlite/schema/schema.ts

@@ -409,6 +409,9 @@ export const clientSitesAssociationsCache = sqliteTable(
         isRelayed: integer("isRelayed", { mode: "boolean" })
             .notNull()
             .default(false),
+        isJitMode: integer("isJitMode", { mode: "boolean" })
+            .notNull()
+            .default(false),
         endpoint: text("endpoint"),
         publicKey: text("publicKey") // this will act as the session's public key for hole punching so we can track when it changes
     }

server/private/routers/ssh/signSshKey.ts

@@ -29,7 +29,6 @@ import HttpCode from "@server/types/HttpCode";
 import createHttpError from "http-errors";
 import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
-import { OpenAPITags, registry } from "@server/openApi";
 import { eq, or, and } from "drizzle-orm";
 import { canUserAccessSiteResource } from "@server/auth/canUserAccessSiteResource";
 import { signPublicKey, getOrgCAKeys } from "@server/lib/sshCA";
@@ -64,6 +63,7 @@ export type SignSshKeyResponse = {
     sshUsername: string;
     sshHost: string;
     resourceId: number;
+    siteId: number;
     keyId: string;
     validPrincipals: string[];
     validAfter: string;
@@ -453,6 +453,7 @@ export async function signSshKey(
                 sshUsername: usernameToUse,
                 sshHost: sshHost,
                 resourceId: resource.siteResourceId,
+                siteId: resource.siteId,
                 keyId: cert.keyId,
                 validPrincipals: cert.validPrincipals,
                 validAfter: cert.validAfter.toISOString(),

server/private/routers/ws/ws.ts

@@ -76,7 +76,7 @@ const processMessage = async (
             clientId,
             message.type, // Pass message type for granular limiting
             100, // max requests per window
-            20, // max requests per message type per window
+            100, // max requests per message type per window
             60 * 1000 // window in milliseconds
         );
         if (rateLimitResult.isLimited) {

server/routers/newt/buildConfiguration.ts

@@ -1,4 +1,15 @@
-import { clients, clientSiteResourcesAssociationsCache, clientSitesAssociationsCache, db, ExitNode, resources, Site, siteResources, targetHealthCheck, targets } from "@server/db";
+import {
+    clients,
+    clientSiteResourcesAssociationsCache,
+    clientSitesAssociationsCache,
+    db,
+    ExitNode,
+    resources,
+    Site,
+    siteResources,
+    targetHealthCheck,
+    targets
+} from "@server/db";
 import logger from "@server/logger";
 import { initPeerAddHandshake, updatePeer } from "../olm/peers";
 import { eq, and } from "drizzle-orm";
@@ -69,6 +80,7 @@ export async function buildClientConfigurationForNewtClient(
                     //         )
                     //     );
 
+                    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, {
@@ -103,6 +115,7 @@ export async function buildClientConfigurationForNewtClient(
                                 }
                             }
                         );
+                    }
 
                     return {
                         publicKey: client.clients.pubKey!,

server/routers/olm/handleOlmRegisterMessage.ts

@@ -17,6 +17,8 @@ import { getUserDeviceName } from "@server/db/names";
 import { buildSiteConfigurationForOlmClient } from "./buildConfiguration";
 import { OlmErrorCodes, sendOlmError } from "./error";
 import { handleFingerprintInsertion } from "./fingerprintingUtils";
+import { Alias } from "@server/lib/ip";
+import { build } from "@server/build";
 
 export const handleOlmRegisterMessage: MessageHandler = async (context) => {
     logger.info("Handling register olm message!");
@@ -207,6 +209,32 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         }
     }
 
+    // Get all sites data
+    const sitesCountResult = await db
+        .select({ count: count() })
+        .from(sites)
+        .innerJoin(
+            clientSitesAssociationsCache,
+            eq(sites.siteId, clientSitesAssociationsCache.siteId)
+        )
+        .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
+
+    // Extract the count value from the result array
+    const sitesCount =
+        sitesCountResult.length > 0 ? sitesCountResult[0].count : 0;
+
+    // Prepare an array to store site configurations
+    logger.debug(`Found ${sitesCount} sites for client ${client.clientId}`);
+
+    let jitMode = false;
+    if (sitesCount > 250 && build == "saas") {
+        // THIS IS THE MAX ON THE BUSINESS TIER
+        // we have too many sites
+        // If we have too many sites we need to drop into fully JIT mode by not sending any of the sites
+        logger.info("Too many sites (%d), dropping into JIT mode", sitesCount)
+        jitMode = true;
+    }
+
     logger.debug(
         `Olm client ID: ${client.clientId}, Public Key: ${publicKey}, Relay: ${relay}`
     );
@@ -233,28 +261,12 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         await db
             .update(clientSitesAssociationsCache)
             .set({
-                isRelayed: relay == true
+                isRelayed: relay == true,
+                isJitMode: jitMode
             })
             .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
     }
 
-    // Get all sites data
-    const sitesCountResult = await db
-        .select({ count: count() })
-        .from(sites)
-        .innerJoin(
-            clientSitesAssociationsCache,
-            eq(sites.siteId, clientSitesAssociationsCache.siteId)
-        )
-        .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
-
-    // Extract the count value from the result array
-    const sitesCount =
-        sitesCountResult.length > 0 ? sitesCountResult[0].count : 0;
-
-    // Prepare an array to store site configurations
-    logger.debug(`Found ${sitesCount} sites for client ${client.clientId}`);
-
     // this prevents us from accepting a register from an olm that has not hole punched yet.
     // the olm will pump the register so we can keep checking
     // TODO: I still think there is a better way to do this rather than locking it out here but ???
@@ -265,18 +277,25 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         return;
     }
 
+    let siteConfigurations: {
+        siteId: number;
+        name: string;
+        endpoint: string;
+        publicKey: string | null;
+        serverIP: string | null;
+        serverPort: number | null;
+        remoteSubnets: string[];
+        aliases: Alias[];
+    }[] = [];
+
+    if (!jitMode) {
         // NOTE: its important that the client here is the old client and the public key is the new key
-    const siteConfigurations = await buildSiteConfigurationForOlmClient(
+        siteConfigurations = await buildSiteConfigurationForOlmClient(
             client,
             publicKey,
             relay
         );
-
-    // REMOVED THIS SO IT CREATES THE INTERFACE AND JUST WAITS FOR THE SITES
-    // if (siteConfigurations.length === 0) {
-    //     logger.warn("No valid site configurations found");
-    //     return;
-    // }
+    }
 
     // Return connect message with all site configurations
     return {

server/routers/olm/handleOlmRelayMessage.ts

@@ -18,7 +18,7 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
     }
 
     if (!olm.clientId) {
-        logger.warn("Olm has no site!"); // TODO: Maybe we create the site here?
+        logger.warn("Olm has no client!");
         return;
     }
 
@@ -41,7 +41,7 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
         return;
     }
 
-    const { siteId } = message.data;
+    const { siteId, chainId } = message.data;
 
     // Get the site
     const [site] = await db
@@ -90,7 +90,8 @@ export const handleOlmRelayMessage: MessageHandler = async (context) => {
             data: {
                 siteId: siteId,
                 relayEndpoint: exitNode.endpoint,
-                relayPort: config.getRawConfig().gerbil.clients_start_port
+                relayPort: config.getRawConfig().gerbil.clients_start_port,
+                chainId
             }
         },
         broadcast: false,

server/routers/olm/handleOlmServerInitAddPeerHandshake.ts

@@ -0,0 +1,241 @@
+import {
+    clientSiteResourcesAssociationsCache,
+    clientSitesAssociationsCache,
+    db,
+    exitNodes,
+    Site,
+    siteResources
+} from "@server/db";
+import { MessageHandler } from "@server/routers/ws";
+import { clients, Olm, sites } from "@server/db";
+import { and, eq, or } from "drizzle-orm";
+import logger from "@server/logger";
+import { initPeerAddHandshake } from "./peers";
+
+export const handleOlmServerInitAddPeerHandshake: MessageHandler = async (
+    context
+) => {
+    logger.info("Handling register olm message!");
+    const { message, client: c, sendToClient } = context;
+    const olm = c as Olm;
+
+    if (!olm) {
+        logger.warn("Olm not found");
+        return;
+    }
+
+    if (!olm.clientId) {
+        logger.warn("Olm has no client!"); // TODO: Maybe we create the site here?
+        return;
+    }
+
+    const clientId = olm.clientId;
+
+    const [client] = await db
+        .select()
+        .from(clients)
+        .where(eq(clients.clientId, clientId))
+        .limit(1);
+
+    if (!client) {
+        logger.warn("Client not found");
+        return;
+    }
+
+    const { siteId, resourceId, chainId } = message.data;
+
+    let site: Site | null = null;
+    if (siteId) {
+        // get the site
+        const [siteRes] = await db
+            .select()
+            .from(sites)
+            .where(eq(sites.siteId, siteId))
+            .limit(1);
+        if (siteRes) {
+            site = siteRes;
+        }
+    }
+
+    if (resourceId && !site) {
+        const resources = await db
+            .select()
+            .from(siteResources)
+            .where(
+                and(
+                    or(
+                        eq(siteResources.niceId, resourceId),
+                        eq(siteResources.alias, resourceId)
+                    ),
+                    eq(siteResources.orgId, client.orgId)
+                )
+            );
+
+        if (!resources || resources.length === 0) {
+            logger.error(`handleOlmServerPeerAddMessage: Resource not found`);
+            // cancel the request from the olm side to not keep doing this
+            await sendToClient(
+                olm.olmId,
+                {
+                    type: "olm/wg/peer/chain/cancel",
+                    data: {
+                        chainId
+                    }
+                },
+                { incrementConfigVersion: false }
+            ).catch((error) => {
+                logger.warn(`Error sending message:`, error);
+            });
+            return;
+        }
+
+        if (resources.length > 1) {
+            // error but this should not happen because the nice id cant contain a dot and the alias has to have a dot and both have to be unique within the org so there should never be multiple matches
+            logger.error(
+                `handleOlmServerPeerAddMessage: Multiple resources found matching the criteria`
+            );
+            return;
+        }
+
+        const resource = resources[0];
+
+        const currentResourceAssociationCaches = await db
+            .select()
+            .from(clientSiteResourcesAssociationsCache)
+            .where(
+                and(
+                    eq(
+                        clientSiteResourcesAssociationsCache.siteResourceId,
+                        resource.siteResourceId
+                    ),
+                    eq(
+                        clientSiteResourcesAssociationsCache.clientId,
+                        client.clientId
+                    )
+                )
+            );
+
+        if (currentResourceAssociationCaches.length === 0) {
+            logger.error(
+                `handleOlmServerPeerAddMessage: Client ${client.clientId} does not have access to resource ${resource.siteResourceId}`
+            );
+            // cancel the request from the olm side to not keep doing this
+            await sendToClient(
+                olm.olmId,
+                {
+                    type: "olm/wg/peer/chain/cancel",
+                    data: {
+                        chainId
+                    }
+                },
+                { incrementConfigVersion: false }
+            ).catch((error) => {
+                logger.warn(`Error sending message:`, error);
+            });
+            return;
+        }
+
+        const siteIdFromResource = resource.siteId;
+
+        // get the site
+        const [siteRes] = await db
+            .select()
+            .from(sites)
+            .where(eq(sites.siteId, siteIdFromResource));
+        if (!siteRes) {
+            logger.error(
+                `handleOlmServerPeerAddMessage: Site with ID ${site} not found`
+            );
+            return;
+        }
+
+        site = siteRes;
+    }
+
+    if (!site) {
+        logger.error(`handleOlmServerPeerAddMessage: Site not found`);
+        return;
+    }
+
+    // check if the client can access this site using the cache
+    const currentSiteAssociationCaches = await db
+        .select()
+        .from(clientSitesAssociationsCache)
+        .where(
+            and(
+                eq(clientSitesAssociationsCache.clientId, client.clientId),
+                eq(clientSitesAssociationsCache.siteId, site.siteId)
+            )
+        );
+
+    if (currentSiteAssociationCaches.length === 0) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Client ${client.clientId} does not have access to site ${site.siteId}`
+        );
+        // cancel the request from the olm side to not keep doing this
+        await sendToClient(
+            olm.olmId,
+            {
+                type: "olm/wg/peer/chain/cancel",
+                data: {
+                    chainId
+                }
+            },
+            { incrementConfigVersion: false }
+        ).catch((error) => {
+            logger.warn(`Error sending message:`, error);
+        });
+        return;
+    }
+
+    if (!site.exitNodeId) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Site with ID ${site.siteId} has no exit node`
+        );
+        // cancel the request from the olm side to not keep doing this
+        await sendToClient(
+            olm.olmId,
+            {
+                type: "olm/wg/peer/chain/cancel",
+                data: {
+                    chainId
+                }
+            },
+            { incrementConfigVersion: false }
+        ).catch((error) => {
+            logger.warn(`Error sending message:`, error);
+        });
+        return;
+    }
+
+    // get the exit node from the side
+    const [exitNode] = await db
+        .select()
+        .from(exitNodes)
+        .where(eq(exitNodes.exitNodeId, site.exitNodeId));
+
+    if (!exitNode) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Site with ID ${site.siteId} has no exit node`
+        );
+        return;
+    }
+
+    // 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.clientId,
+        {
+            siteId: site.siteId,
+            exitNode: {
+                publicKey: exitNode.publicKey,
+                endpoint: exitNode.endpoint
+            }
+        },
+        olm.olmId,
+        chainId
+    );
+
+    return;
+};

server/routers/olm/handleOlmServerPeerAddMessage.ts

@@ -54,7 +54,7 @@ export const handleOlmServerPeerAddMessage: MessageHandler = async (
         return;
     }
 
-    const { siteId } = message.data;
+    const { siteId, chainId } = message.data;
 
     // get the site
     const [site] = await db
@@ -179,7 +179,8 @@ export const handleOlmServerPeerAddMessage: MessageHandler = async (
                 ),
                 aliases: generateAliasConfig(
                     allSiteResources.map(({ siteResources }) => siteResources)
-                )
+                ),
+                chainId: chainId,
             }
         },
         broadcast: false,

server/routers/olm/handleOlmUnRelayMessage.ts

@@ -17,7 +17,7 @@ export const handleOlmUnRelayMessage: MessageHandler = async (context) => {
     }
 
     if (!olm.clientId) {
-        logger.warn("Olm has no site!"); // TODO: Maybe we create the site here?
+        logger.warn("Olm has no client!");
         return;
     }
 
@@ -40,7 +40,7 @@ export const handleOlmUnRelayMessage: MessageHandler = async (context) => {
         return;
     }
 
-    const { siteId } = message.data;
+    const { siteId, chainId } = message.data;
 
     // Get the site
     const [site] = await db
@@ -87,7 +87,8 @@ export const handleOlmUnRelayMessage: MessageHandler = async (context) => {
             type: "olm/wg/peer/unrelay",
             data: {
                 siteId: siteId,
-                endpoint: site.endpoint
+                endpoint: site.endpoint,
+                chainId
             }
         },
         broadcast: false,

server/routers/olm/index.ts

@@ -11,3 +11,4 @@ export * from "./handleOlmServerPeerAddMessage";
 export * from "./handleOlmUnRelayMessage";
 export * from "./recoverOlmWithFingerprint";
 export * from "./handleOlmDisconnectingMessage";
+export * from "./handleOlmServerInitAddPeerHandshake";

server/routers/olm/peers.ts

@@ -1,8 +1,8 @@
 import { sendToClient } from "#dynamic/routers/ws";
-import { db, olms } from "@server/db";
+import { clientSitesAssociationsCache, db, olms } from "@server/db";
 import config from "@server/lib/config";
 import logger from "@server/logger";
-import { eq } from "drizzle-orm";
+import { and, eq } from "drizzle-orm";
 import { Alias } from "yaml";
 
 export async function addPeer(
@@ -149,7 +149,8 @@ export async function initPeerAddHandshake(
             endpoint: string;
         };
     },
-    olmId?: string
+    olmId?: string,
+    chainId?: string
 ) {
     if (!olmId) {
         const [olm] = await db
@@ -173,7 +174,8 @@ export async function initPeerAddHandshake(
                     publicKey: peer.exitNode.publicKey,
                     relayPort: config.getRawConfig().gerbil.clients_start_port,
                     endpoint: peer.exitNode.endpoint
-                }
+                },
+                chainId
             }
         },
         { incrementConfigVersion: true }
@@ -181,6 +183,17 @@ export async function initPeerAddHandshake(
         logger.warn(`Error sending message:`, error);
     });
 
+    // update the clientSiteAssociationsCache to make the isJitMode flag false so that JIT mode is disabled for this site if it restarts or something after the connection
+    await db
+        .update(clientSitesAssociationsCache)
+        .set({ isJitMode: false })
+        .where(
+            and(
+                eq(clientSitesAssociationsCache.clientId, clientId),
+                eq(clientSitesAssociationsCache.siteId, peer.siteId)
+            )
+        );
+
     logger.info(
         `Initiated peer add handshake for site ${peer.siteId} to olm ${olmId}`
     );

server/routers/ws/messageHandlers.ts

@@ -15,7 +15,8 @@ import {
     startOlmOfflineChecker,
     handleOlmServerPeerAddMessage,
     handleOlmUnRelayMessage,
-    handleOlmDisconnecingMessage
+    handleOlmDisconnecingMessage,
+    handleOlmServerInitAddPeerHandshake
 } from "../olm";
 import { handleHealthcheckStatusMessage } from "../target";
 import { handleRoundTripMessage } from "./handleRoundTripMessage";
@@ -23,6 +24,7 @@ import { MessageHandler } from "./types";
 
 export const messageHandlers: Record<string, MessageHandler> = {
     "olm/wg/server/peer/add": handleOlmServerPeerAddMessage,
+    "olm/wg/server/peer/init": handleOlmServerInitAddPeerHandshake,
     "olm/wg/register": handleOlmRegisterMessage,
     "olm/wg/relay": handleOlmRelayMessage,
     "olm/wg/unrelay": handleOlmUnRelayMessage,