Allow JIT based on site or resource

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";

server/routers/olm/handleOlmRegisterMessage.ts

@@ -265,12 +265,14 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         return;
     }
 
-    // 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
-    );
+    // // 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
+    // );
+
+    const siteConfigurations: any = [];
 
     // REMOVED THIS SO IT CREATES THE INTERFACE AND JUST WAITS FOR THE SITES
     // if (siteConfigurations.length === 0) {

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;
     }
 

server/routers/olm/handleOlmServerInitAddPeerHandshake.ts

@@ -0,0 +1,142 @@
+import {
+    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 } = 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`);
+            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 siteIdFromResource = resources[0].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;
+    }
+
+    if (!site.exitNodeId) {
+        logger.error(
+            `handleOlmServerPeerAddMessage: Site with ID ${site.siteId} has no exit node`
+        );
+        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
+    );
+
+    return;
+};

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;
     }
 

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/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,