Working on jit

server/lib/telemetry.ts

@@ -21,13 +21,13 @@ class TelemetryClient {
         this.enabled = enabled;
         const dev = process.env.ENVIRONMENT !== "prod";
 
-        if (dev) {
-            return;
-        }
+        // if (dev) {
+        //     return;
+        // }
 
-        if (build === "saas") {
-            return;
-        }
+        // if (build === "saas") {
+        //     return;
+        // }
 
         if (this.enabled) {
             this.client = new PostHog(

server/routers/olm/buildConfiguration.ts

@@ -1,5 +1,17 @@
-import { Client, clientSiteResourcesAssociationsCache, clientSitesAssociationsCache, db, exitNodes, siteResources, sites } from "@server/db";
-import { generateAliasConfig, generateRemoteSubnets } from "@server/lib/ip";
+import {
+    Client,
+    clientSiteResourcesAssociationsCache,
+    clientSitesAssociationsCache,
+    db,
+    exitNodes,
+    siteResources,
+    sites
+} from "@server/db";
+import {
+    Alias,
+    generateAliasConfig,
+    generateRemoteSubnets
+} from "@server/lib/ip";
 import logger from "@server/logger";
 import { and, eq } from "drizzle-orm";
 import { addPeer, deletePeer } from "../newt/peers";
@@ -8,9 +20,19 @@ import config from "@server/lib/config";
 export async function buildSiteConfigurationForOlmClient(
     client: Client,
     publicKey: string | null,
-    relay: boolean
+    relay: boolean,
+    jitMode: boolean = false
 ) {
-    const siteConfigurations = [];
+    const siteConfigurations: {
+        siteId: number;
+        name: string | null;
+        endpoint: string | null;
+        publicKey: string | null;
+        serverIP: string | null;
+        serverPort: number | null;
+        remoteSubnets: string[];
+        aliases: Alias[];
+    }[] = [];
 
     // Get all sites data
     const sitesData = await db
@@ -27,6 +49,46 @@ export async function buildSiteConfigurationForOlmClient(
         sites: site,
         clientSitesAssociationsCache: association
     } of sitesData) {
+        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.clientId
+                    )
+                )
+            );
+
+        if (jitMode) {
+            // Add site configuration to the array
+            siteConfigurations.push({
+                siteId: site.siteId,
+                name: null, // this is just to sync the aliases
+                endpoint: null,
+                publicKey: null,
+                serverIP: null,
+                serverPort: null,
+                // remoteSubnets: generateRemoteSubnets(
+                //     allSiteResources.map(({ siteResources }) => siteResources)
+                // ),
+                remoteSubnets: [],
+                aliases: generateAliasConfig(
+                    allSiteResources.map(({ siteResources }) => siteResources)
+                )
+            });
+            continue;
+        }
+
         if (!site.exitNodeId) {
             logger.warn(
                 `Site ${site.siteId} does not have exit node, skipping`
@@ -103,26 +165,6 @@ export async function buildSiteConfigurationForOlmClient(
             relayEndpoint = `${exitNode.endpoint}:${config.getRawConfig().gerbil.clients_start_port}`;
         }
 
-        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.clientId
-                    )
-                )
-            );
-
         // Add site configuration to the array
         siteConfigurations.push({
             siteId: site.siteId,

server/routers/olm/handleOlmRegisterMessage.ts

@@ -226,12 +226,12 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
     // Prepare an array to store site configurations
     logger.debug(`Found ${sitesCount} sites for client ${client.clientId}`);
 
-    let jitMode = false;
+    let jitMode = true;
     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)
+        logger.info("Too many sites (%d), dropping into JIT mode", sitesCount);
         jitMode = true;
     }
 
@@ -277,25 +277,13 @@ 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
-        siteConfigurations = await buildSiteConfigurationForOlmClient(
-            client,
-            publicKey,
-            relay
-        );
-    }
+   // NOTE: its important that the client here is the old client and the public key is the new key
+    const siteConfigurations = await buildSiteConfigurationForOlmClient(
+        client,
+        publicKey,
+        relay,
+        jitMode
+    );
 
     // Return connect message with all site configurations
     return {