pangolin/server/lib/deleteOrg.ts

import {
    clients,
    clientSiteResourcesAssociationsCache,
    clientSitesAssociationsCache,
    db,
    domains,
    exitNodeOrgs,
    exitNodes,
    olms,
    orgDomains,
    orgs,
    remoteExitNodes,
    resources,
    sites,
    userOrgs
} from "@server/db";
import { newts, newtSessions } from "@server/db";
import { eq, and, inArray, sql, count, countDistinct } from "drizzle-orm";
import HttpCode from "@server/types/HttpCode";
import createHttpError from "http-errors";
import logger from "@server/logger";
import { sendToClient } from "#dynamic/routers/ws";
import { deletePeer } from "@server/routers/gerbil/peers";
import { OlmErrorCodes } from "@server/routers/olm/error";
import { sendTerminateClient } from "@server/routers/client/terminate";
import { cleanupSiteAssociations } from "@server/lib/rebuildClientAssociations";
import { usageService } from "./billing/usageService";
import { LimitId } from "./billing";

export type DeleteOrgByIdResult = {
    deletedNewtIds: string[];
    olmsToTerminate: string[];
};

/**
 * Deletes one organization and its related data. Returns ids for termination
 * messages; caller should call sendTerminationMessages with the result.
 * Throws if org not found.
 */
export async function deleteOrgById(
    orgId: string
): Promise<DeleteOrgByIdResult> {
    const [org] = await db
        .select()
        .from(orgs)
        .where(eq(orgs.orgId, orgId))
        .limit(1);

    if (!org) {
        throw createHttpError(
            HttpCode.NOT_FOUND,
            `Organization with ID ${orgId} not found`
        );
    }

    const orgSites = await db
        .select()
        .from(sites)
        .where(eq(sites.orgId, orgId));

    const orgClients = await db
        .select()
        .from(clients)
        .where(eq(clients.orgId, orgId));

    const deletedNewtIds: string[] = [];
    const olmsToTerminate: string[] = [];

    let domainCount: number | null = null;
    let siteCount: number | null = null;
    let userCount: number | null = null;
    let remoteExitNodeCount: number | null = null;
    // billingOrgId has no FK, so the billing org may already be gone (e.g.
    // a partially failed account deletion). Writing usage for it would
    // violate the usage FK and abort the transaction.
    let billingOrgExists = false;

    await db.transaction(async (trx) => {
        for (const site of orgSites) {
            if (site.type == "wireguard") {
                if (site.pubKey) {
                    await deletePeer(site.exitNodeId!, site.pubKey);
                }
            } else if (site.type == "newt") {
                // Same per-site cleanup deleteSite.ts performs, so that
                // removing an org does not leave other orgs' clients
                // holding stale associations for these sites.
                await cleanupSiteAssociations(site, trx);
                const [deletedNewt] = await trx
                    .delete(newts)
                    .where(eq(newts.siteId, site.siteId))
                    .returning();
                if (deletedNewt) {
                    deletedNewtIds.push(deletedNewt.newtId);
                    await trx
                        .delete(newtSessions)
                        .where(eq(newtSessions.newtId, deletedNewt.newtId));
                }
            }
            logger.info(`Deleting site ${site.siteId}`);
            await trx.delete(sites).where(eq(sites.siteId, site.siteId));
        }
        for (const client of orgClients) {
            if (client.exitNodeId && client.pubKey) {
                await deletePeer(client.exitNodeId, client.pubKey);
            }
            const [olm] = await trx
                .select()
                .from(olms)
                .where(eq(olms.clientId, client.clientId))
                .limit(1);
            if (olm) {
                olmsToTerminate.push(olm.olmId);
            }
            logger.info(`Deleting client ${client.clientId}`);
            await trx
                .delete(clients)
                .where(eq(clients.clientId, client.clientId));
            await trx
                .delete(clientSiteResourcesAssociationsCache)
                .where(
                    eq(
                        clientSiteResourcesAssociationsCache.clientId,
                        client.clientId
                    )
                );
            await trx
                .delete(clientSitesAssociationsCache)
                .where(
                    eq(clientSitesAssociationsCache.clientId, client.clientId)
                );
        }

        await trx.delete(resources).where(eq(resources.orgId, orgId));

        const allOrgDomains = await trx
            .select()
            .from(orgDomains)
            .innerJoin(domains, eq(orgDomains.domainId, domains.domainId))
            .where(
                and(
                    eq(orgDomains.orgId, orgId),
                    eq(domains.configManaged, false)
                )
            );
        logger.info(`Found ${allOrgDomains.length} domains to delete`);
        const domainIdsToDelete: string[] = [];
        for (const orgDomain of allOrgDomains) {
            const domainId = orgDomain.domains.domainId;
            const [orgCount] = await trx
                .select({ count: count() })
                .from(orgDomains)
                .where(eq(orgDomains.domainId, domainId));
            logger.info(
                `Found ${orgCount.count} orgs using domain ${domainId}`
            );
            if (orgCount.count === 1) {
                domainIdsToDelete.push(domainId);
            }
        }
        logger.info(`Found ${domainIdsToDelete.length} domains to delete`);
        if (domainIdsToDelete.length > 0) {
            await trx
                .delete(domains)
                .where(inArray(domains.domainId, domainIdsToDelete));
        }

        const [billingOrg] = org.billingOrgId
            ? await trx
                  .select({ orgId: orgs.orgId })
                  .from(orgs)
                  .where(eq(orgs.orgId, org.billingOrgId))
                  .limit(1)
            : [org];
        billingOrgExists = !!billingOrg;
        if (billingOrgExists) {
            await usageService.add(orgId, LimitId.ORGANIZATIONS, -1, trx); // here we are decreasing the org count BEFORE deleting the org because we need to still be able to get the org to get the billing org inside of here
        }

        await trx.delete(orgs).where(eq(orgs.orgId, orgId));

        if (org.billingOrgId) {
            const billingOrgs = await trx
                .select()
                .from(orgs)
                .where(eq(orgs.billingOrgId, org.billingOrgId));

            if (billingOrgs.length > 0) {
                const billingOrgIds = billingOrgs.map((org) => org.orgId);

                const [domainCountRes] = await trx
                    .select({ count: count() })
                    .from(orgDomains)
                    .where(inArray(orgDomains.orgId, billingOrgIds));

                domainCount = domainCountRes.count;

                const [siteCountRes] = await trx
                    .select({ count: count() })
                    .from(sites)
                    .where(inArray(sites.orgId, billingOrgIds));

                siteCount = siteCountRes.count;

                const [userCountRes] = await trx
                    .select({ count: countDistinct(userOrgs.userId) })
                    .from(userOrgs)
                    .where(inArray(userOrgs.orgId, billingOrgIds));

                userCount = userCountRes.count;

                const [remoteExitNodeCountRes] = await trx
                    .select({ count: countDistinct(exitNodeOrgs.exitNodeId) })
                    .from(exitNodeOrgs)
                    .where(inArray(exitNodeOrgs.orgId, billingOrgIds));

                remoteExitNodeCount = remoteExitNodeCountRes.count;
            }
        }
    });

    if (
        org.billingOrgId &&
        org.billingOrgId !== orgId &&
        billingOrgExists
    ) {
        usageService.updateCount(
            org.billingOrgId,
            LimitId.DOMAINS,
            domainCount ?? 0
        );
        usageService.updateCount(
            org.billingOrgId,
            LimitId.SITES,
            siteCount ?? 0
        );
        usageService.updateCount(
            org.billingOrgId,
            LimitId.USERS,
            userCount ?? 0
        );
        usageService.updateCount(
            org.billingOrgId,
            LimitId.REMOTE_EXIT_NODES,
            remoteExitNodeCount ?? 0
        );
    }

    return { deletedNewtIds, olmsToTerminate };
}

export function sendTerminationMessages(result: DeleteOrgByIdResult): void {
    for (const newtId of result.deletedNewtIds) {
        sendToClient(newtId, { type: `newt/wg/terminate`, data: {} }).catch(
            (error) => {
                logger.error(
                    "Failed to send termination message to newt:",
                    error
                );
            }
        );
    }
    for (const olmId of result.olmsToTerminate) {
        sendTerminateClient(0, OlmErrorCodes.TERMINATED_REKEYED, olmId).catch(
            (error) => {
                logger.error(
                    "Failed to send termination message to olm:",
                    error
                );
            }
        );
    }
}