HiveGrid-Wireframes/hivegrid-cloud/apps/web/app/routeModules/admin/mfa-challenge.impl.ts

import { redirect, type ActionFunctionArgs, type LoaderFunctionArgs } from "react-router";
import { serverOnly$ } from "vite-env-only/macros";
import { loadRequiredWebConfig } from "@toolctrl/config";
import { getTranslations } from "~/locale/i18next.server";
import { ADMIN_MFA_DEFAULT_REDIRECT_PATH, ADMIN_MFA_ENROLLMENT_PATH } from "~/utils/auth/mfa/paths";
import { setMfaStepUpState, type MfaStepUpMethod } from "~/utils/auth/mfa/step-up.server";
import { verifyTotpCode } from "~/services/auth/totp.server";
import { checkMfaChallengeRateLimit, mfaChallengeRateLimitKey, recordFailedMfaChallengeAttempt, resetMfaChallengeRateLimit } from "~/utils/auth/mfa/challenge-rate-limit.server";
import { absoluteAdminPath } from "~/utils/app/CrossHostUrls";
import { getHostRoutingContextFromRequest } from "~/utils/app/HostRouting.server";
import { safeRedirectPath } from "~/utils/app/RedirectSafety.server";
import { consumeUserMfaRecoveryCode, getUser, getUserTotpStatus } from "~/utils/db/users.db.server";
import { db } from "~/utils/db.server";
import { commitSession, getUserSession, validateCSRFToken } from "~/utils/session.server";

export type LoaderData = {
  redirectTo: string;
  totpEnabled: boolean;
  enrollmentPath: string;
};

export type ActionData = {
  error?: string;
  redirectTo: string;
};

function isAdminHostRequest(request: Request) {
  return getHostRoutingContextFromRequest(request) === "admin";
}

function getCanonicalAdminChallengeUrl(request: Request) {
  const url = new URL(request.url);
  return absoluteAdminPath(loadRequiredWebConfig().routing, `mfa-challenge${url.search}`);
}

function safeRedirectTo(value: FormDataEntryValue | string | null | undefined) {
  return safeRedirectPath(typeof value === "string" ? value : null, { fallbackPath: ADMIN_MFA_DEFAULT_REDIRECT_PATH });
}

function isSixDigitTotpCode(code: FormDataEntryValue | null): code is string {
  return typeof code === "string" && /^\d{6}$/.test(code);
}

function normalizeRecoveryCode(code: FormDataEntryValue | null) {
  if (typeof code !== "string") {
    return "";
  }
  return code.trim().toUpperCase();
}

async function getActiveTotpSecret(userId: string) {
  const user = await db.user.findUnique({
    where: { id: userId },
    select: { totpSecretEncrypted: true, totpEnabledAt: true },
  });

  if (!user?.totpSecretEncrypted || !user.totpEnabledAt) {
    return null;
  }

  return user.totpSecretEncrypted;
}

export const loader = serverOnly$(async ({ request }: LoaderFunctionArgs) => {
  if (!isAdminHostRequest(request)) {
    throw redirect(getCanonicalAdminChallengeUrl(request));
  }

  const session = await getUserSession(request);
  const userId = session.get("userId");
  if (typeof userId !== "string" || userId.length === 0) {
    throw redirect("/login");
  }

  const url = new URL(request.url);
  const { t } = await getTranslations(request);
  const redirectTo = safeRedirectTo(url.searchParams.get("redirectTo"));
  const user = await getUser(userId);
  if (!user?.admin) {
    throw Response.json({ error: t("admin.mfaChallenge.errors.adminAccessRequired"), redirectTo } satisfies ActionData, { status: 403 });
  }

  const totpStatus = await getUserTotpStatus(userId);
  if (!totpStatus?.enabled) {
    throw redirect(ADMIN_MFA_ENROLLMENT_PATH);
  }

  return Response.json({ redirectTo, totpEnabled: true, enrollmentPath: ADMIN_MFA_ENROLLMENT_PATH } satisfies LoaderData);
})!;

export const action = serverOnly$(async ({ request }: ActionFunctionArgs) => {
  const { t } = await getTranslations(request);
  if (!isAdminHostRequest(request)) {
    return Response.json(
      { error: t("admin.mfaChallenge.errors.adminHostRequired"), redirectTo: getCanonicalAdminChallengeUrl(request) } satisfies ActionData,
      { status: 403 }
    );
  }

  const session = await getUserSession(request);
  const userId = session.get("userId");
  if (typeof userId !== "string" || userId.length === 0) {
    throw redirect("/login");
  }

  await validateCSRFToken(request);
  const form = await request.formData();
  const redirectTo = safeRedirectTo(form.get("redirectTo") ?? new URL(request.url).searchParams.get("redirectTo"));
  const user = await getUser(userId);
  if (!user?.admin) {
    return Response.json({ error: t("admin.mfaChallenge.errors.adminAccessRequired"), redirectTo } satisfies ActionData, { status: 403 });
  }

  const totpStatus = await getUserTotpStatus(user.id);
  if (!totpStatus?.enabled) {
    throw redirect(ADMIN_MFA_ENROLLMENT_PATH);
  }

  const rateLimitKey = mfaChallengeRateLimitKey(request, user.id);
  const rateLimit = await checkMfaChallengeRateLimit(rateLimitKey);
  if (!rateLimit.allowed) {
    return Response.json({ error: t("admin.mfaChallenge.errors.tooManyAttempts"), redirectTo } satisfies ActionData, {
      status: 429,
      headers: { "Retry-After": String(rateLimit.retryAfterSeconds) },
    });
  }

  const code = form.get("code");
  let method: MfaStepUpMethod | null = null;

  if (isSixDigitTotpCode(code)) {
    const encryptedSecret = await getActiveTotpSecret(user.id);
    let totpMatches = false;
    try {
      totpMatches = Boolean(encryptedSecret && verifyTotpCode(encryptedSecret, code));
    } catch {
      totpMatches = false;
    }
    if (totpMatches) {
      method = "totp";
    }
  }

  if (!method) {
    const recoveryCode = normalizeRecoveryCode(form.get("recoveryCode"));
    if (recoveryCode.length > 0) {
      const result = await consumeUserMfaRecoveryCode(user.id, recoveryCode);
      if (result?.consumed) {
        method = "recovery-code";
      }
    }
  }

  if (!method) {
    await recordFailedMfaChallengeAttempt(rateLimitKey);
    return Response.json({ error: t("admin.mfaChallenge.errors.invalidCode"), redirectTo } satisfies ActionData, { status: 400 });
  }

  await resetMfaChallengeRateLimit(rateLimitKey);
  setMfaStepUpState(session, { userId: user.id, method, redirectTo });
  throw redirect(redirectTo, {
    headers: {
      "Set-Cookie": await commitSession(session),
    },
  });
})!;