Further optimizations

server/routers/olm/getOlmToken.ts

@@ -8,7 +8,7 @@ import {
     ExitNode,
     exitNodes,
     sites,
-    clientSitesAssociationsCache,
+    clientSitesAssociationsCache
 } from "@server/db";
 import { olms } from "@server/db";
 import HttpCode from "@server/types/HttpCode";
@@ -28,6 +28,7 @@ import { verifyPassword } from "@server/auth/password";
 import logger from "@server/logger";
 import config from "@server/lib/config";
 import { APP_VERSION } from "@server/lib/consts";
+import { build } from "@server/build";
 
 export const olmGetTokenBodySchema = z.object({
     olmId: z.string(),
@@ -220,6 +221,22 @@ export async function getOlmToken(
             )
             .where(eq(clientSitesAssociationsCache.clientId, clientIdToUse!));
 
+        if (clientSites.length > 250 && build == "saas") {
+            // set all of the cache rows isJitMode to true
+            await db
+                .update(clientSitesAssociationsCache)
+                .set({ isJitMode: true })
+                .where(
+                    and(
+                        eq(
+                            clientSitesAssociationsCache.clientId,
+                            clientIdToUse!
+                        ),
+                        eq(clientSitesAssociationsCache.isJitMode, false)
+                    )
+                );
+        }
+
         // Extract unique exit node IDs
         const exitNodeIds = Array.from(
             new Set(

server/routers/olm/handleOlmRegisterMessage.ts

@@ -7,7 +7,7 @@ import {
     olms,
     sites
 } from "@server/db";
-import { count, eq } from "drizzle-orm";
+import { and, count, eq, ne, or } from "drizzle-orm";
 import logger from "@server/logger";
 import { checkOrgAccessPolicy } from "#dynamic/lib/checkOrgAccessPolicy";
 import { validateSessionToken } from "@server/auth/sessions/app";
@@ -301,7 +301,18 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
                 isRelayed: relay == true,
                 isJitMode: jitMode
             })
-            .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
+            .where(
+                and(
+                    eq(clientSitesAssociationsCache.clientId, client.clientId),
+                    or(
+                        ne(
+                            clientSitesAssociationsCache.isRelayed,
+                            relay == true
+                        ),
+                        ne(clientSitesAssociationsCache.isJitMode, jitMode)
+                    )
+                )
+            );
     }
 
     // this prevents us from accepting a register from an olm that has not hole punched yet.