Update rekey endpoints and pages

server/cleanup.ts

@@ -1,4 +1,4 @@
-import { cleanup as wsCleanup } from "@server/routers/ws";
+import { cleanup as wsCleanup } from "#dynamic/routers/ws";
 
 async function cleanup() {
     await wsCleanup();

server/private/routers/external.ts

@@ -30,7 +30,7 @@ import {
     verifyUserHasAction,
     verifyUserIsServerAdmin,
     verifySiteAccess,
-    verifyClientAccess,
+    verifyClientAccess
 } from "@server/middlewares";
 import { ActionsEnum } from "@server/auth/actions";
 import {
@@ -409,6 +409,8 @@ authenticated.get(
 
 authenticated.post(
     "/re-key/:clientId/regenerate-client-secret",
+    verifyValidLicense,
+    verifyValidSubscription,
     verifyClientAccess,
     verifyUserHasAction(ActionsEnum.reGenerateSecret),
     reKey.reGenerateClientSecret
@@ -416,15 +418,18 @@ authenticated.post(
 
 authenticated.post(
     "/re-key/:siteId/regenerate-site-secret",
+    verifyValidLicense,
+    verifyValidSubscription,
     verifySiteAccess,
     verifyUserHasAction(ActionsEnum.reGenerateSecret),
     reKey.reGenerateSiteSecret
 );
 
 authenticated.put(
-    "/re-key/:orgId/reGenerate-remote-exit-node-secret",
+    "/re-key/:orgId/regenerate-remote-exit-node-secret",
     verifyValidLicense,
+    verifyValidSubscription,
     verifyOrgAccess,
-    verifyUserHasAction(ActionsEnum.updateRemoteExitNode),
+    verifyUserHasAction(ActionsEnum.reGenerateSecret),
     reKey.reGenerateExitNodeSecret
 );

server/private/routers/re-key/reGenerateClientSecret.ts

@@ -13,7 +13,7 @@
 
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
-import { db, olms, } from "@server/db";
+import { db, olms } from "@server/db";
 import { clients } from "@server/db";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
@@ -23,16 +23,16 @@ import { eq, and } from "drizzle-orm";
 import { fromError } from "zod-validation-error";
 import { OpenAPITags, registry } from "@server/openApi";
 import { hashPassword } from "@server/auth/password";
+import { disconnectClient, sendToClient } from "#dynamic/routers/ws";
 
 const reGenerateSecretParamsSchema = z.strictObject({
-        clientId: z.string().transform(Number).pipe(z.int().positive())
+    clientId: z.string().transform(Number).pipe(z.int().positive())
-    });
+});
 
 const reGenerateSecretBodySchema = z.strictObject({
-        olmId: z.string().min(1).optional(),
+    // olmId: z.string().min(1).optional(),
-        secret: z.string().min(1).optional(),
+    secret: z.string().min(1)
-
+});
-    });
 
 export type ReGenerateSecretBody = z.infer<typeof reGenerateSecretBodySchema>;
 
@@ -54,7 +54,6 @@ registry.registerPath({
     responses: {}
 });
 
-
 export async function reGenerateClientSecret(
     req: Request,
     res: Response,
@@ -71,7 +70,7 @@ export async function reGenerateClientSecret(
             );
         }
 
-        const { olmId, secret } = parsedBody.data;
+        const { secret } = parsedBody.data;
 
         const parsedParams = reGenerateSecretParamsSchema.safeParse(req.params);
         if (!parsedParams.success) {
@@ -85,11 +84,7 @@ export async function reGenerateClientSecret(
 
         const { clientId } = parsedParams.data;
 
-        let secretHash = undefined;
+        const secretHash = await hashPassword(secret);
-        if (secret) {
-            secretHash = await hashPassword(secret);
-        }
-
 
         // Fetch the client to make sure it exists and the user has access to it
         const [client] = await db
@@ -107,24 +102,51 @@ export async function reGenerateClientSecret(
             );
         }
 
-        const [existingOlm] = await db
+        const existingOlms = await db
             .select()
             .from(olms)
-            .where(eq(olms.clientId, clientId))
+            .where(eq(olms.clientId, clientId));
-            .limit(1);
 
-        if (existingOlm && olmId && secretHash) {
+        if (existingOlms.length === 0) {
-            await db
+            return next(
-                .update(olms)
+                createHttpError(
-                .set({
+                    HttpCode.NOT_FOUND,
-                    olmId,
+                    `No OLM found for client ID ${clientId}`
-                    secretHash
+                )
-                })
+            );
-                .where(eq(olms.clientId, clientId));
         }
 
+        if (existingOlms.length > 1) {
+            return next(
+                createHttpError(
+                    HttpCode.INTERNAL_SERVER_ERROR,
+                    `Multiple OLM entries found for client ID ${clientId}`
+                )
+            );
+        }
+
+        await db
+            .update(olms)
+            .set({
+                secretHash
+            })
+            .where(eq(olms.olmId, existingOlms[0].olmId));
+
+        const payload = {
+            type: `olm/terminate`,
+            data: {}
+        };
+        // Don't await this to prevent blocking the response
+        sendToClient(existingOlms[0].olmId, payload).catch((error) => {
+            logger.error("Failed to send termination message to olm:", error);
+        });
+
+        disconnectClient(existingOlms[0].olmId).catch((error) => {
+            logger.error("Failed to disconnect olm after re-key:", error);
+        });
+
         return response(res, {
-            data: existingOlm,
+            data: existingOlms,
             success: true,
             error: false,
             message: "Credentials regenerated successfully",

server/private/routers/re-key/reGenerateExitNodeSecret.ts

@@ -24,16 +24,16 @@ import logger from "@server/logger";
 import { and, eq } from "drizzle-orm";
 import { UpdateRemoteExitNodeResponse } from "@server/routers/remoteExitNode/types";
 import { OpenAPITags, registry } from "@server/openApi";
+import { disconnectClient } from "@server/routers/ws";
 
 export const paramsSchema = z.object({
     orgId: z.string()
 });
 
 const bodySchema = z.strictObject({
-        remoteExitNodeId: z.string().length(15),
+    remoteExitNodeId: z.string().length(15),
-        secret: z.string().length(48)
+    secret: z.string().length(48)
-    });
+});
-
 
 registry.registerPath({
     method: "post",
@@ -81,12 +81,6 @@ export async function reGenerateExitNodeSecret(
 
         const { remoteExitNodeId, secret } = parsedBody.data;
 
-        if (req.user && !req.userOrgRoleId) {
-            return next(
-                createHttpError(HttpCode.FORBIDDEN, "User does not have a role")
-            );
-        }
-
         const [existingRemoteExitNode] = await db
             .select()
             .from(remoteExitNodes)
@@ -94,7 +88,10 @@ export async function reGenerateExitNodeSecret(
 
         if (!existingRemoteExitNode) {
             return next(
-                createHttpError(HttpCode.NOT_FOUND, "Remote Exit Node does not exist")
+                createHttpError(
+                    HttpCode.NOT_FOUND,
+                    "Remote Exit Node does not exist"
+                )
             );
         }
 
@@ -105,15 +102,21 @@ export async function reGenerateExitNodeSecret(
             .set({ secretHash })
             .where(eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId));
 
+        disconnectClient(existingRemoteExitNode.remoteExitNodeId).catch(
+            (error) => {
+                logger.error("Failed to disconnect newt after re-key:", error);
+            }
+        );
+
         return response<UpdateRemoteExitNodeResponse>(res, {
             data: {
                 remoteExitNodeId,
-                secret,
+                secret
             },
             success: true,
             error: false,
             message: "Remote Exit Node secret updated successfully",
-            status: HttpCode.OK,
+            status: HttpCode.OK
         });
     } catch (e) {
         logger.error("Failed to update remoteExitNode", e);

server/private/routers/re-key/reGenerateSiteSecret.ts

@@ -13,7 +13,7 @@
 
 import { Request, Response, NextFunction } from "express";
 import { z } from "zod";
-import { db, newts, sites } from "@server/db";
+import { db, Newt, newts, sites } from "@server/db";
 import { eq } from "drizzle-orm";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
@@ -22,38 +22,37 @@ import logger from "@server/logger";
 import { fromError } from "zod-validation-error";
 import { OpenAPITags, registry } from "@server/openApi";
 import { hashPassword } from "@server/auth/password";
-import { addPeer } from "@server/routers/gerbil/peers";
+import { addPeer, deletePeer } from "@server/routers/gerbil/peers";
-
+import { getAllowedIps } from "@server/routers/target/helpers";
+import { disconnectClient, sendToClient } from "#dynamic/routers/ws";
 
 const updateSiteParamsSchema = z.strictObject({
-        siteId: z.string().transform(Number).pipe(z.int().positive())
+    siteId: z.string().transform(Number).pipe(z.int().positive())
-    });
+});
 
 const updateSiteBodySchema = z.strictObject({
-        type: z.enum(["newt", "wireguard"]),
+    type: z.enum(["newt", "wireguard"]),
-        newtId: z.string().min(1).max(255).optional(),
+    secret: z.string().min(1).max(255).optional(),
-        newtSecret: z.string().min(1).max(255).optional(),
+    pubKey: z.string().optional()
-        exitNodeId: z.int().positive().optional(),
+});
-        pubKey: z.string().optional(),
-        subnet: z.string().optional(),
-    });
 
 registry.registerPath({
     method: "post",
     path: "/re-key/{siteId}/regenerate-site-secret",
-    description: "Regenerate a site's Newt or WireGuard credentials by its site ID.",
+    description:
+        "Regenerate a site's Newt or WireGuard credentials by its site ID.",
     tags: [OpenAPITags.Site],
     request: {
         params: updateSiteParamsSchema,
         body: {
             content: {
                 "application/json": {
-                    schema: updateSiteBodySchema,
+                    schema: updateSiteBodySchema
-                },
+                }
-            },
+            }
-        },
+        }
     },
-    responses: {},
+    responses: {}
 });
 
 export async function reGenerateSiteSecret(
@@ -65,74 +64,141 @@ export async function reGenerateSiteSecret(
         const parsedParams = updateSiteParamsSchema.safeParse(req.params);
         if (!parsedParams.success) {
             return next(
-                createHttpError(HttpCode.BAD_REQUEST, fromError(parsedParams.error).toString())
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
             );
         }
 
         const parsedBody = updateSiteBodySchema.safeParse(req.body);
         if (!parsedBody.success) {
             return next(
-                createHttpError(HttpCode.BAD_REQUEST, fromError(parsedBody.error).toString())
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedBody.error).toString()
+                )
             );
         }
 
         const { siteId } = parsedParams.data;
-        const { type, exitNodeId, pubKey, subnet, newtId, newtSecret } = parsedBody.data;
+        const { type, pubKey, secret } = parsedBody.data;
-
-        let updatedSite = undefined;
 
+        let existingNewt: Newt | null = null;
         if (type === "newt") {
-            if (!newtSecret) {
+            if (!secret) {
-                return next(
-                    createHttpError(HttpCode.BAD_REQUEST, "newtSecret is required for newt sites")
-                );
-            }
-
-            const secretHash = await hashPassword(newtSecret);
-
-            updatedSite = await db
-                .update(newts)
-                .set({
-                    newtId,
-                    secretHash,
-                })
-                .where(eq(newts.siteId, siteId))
-                .returning();
-
-            logger.info(`Regenerated Newt credentials for site ${siteId}`);
-
-        } else if (type === "wireguard") {
-            if (!pubKey) {
-                return next(
-                    createHttpError(HttpCode.BAD_REQUEST, "Public key is required for wireguard sites")
-                );
-            }
-
-            if (!exitNodeId) {
                 return next(
                     createHttpError(
                         HttpCode.BAD_REQUEST,
-                        "Exit node ID is required for wireguard sites"
+                        "newtSecret is required for newt sites"
+                    )
+                );
+            }
+
+            const secretHash = await hashPassword(secret);
+
+            // get the newt to verify it exists
+             const existingNewts = await db
+                .select()
+                .from(newts)
+                .where(eq(newts.siteId, siteId));
+
+            if (existingNewts.length === 0) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        `No Newt found for site ID ${siteId}`
+                    )
+                );
+            }
+
+            if (existingNewts.length > 1) {
+                return next(
+                    createHttpError(
+                        HttpCode.INTERNAL_SERVER_ERROR,
+                        `Multiple Newts found for site ID ${siteId}`
+                    )
+                );
+            }
+
+            existingNewt = existingNewts[0];
+
+            // update the secret on the existing newt
+            await db
+                .update(newts)
+                .set({
+                    secretHash
+                })
+                .where(eq(newts.newtId, existingNewts[0].newtId));
+
+            const payload = {
+                type: `newt/wg/terminate`,
+                data: {}
+            };
+            // Don't await this to prevent blocking the response
+            sendToClient(existingNewts[0].newtId, payload).catch((error) => {
+                logger.error(
+                    "Failed to send termination message to newt:",
+                    error
+                );
+            });
+
+            disconnectClient(existingNewts[0].newtId).catch((error) => {
+                logger.error("Failed to disconnect newt after re-key:", error);
+            });
+
+            logger.info(`Regenerated Newt credentials for site ${siteId}`);
+        } else if (type === "wireguard") {
+            if (!pubKey) {
+                return next(
+                    createHttpError(
+                        HttpCode.BAD_REQUEST,
+                        "Public key is required for wireguard sites"
                     )
                 );
             }
 
             try {
-                updatedSite = await db.transaction(async (tx) => {
+                const [site] = await db
-                    await addPeer(exitNodeId, {
+                    .select()
+                    .from(sites)
+                    .where(eq(sites.siteId, siteId))
+                    .limit(1);
+
+                if (!site) {
+                    return next(
+                        createHttpError(
+                            HttpCode.NOT_FOUND,
+                            `Site with ID ${siteId} not found`
+                        )
+                    );
+                }
+
+                await db
+                    .update(sites)
+                    .set({ pubKey })
+                    .where(eq(sites.siteId, siteId));
+
+                if (!site) {
+                    return next(
+                        createHttpError(
+                            HttpCode.NOT_FOUND,
+                            `Site with ID ${siteId} not found`
+                        )
+                    );
+                }
+
+                if (site.exitNodeId && site.subnet) {
+                    await deletePeer(site.exitNodeId, site.pubKey!); // the old pubkey
+                    await addPeer(site.exitNodeId, {
                         publicKey: pubKey,
-                        allowedIps: subnet ? [subnet] : [],
+                        allowedIps: await getAllowedIps(site.siteId)
                     });
-                    const result = await tx
+                }
-                        .update(sites)
-                        .set({ pubKey })
-                        .where(eq(sites.siteId, siteId))
-                        .returning();
 
-                    return result;
+                logger.info(
-                });
+                    `Regenerated WireGuard credentials for site ${siteId}`
-
+                );
-                logger.info(`Regenerated WireGuard credentials for site ${siteId}`);
             } catch (err) {
                 logger.error(
                     `Transaction failed while regenerating WireGuard secret for site ${siteId}`,
@@ -148,17 +214,19 @@ export async function reGenerateSiteSecret(
         }
 
         return response(res, {
-            data: updatedSite,
+            data: existingNewt,
             success: true,
             error: false,
             message: "Credentials regenerated successfully",
-            status: HttpCode.OK,
+            status: HttpCode.OK
         });
-
     } catch (error) {
         logger.error("Unexpected error in reGenerateSiteSecret", error);
         return next(
-            createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An unexpected error occurred")
+            createHttpError(
+                HttpCode.INTERNAL_SERVER_ERROR,
+                "An unexpected error occurred"
+            )
         );
     }
 }

server/routers/olm/getOlmToken.ts

@@ -104,43 +104,13 @@ export async function getOlmToken(
         const resToken = generateSessionToken();
         await createOlmSession(resToken, existingOlm.olmId);
 
-        let orgIdToUse = orgId;
         let clientIdToUse;
-        if (!orgIdToUse) {
+        if (orgId) {
-            if (!existingOlm.clientId) {
-                return next(
-                    createHttpError(
-                        HttpCode.BAD_REQUEST,
-                        "Olm is not associated with a client, orgId is required"
-                    )
-                );
-            }
-
-            const [client] = await db
-                .select()
-                .from(clients)
-                .where(eq(clients.clientId, existingOlm.clientId))
-                .limit(1);
-
-            if (!client) {
-                return next(
-                    createHttpError(
-                        HttpCode.BAD_REQUEST,
-                        "Olm's associated client not found, orgId is required"
-                    )
-                );
-            }
-
-            orgIdToUse = client.orgId;
-            clientIdToUse = client.clientId;
-        } else {
             // we did provide the org
             const [client] = await db
                 .select()
                 .from(clients)
-                .where(
+                .where(and(eq(clients.orgId, orgId), eq(clients.olmId, olmId))) // we want to lock on to the client with this olmId otherwise it can get assigned to a random one
-                    and(eq(clients.orgId, orgIdToUse), eq(clients.olmId, olmId))
-                ) // we want to lock on to the client with this olmId otherwise it can get assigned to a random one
                 .limit(1);
 
             if (!client) {
@@ -167,6 +137,32 @@ export async function getOlmToken(
                     .where(eq(olms.olmId, existingOlm.olmId));
             }
 
+            clientIdToUse = client.clientId;
+        } else {
+            if (!existingOlm.clientId) {
+                return next(
+                    createHttpError(
+                        HttpCode.BAD_REQUEST,
+                        "Olm is not associated with a client, orgId is required"
+                    )
+                );
+            }
+
+            const [client] = await db
+                .select()
+                .from(clients)
+                .where(eq(clients.clientId, existingOlm.clientId))
+                .limit(1);
+
+            if (!client) {
+                return next(
+                    createHttpError(
+                        HttpCode.BAD_REQUEST,
+                        "Olm's associated client not found, orgId is required"
+                    )
+                );
+            }
+
             clientIdToUse = client.clientId;
         }
 

src/app/[orgId]/settings/(private)/remote-exit-nodes/[remoteExitNodeId]/credentials/page.tsx

@@ -64,7 +64,7 @@ export default function CredentialsPage() {
         setCredentials(data);
 
         await api.put<AxiosResponse<QuickStartRemoteExitNodeResponse>>(
-            `/re-key/${orgId}/reGenerate-remote-exit-node-secret`,
+            `/re-key/${orgId}/regenerate-remote-exit-node-secret`,
             {
                 remoteExitNodeId: remoteExitNode.remoteExitNodeId,
                 secret: data.secret

src/app/[orgId]/settings/clients/machine/[clientId]/credentials/page.tsx

@@ -60,7 +60,6 @@ export default function CredentialsPage() {
             await api.post(
                 `/re-key/${client?.clientId}/regenerate-client-secret`,
                 {
-                    olmId: data.olmId,
                     secret: data.olmSecret
                 }
             );

src/app/[orgId]/settings/sites/[niceId]/credentials/page.tsx

@@ -103,8 +103,6 @@ PersistentKeepalive = 5`;
 
             await api.post(`/re-key/${site?.siteId}/regenerate-site-secret`, {
                 type: "wireguard",
-                subnet: res.data.data.subnet,
-                exitNodeId: res.data.data.exitNodeId,
                 pubKey: generatedPublicKey
             });
         }
@@ -119,8 +117,7 @@ PersistentKeepalive = 5`;
                     `/re-key/${site?.siteId}/regenerate-site-secret`,
                     {
                         type: "newt",
-                        newtId: data.newtId,
+                        secret: data.newtSecret
-                        newtSecret: data.newtSecret
                     }
                 );
             }