Attempt to handle creating/deleting clients and role

server/routers/auth/securityKey.ts

@@ -52,7 +52,7 @@ setInterval(async () => {
         await db
             .delete(webauthnChallenge)
             .where(lt(webauthnChallenge.expiresAt, now));
-        logger.debug("Cleaned up expired security key challenges");
+        // logger.debug("Cleaned up expired security key challenges");
     } catch (error) {
         logger.error("Failed to clean up expired security key challenges", error);
     }

server/routers/client/createClient.ts

@@ -24,18 +24,19 @@ import { isIpInCidr } from "@server/lib/ip";
 import { listExitNodes } from "#dynamic/lib/exitNodes";
 import { generateId } from "@server/auth/sessions/app";
 import { OpenAPITags, registry } from "@server/openApi";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
 
 const createClientParamsSchema = z.strictObject({
-        orgId: z.string()
-    });
+    orgId: z.string()
+});
 
 const createClientSchema = z.strictObject({
-        name: z.string().min(1).max(255),
-        olmId: z.string(),
-        secret: z.string(),
-        subnet: z.string(),
-        type: z.enum(["olm"])
-    });
+    name: z.string().min(1).max(255),
+    olmId: z.string(),
+    secret: z.string(),
+    subnet: z.string(),
+    type: z.enum(["olm"])
+});
 
 export type CreateClientBody = z.infer<typeof createClientSchema>;
 
@@ -186,6 +187,7 @@ 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);
@@ -204,7 +206,7 @@ export async function createClient(
                 );
             }
 
-            const [newClient] = await trx
+            [newClient] = await trx
                 .insert(clients)
                 .values({
                     exitNodeId: randomExitNode.exitNodeId,
@@ -244,13 +246,15 @@ export async function createClient(
                 dateCreated: moment().toISOString()
             });
 
-            return response<CreateClientResponse>(res, {
-                data: newClient,
-                success: true,
-                error: false,
-                message: "Site created successfully",
-                status: HttpCode.CREATED
-            });
+            await rebuildClientAssociationsFromClient(newClient, trx);
+        });
+
+        return response<CreateClientResponse>(res, {
+            data: newClient,
+            success: true,
+            error: false,
+            message: "Site created successfully",
+            status: HttpCode.CREATED
         });
     } catch (error) {
         logger.error(error);

server/routers/client/createUserClient.ts

@@ -21,6 +21,7 @@ import { isValidIP } from "@server/lib/validators";
 import { isIpInCidr } from "@server/lib/ip";
 import { listExitNodes } from "#dynamic/lib/exitNodes";
 import { OpenAPITags, registry } from "@server/openApi";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
 
 const paramsSchema = z
     .object({
@@ -191,6 +192,7 @@ 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);
@@ -209,7 +211,7 @@ export async function createUserClient(
                 );
             }
 
-            const [newClient] = await trx
+            [newClient] = await trx
                 .insert(clients)
                 .values({
                     exitNodeId: randomExitNode.exitNodeId,
@@ -232,13 +234,15 @@ export async function createUserClient(
                 clientId: newClient.clientId
             });
 
-            return response<CreateClientAndOlmResponse>(res, {
-                data: newClient,
-                success: true,
-                error: false,
-                message: "Site created successfully",
-                status: HttpCode.CREATED
-            });
+            await rebuildClientAssociationsFromClient(newClient, trx);
+        });
+
+        return response<CreateClientAndOlmResponse>(res, {
+            data: newClient,
+            success: true,
+            error: false,
+            message: "Site created successfully",
+            status: HttpCode.CREATED
         });
     } catch (error) {
         logger.error(error);

server/routers/client/deleteClient.ts

@@ -1,6 +1,6 @@
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
-import { db } from "@server/db";
+import { db, olms } from "@server/db";
 import { clients, clientSitesAssociationsCache } from "@server/db";
 import { eq } from "drizzle-orm";
 import response from "@server/lib/response";
@@ -9,10 +9,12 @@ import createHttpError from "http-errors";
 import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { OpenAPITags, registry } from "@server/openApi";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
+import { sendTerminateClient } from "./terminate";
 
 const deleteClientSchema = z.strictObject({
-        clientId: z.string().transform(Number).pipe(z.int().positive())
-    });
+    clientId: z.string().transform(Number).pipe(z.int().positive())
+});
 
 registry.registerPath({
     method: "delete",
@@ -68,19 +70,27 @@ export async function deleteClient(
         }
 
         await db.transaction(async (trx) => {
-            // Delete the client-site associations first
-            await trx
-                .delete(clientSitesAssociationsCache)
-                .where(eq(clientSitesAssociationsCache.clientId, clientId));
-
             // Then delete the client itself
-            await trx.delete(clients).where(eq(clients.clientId, clientId));
+            const [deletedClient] = await trx
+                .delete(clients)
+                .where(eq(clients.clientId, clientId))
+                .returning();
 
-            // this is a machine client
+            const [olm] = await trx
+                .select()
+                .from(olms)
+                .where(eq(olms.clientId, clientId))
+                .limit(1);
+
+            // this is a machine client so we also delete the olm
             if (!client.userId && client.olmId) {
-                await trx
-                    .delete(clients)
-                    .where(eq(clients.olmId, client.olmId));
+                await trx.delete(olms).where(eq(olms.olmId, client.olmId));
+            }
+
+            await rebuildClientAssociationsFromClient(deletedClient, trx);
+
+            if (olm) {
+                await sendTerminateClient(deletedClient.clientId, olm.olmId); //  the olmId needs to be provided because it cant look it up after deletion
             }
         });
 

server/routers/client/terminate.ts

@@ -0,0 +1,22 @@
+import { sendToClient } from "#dynamic/routers/ws";
+import { db, olms } from "@server/db";
+import { eq } from "drizzle-orm";
+
+export async function sendTerminateClient(clientId: number, olmId?: string | null) {
+    if (!olmId) {
+        const [olm] = await db
+            .select()
+            .from(olms)
+            .where(eq(olms.clientId, clientId))
+            .limit(1);
+        if (!olm) {
+            throw new Error(`Olm with ID ${clientId} not found`);
+        }
+        olmId = olm.olmId;
+    }
+
+    await sendToClient(olmId, {
+        type: `olm/terminate`,
+        data: {}
+    });
+}

server/routers/olm/deleteUserOlm.ts

@@ -1,5 +1,5 @@
 import { NextFunction, Request, Response } from "express";
-import { db } from "@server/db";
+import { Client, db } from "@server/db";
 import { olms, clients, clientSitesAssociationsCache } from "@server/db";
 import { eq } from "drizzle-orm";
 import HttpCode from "@server/types/HttpCode";
@@ -9,6 +9,8 @@ import { z } from "zod";
 import { fromError } from "zod-validation-error";
 import logger from "@server/logger";
 import { OpenAPITags, registry } from "@server/openApi";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
+import { sendTerminateClient } from "../client/terminate";
 
 const paramsSchema = z
     .object({
@@ -54,20 +56,30 @@ export async function deleteUserOlm(
                 .from(clients)
                 .where(eq(clients.olmId, olmId));
 
-            // Delete client-site associations for each associated client
-            for (const client of associatedClients) {
-                await trx
-                    .delete(clientSitesAssociationsCache)
-                    .where(eq(clientSitesAssociationsCache.clientId, client.clientId));
-            }
-
+            let deletedClient: Client | null = null;
             // Delete all associated clients
             if (associatedClients.length > 0) {
-                await trx.delete(clients).where(eq(clients.olmId, olmId));
+                [deletedClient] = await trx
+                    .delete(clients)
+                    .where(eq(clients.olmId, olmId))
+                    .returning();
             }
 
             // Finally, delete the OLM itself
-            await trx.delete(olms).where(eq(olms.olmId, olmId));
+            const [olm] = await trx
+                .delete(olms)
+                .where(eq(olms.olmId, olmId))
+                .returning();
+
+            if (deletedClient) {
+                await rebuildClientAssociationsFromClient(deletedClient, trx);
+                if (olm) {
+                    await sendTerminateClient(
+                        deletedClient.clientId,
+                        olm.olmId
+                    ); //  the olmId needs to be provided because it cant look it up after deletion
+                }
+            }
         });
 
         return response(res, {

server/routers/olm/getUserOlm.ts

@@ -1,6 +1,6 @@
 import { NextFunction, Request, Response } from "express";
 import { db } from "@server/db";
-import { olms, clients, clientSites } from "@server/db";
+import { olms } from "@server/db";
 import { eq, and } from "drizzle-orm";
 import HttpCode from "@server/types/HttpCode";
 import createHttpError from "http-errors";

server/routers/olm/handleOlmRegisterMessage.ts

@@ -31,6 +31,7 @@ import {
     getNextAvailableClientSubnet
 } from "@server/lib/ip";
 import { generateRemoteSubnets } from "@server/lib/ip";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
 
 export const handleOlmRegisterMessage: MessageHandler = async (context) => {
     logger.info("Handling register olm message!");
@@ -60,6 +61,7 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
                     olm.name || "User Device",
                     // doNotCreateNewClient ? true : false
                     true // for now never create a new client automatically because we create the users clients when they are added to the org
+                    // this means that the rebuildClientAssociationsFromClient call below issue is not a problem 
                 );
 
             client = clientRes;
@@ -99,6 +101,12 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
             .from(clients)
             .where(eq(clients.clientId, olm.clientId))
             .limit(1);
+
+        [org] = await db
+            .select()
+            .from(orgs)
+            .where(eq(orgs.orgId, client.orgId))
+            .limit(1);
     }
 
     if (!client) {
@@ -205,13 +213,6 @@ export const handleOlmRegisterMessage: MessageHandler = async (context) => {
         `Found ${sitesData.length} sites for client ${client.clientId}`
     );
 
-    if (sitesData.length === 0) {
-        sendToClient(olm.olmId, {
-            type: "olm/register/no-sites",
-            data: {}
-        });
-    }
-
     // Process each site
     for (const { sites: site } of sitesData) {
         if (!site.exitNodeId) {
@@ -462,6 +463,8 @@ async function getOrCreateOrgClient(
             });
         }
 
+        await rebuildClientAssociationsFromClient(newClient, trx); // TODO: this will try to messages to the olm which has not connected yet - is that a problem?
+
         client = newClient;
     }
 

server/routers/siteResource/addClientToSiteResource.ts

@@ -8,7 +8,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const addClientToSiteResourceBodySchema = z
     .object({
@@ -136,7 +136,7 @@ export async function addClientToSiteResource(
                 siteResourceId
             });
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/addRoleToSiteResource.ts

@@ -9,7 +9,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const addRoleToSiteResourceBodySchema = z
     .object({
@@ -146,7 +146,7 @@ export async function addRoleToSiteResource(
                 siteResourceId
             });
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/addUserToSiteResource.ts

@@ -9,7 +9,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const addUserToSiteResourceBodySchema = z
     .object({
@@ -115,7 +115,7 @@ export async function addUserToSiteResource(
                 siteResourceId
             });
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/createSiteResource.ts

@@ -17,7 +17,7 @@ import { fromError } from "zod-validation-error";
 import logger from "@server/logger";
 import { OpenAPITags, registry } from "@server/openApi";
 import { getUniqueSiteResourceName } from "@server/db/names";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 import { getNextAvailableAliasAddress } from "@server/lib/ip";
 
 const createSiteResourceParamsSchema = z.strictObject({
@@ -278,7 +278,7 @@ export async function createSiteResource(
                 );
             }
 
-            await rebuildClientAssociations(newSiteResource, trx); // we need to call this because we added to the admin role
+            await rebuildClientAssociationsFromSiteResource(newSiteResource, trx); // we need to call this because we added to the admin role
         });
 
         if (!newSiteResource) {

server/routers/siteResource/deleteSiteResource.ts

@@ -9,7 +9,7 @@ import { eq, and } from "drizzle-orm";
 import { fromError } from "zod-validation-error";
 import logger from "@server/logger";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const deleteSiteResourceParamsSchema = z.strictObject({
     siteResourceId: z.string().transform(Number).pipe(z.int().positive()),
@@ -106,7 +106,7 @@ export async function deleteSiteResource(
                 );
             }
 
-            await rebuildClientAssociations(removedSiteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(removedSiteResource, trx);
         });
 
         logger.info(

server/routers/siteResource/removeClientFromSiteResource.ts

@@ -8,7 +8,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const removeClientFromSiteResourceBodySchema = z
     .object({
@@ -142,7 +142,7 @@ export async function removeClientFromSiteResource(
                     )
                 );
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/removeRoleFromSiteResource.ts

@@ -9,7 +9,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const removeRoleFromSiteResourceBodySchema = z
     .object({
@@ -151,7 +151,7 @@ export async function removeRoleFromSiteResource(
                     )
                 );
 
-             await rebuildClientAssociations(siteResource, trx);
+             await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/removeUserFromSiteResource.ts

@@ -9,7 +9,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const removeUserFromSiteResourceBodySchema = z
     .object({
@@ -121,7 +121,7 @@ export async function removeUserFromSiteResource(
                     )
                 );
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/setSiteResourceClients.ts

@@ -8,7 +8,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, inArray } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const setSiteResourceClientsBodySchema = z
     .object({
@@ -124,7 +124,7 @@ export async function setSiteResourceClients(
                     .values(clientIds.map((clientId) => ({ clientId, siteResourceId })));
             }
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/setSiteResourceRoles.ts

@@ -9,7 +9,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq, and, ne, inArray } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const setSiteResourceRolesBodySchema = z
     .object({
@@ -147,7 +147,7 @@ export async function setSiteResourceRoles(
                     .values(roleIds.map((roleId) => ({ roleId, siteResourceId })));
             }
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/setSiteResourceUsers.ts

@@ -9,7 +9,7 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { eq } from "drizzle-orm";
 import { OpenAPITags, registry } from "@server/openApi";
-import { rebuildClientAssociations } from "@server/lib/rebuildClientAssociations";
+import { rebuildClientAssociationsFromSiteResource } from "@server/lib/rebuildClientAssociations";
 
 const setSiteResourceUsersBodySchema = z
     .object({
@@ -102,7 +102,7 @@ export async function setSiteResourceUsers(
                     .values(userIds.map((userId) => ({ userId, siteResourceId })));
             }
 
-            await rebuildClientAssociations(siteResource, trx);
+            await rebuildClientAssociationsFromSiteResource(siteResource, trx);
         });
 
         return response(res, {

server/routers/siteResource/updateSiteResource.ts

@@ -25,7 +25,7 @@ import {
 } from "@server/lib/ip";
 import {
     getClientSiteResourceAccess,
-    rebuildClientAssociations
+    rebuildClientAssociationsFromSiteResource
 } from "@server/lib/rebuildClientAssociations";
 
 const updateSiteResourceParamsSchema = z.strictObject({
@@ -224,7 +224,7 @@ export async function updateSiteResource(
                     );
             }
 
-            const { mergedAllClients } = await rebuildClientAssociations(
+            const { mergedAllClients } = await rebuildClientAssociationsFromSiteResource(
                 existingSiteResource, // we want to rebuild based on the existing resource then we will apply the change to the destination below
                 trx
             );

server/routers/user/addUserRole.ts

@@ -1,6 +1,6 @@
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
-import { db } from "@server/db";
+import { clients, db, UserOrg } from "@server/db";
 import { userOrgs, roles } from "@server/db";
 import { eq, and } from "drizzle-orm";
 import response from "@server/lib/response";
@@ -10,11 +10,12 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import stoi from "@server/lib/stoi";
 import { OpenAPITags, registry } from "@server/openApi";
+import { rebuildClientAssociationsFromClient } from "@server/lib/rebuildClientAssociations";
 
 const addUserRoleParamsSchema = z.strictObject({
-        userId: z.string(),
-        roleId: z.string().transform(stoi).pipe(z.number())
-    });
+    userId: z.string(),
+    roleId: z.string().transform(stoi).pipe(z.number())
+});
 
 export type AddUserRoleResponse = z.infer<typeof addUserRoleParamsSchema>;
 
@@ -72,7 +73,9 @@ export async function addUserRole(
         const existingUser = await db
             .select()
             .from(userOrgs)
-            .where(and(eq(userOrgs.userId, userId), eq(userOrgs.orgId, role.orgId)))
+            .where(
+                and(eq(userOrgs.userId, userId), eq(userOrgs.orgId, role.orgId))
+            )
             .limit(1);
 
         if (existingUser.length === 0) {
@@ -108,14 +111,39 @@ export async function addUserRole(
             );
         }
 
-        const newUserRole = await db
-            .update(userOrgs)
-            .set({ roleId })
-            .where(and(eq(userOrgs.userId, userId), eq(userOrgs.orgId, role.orgId)))
-            .returning();
+        let newUserRole: UserOrg | null = null;
+        await db.transaction(async (trx) => {
+            [newUserRole] = await trx
+                .update(userOrgs)
+                .set({ roleId })
+                .where(
+                    and(
+                        eq(userOrgs.userId, userId),
+                        eq(userOrgs.orgId, role.orgId)
+                    )
+                )
+                .returning();
+
+            // get the client associated with this user in this org
+            const [orgClient] = await trx
+                .select()
+                .from(clients)
+                .where(
+                    and(
+                        eq(clients.userId, userId),
+                        eq(clients.orgId, role.orgId)
+                    )
+                )
+                .limit(1);
+
+            if (orgClient) {
+                // we just changed the user's role, so we need to rebuild client associations and what they have access to
+                await rebuildClientAssociationsFromClient(orgClient, trx);
+            }
+        });
 
         return response(res, {
-            data: newUserRole[0],
+            data: newUserRole,
             success: true,
             error: false,
             message: "Role added to user successfully",

server/routers/user/adminRemoveUser.ts

@@ -8,10 +8,11 @@ import HttpCode from "@server/types/HttpCode";
 import createHttpError from "http-errors";
 import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
+import { calculateUserClientsForOrgs } from "@server/lib/calculateUserClientsForOrgs";
 
 const removeUserSchema = z.strictObject({
-        userId: z.string()
-    });
+    userId: z.string()
+});
 
 export async function adminRemoveUser(
     req: Request,
@@ -50,7 +51,11 @@ export async function adminRemoveUser(
             );
         }
 
-        await db.delete(users).where(eq(users.userId, userId));
+        await db.transaction(async (trx) => {
+            await trx.delete(users).where(eq(users.userId, userId));
+
+            await calculateUserClientsForOrgs(userId, trx);
+        });
 
         return response(res, {
             data: null,