pangolin/server/routers/idp/oidcAutoProvision.ts

import { Request, Response } from "express";
import { z } from "zod";
import {
    createSession,
    generateId,
    generateSessionToken,
    serializeSessionCookie
} from "@server/auth/sessions/app";
import logger from "@server/logger";
import db from "@server/db";
import {
    Idp,
    idpOrg,
    orgs,
    roles,
    User,
    userOrgs,
    users
} from "@server/db/schemas";
import { eq, and, inArray } from "drizzle-orm";
import jmespath from "jmespath";
import { UserType } from "@server/types/UserTypes";

const extractedRolesSchema = z.array(z.string()).or(z.string()).nullable();

export async function oidcAutoProvision({
    idp,
    claims,
    existingUser,
    userIdentifier,
    email,
    name,
    req,
    res
}: {
    idp: Idp;
    claims: any;
    existingUser?: User;
    userIdentifier: string;
    email?: string;
    name?: string;
    req: Request;
    res: Response;
}) {
    // Get user's roles of all orgs as stated in the ID token claims
    const allOrgs = await db.select().from(orgs);
    const userOrgInfo: { orgId: string; roleId: number }[] = [];

    for (const org of allOrgs) {
        const idpOrgs = await db
            .select()
            .from(idpOrg)
            .where(
                and(eq(idpOrg.idpId, idp.idpId), eq(idpOrg.orgId, org.orgId))
            );
        if (idpOrgs.length === 0) continue;
        const idpOrgRes = idpOrgs[0];

        const orgMapping = hydrateOrgMapping(
            idpOrgRes.orgMapping || idp.defaultOrgMapping,
            org.orgId
        );
        const roleMapping = idpOrgRes.roleMapping || idp.defaultRoleMapping;

        if (orgMapping) {
            const orgId = jmespath.search(claims, orgMapping);
            logger.debug("Extracted org ID", { orgId });
            if (orgId !== true && orgId !== org.orgId) {
                // user not allowed to access this org
                continue;
            }
        }

        if (roleMapping) {
            logger.info("claims", { claims });
            const extractedRoles = extractedRolesSchema.safeParse(
                jmespath.search(claims, roleMapping)
            );
            if (!extractedRoles.success) {
                logger.error("Error extracting roles", {
                    error: extractedRoles.error
                });
                continue;
            }
            const rd = extractedRoles.data;
            if (!rd) {
                continue;
            }
            const rolesFromToken = typeof rd === "string" ? [rd] : rd;
            logger.debug("Extracted roles", { rolesFromToken });
            if (rd.length === 0) {
                continue;
            }

            const rolesFromDb = await db
                .select()
                .from(roles)
                .where(
                    and(
                        eq(roles.orgId, org.orgId),
                        inArray(roles.name, rolesFromToken)
                    )
                );
            if (rolesFromDb.length === 0) {
                logger.error("Role(s) not found", { roles: rolesFromToken });
                continue;
            }
            if (rolesFromDb.length < rolesFromToken.length) {
                logger.warn("Role(s) not found", {
                    roles: rolesFromToken.filter(
                        (r) => !rolesFromDb.some((rdb) => rdb.name === r)
                    )
                });
            }

            rolesFromDb.forEach((r) => {
                userOrgInfo.push({
                    orgId: org.orgId,
                    roleId: r.roleId
                });
            });
        }
    }
    logger.debug("User org info", { userOrgInfo });

    let userId = existingUser?.userId;
    // sync the user with the orgs and roles
    await db.transaction(async (trx) => {
        if (!userId) {
            // create user if it does not exist
            userId = generateId(15);

            await trx.insert(users).values({
                userId,
                username: userIdentifier,
                email: email || null,
                name: name || null,
                type: UserType.OIDC,
                idpId: idp.idpId,
                emailVerified: true, // OIDC users are always verified
                dateCreated: new Date().toISOString()
            });
        } else {
            // update username/email
            await trx
                .update(users)
                .set({
                    username: userIdentifier,
                    email: email || null,
                    name: name || null
                })
                .where(eq(users.userId, userId));
        }

        // get all current user orgs/roles
        const currentUserOrgs = await trx
            .select()
            .from(userOrgs)
            .where(eq(userOrgs.userId, userId));

        // Delete orgs that are no longer valid
        const orgsToDelete = currentUserOrgs
            .filter(
                (currentOrg) =>
                    !userOrgInfo.some(
                        (newOrg) =>
                            newOrg.orgId === currentOrg.orgId &&
                            newOrg.roleId === currentOrg.roleId
                    )
            )
            .map((org) => org.orgId);

        if (orgsToDelete.length > 0) {
            await trx
                .delete(userOrgs)
                .where(
                    and(
                        eq(userOrgs.userId, userId!),
                        inArray(userOrgs.orgId, orgsToDelete)
                    )
                );
        }

        // Add new orgs that don't exist yet
        const orgsToAdd = userOrgInfo.filter(
            (newOrg) =>
                !currentUserOrgs.some(
                    (currentOrg) =>
                        currentOrg.orgId === newOrg.orgId &&
                        currentOrg.roleId === newOrg.roleId
                )
        );
        if (orgsToAdd.length > 0) {
            await trx.insert(userOrgs).values(
                orgsToAdd.map((org) => ({
                    userId: userId!,
                    orgId: org.orgId,
                    roleId: org.roleId
                }))
            );
        }
    });

    const token = generateSessionToken();
    const sess = await createSession(token, userId!);
    const isSecure = req.protocol === "https";
    const cookie = serializeSessionCookie(
        token,
        isSecure,
        new Date(sess.expiresAt)
    );

    res.appendHeader("Set-Cookie", cookie);
}

function hydrateOrgMapping(
    orgMapping: string | null,
    orgId: string
): string | null {
    if (!orgMapping) {
        return null;
    }
    return orgMapping.replaceAll("{{orgId}}", orgId);
}