app/src/__tests__/totp-backup.test.ts

import { describe, it, expect, beforeAll, afterAll } from "vitest";
import crypto from "crypto";
import bcrypt from "bcryptjs";
import * as OTPAuthLib from "otpauth";
import { buildApp, extractCookie } from "./helpers";
import prisma from "../config/database";
import { config } from "../config/env";
import { encrypt } from "../utils/encryption";

let app: Awaited<ReturnType<typeof buildApp>>;

// Fixture prefix so cleanup never touches real data.
const N = "totp_backup_test_";
const TEST_USERNAME = `${N}user`;

// Plaintext backup codes seeded for the fixture user (the enable route stores
// bcrypt hashes of 8-char uppercase hex codes — mirror that format).
const BACKUP_CODE_1 = "AAAA1111";
const BACKUP_CODE_2 = "BBBB2222";

let testUserId: number;

/** Pull the raw Set-Cookie header line for a given cookie name (with attrs). */
function rawSetCookie(res: { headers: Record<string, unknown> }, name: string) {
  const cookies = res.headers["set-cookie"];
  if (!cookies) return undefined;
  const arr = Array.isArray(cookies) ? cookies : [cookies];
  return arr.find((c) => typeof c === "string" && c.startsWith(`${name}=`)) as
    | string
    | undefined;
}

/** Create a pending TOTP login token the same way /login does (sha256 hash). */
async function issueLoginToken(userId: number): Promise<string> {
  const raw = crypto.randomBytes(32).toString("hex");
  const hash = crypto.createHash("sha256").update(raw).digest("hex");
  await prisma.totp_login_tokens.create({
    data: {
      user_id: userId,
      token_hash: hash,
      expires_at: new Date(Date.now() + 5 * 60_000),
    },
  });
  return raw;
}

// /totp/backup-verify carries its own per-IP rate limit
// (config.rateLimit.loginTotp, default 5/min) and this file legitimately makes
// more attempts than that — give each request a distinct source IP.
let ipCounter = 0;
function postBackupVerify(payload: Record<string, unknown>) {
  ipCounter += 1;
  return app.inject({
    method: "POST",
    url: "/api/admin/totp/backup-verify",
    payload,
    remoteAddress: `10.99.0.${ipCounter}`,
  });
}

async function fixtureUser() {
  const user = await prisma.users.findUniqueOrThrow({
    where: { id: testUserId },
  });
  return user;
}

beforeAll(async () => {
  app = await buildApp();

  // Clean up any leftover fixture from a previous failed run.
  const leftovers = await prisma.users.findMany({
    where: { username: { startsWith: N } },
    select: { id: true },
  });
  if (leftovers.length) {
    const ids = leftovers.map((u) => u.id);
    await prisma.refresh_tokens.deleteMany({ where: { user_id: { in: ids } } });
    await prisma.totp_login_tokens.deleteMany({
      where: { user_id: { in: ids } },
    });
    await prisma.users.deleteMany({ where: { username: { startsWith: N } } });
  }

  const role = await prisma.roles.findFirst();
  if (!role) throw new Error("Test setup: no roles found — seed the database");

  const hashedCodes = [
    await bcrypt.hash(BACKUP_CODE_1, config.security.bcryptCost),
    await bcrypt.hash(BACKUP_CODE_2, config.security.bcryptCost),
  ];

  const user = await prisma.users.create({
    data: {
      username: TEST_USERNAME,
      email: `${N}user@test.local`,
      password_hash: await bcrypt.hash(
        "irrelevant",
        config.security.bcryptCost,
      ),
      first_name: "TotpBackup",
      last_name: "Test",
      is_active: true,
      totp_enabled: true,
      totp_secret: encrypt(new OTPAuthLib.Secret().base32),
      totp_backup_codes: JSON.stringify(hashedCodes),
      role_id: role.id,
    },
  });
  testUserId = user.id;
});

afterAll(async () => {
  await prisma.refresh_tokens
    .deleteMany({ where: { user_id: testUserId } })
    .catch(() => {});
  await prisma.totp_login_tokens
    .deleteMany({ where: { user_id: testUserId } })
    .catch(() => {});
  await prisma.users
    .deleteMany({ where: { username: { startsWith: N } } })
    .catch(() => {});
  await app.close();
});

describe("POST /api/admin/totp/backup-verify", () => {
  it("returns 400 for missing fields", async () => {
    const res = await postBackupVerify({});
    expect(res.statusCode).toBe(400);
    expect(res.json().success).toBe(false);
  });

  it("returns 401 for an invalid login token", async () => {
    const res = await postBackupVerify({
      login_token: "not-a-real-token",
      backup_code: BACKUP_CODE_1,
    });
    expect(res.statusCode).toBe(401);
    expect(res.json().success).toBe(false);
  });

  it("rejects a wrong backup code without consuming the stored codes", async () => {
    const loginToken = await issueLoginToken(testUserId);
    const res = await postBackupVerify({
      login_token: loginToken,
      backup_code: "ZZZZ9999",
    });
    expect(res.statusCode).toBe(401);
    expect(res.json().success).toBe(false);

    // Both backup codes must still be stored — nothing consumed.
    const user = await fixtureUser();
    expect(JSON.parse(user.totp_backup_codes as string)).toHaveLength(2);

    // Reset the failed-attempt counter so this test can't lock the fixture
    // user for later tests (max_login_attempts comes from system settings).
    await prisma.users.update({
      where: { id: testUserId },
      data: { failed_login_attempts: 0 },
    });
    await prisma.totp_login_tokens.deleteMany({
      where: { user_id: testUserId },
    });
  });

  it("issues tokens for a valid login token + backup code and consumes both (single-use)", async () => {
    const loginToken = await issueLoginToken(testUserId);
    const res = await postBackupVerify({
      login_token: loginToken,
      backup_code: BACKUP_CODE_1,
    });
    expect(res.statusCode).toBe(200);
    const body = res.json();
    expect(body.success).toBe(true);
    expect(typeof body.data.access_token).toBe("string");
    expect(body.data.access_token.length).toBeGreaterThan(0);
    expect(body.data.expires_in).toBe(config.jwt.accessTokenExpiry);
    expect(body.data.user.userId).toBe(testUserId);

    // Same login completion as the TOTP path: httpOnly refresh cookie set.
    const cookie = extractCookie(res, "refresh_token");
    expect(cookie).toBeTruthy();
    const raw = rawSetCookie(res, "refresh_token")!;
    expect(raw).toMatch(/HttpOnly/i);
    expect(raw).toMatch(/SameSite=Strict/i);

    // The used backup code must be removed (single-use).
    const user = await fixtureUser();
    const remaining: string[] = JSON.parse(user.totp_backup_codes as string);
    expect(remaining).toHaveLength(1);

    // The login token must be consumed: replaying it (even with the second,
    // still-valid backup code) is rejected.
    const replay = await postBackupVerify({
      login_token: loginToken,
      backup_code: BACKUP_CODE_2,
    });
    expect(replay.statusCode).toBe(401);
  });

  it("rejects reuse of an already-consumed backup code", async () => {
    const loginToken = await issueLoginToken(testUserId);
    const res = await postBackupVerify({
      login_token: loginToken,
      backup_code: BACKUP_CODE_1,
    });
    expect(res.statusCode).toBe(401);
    expect(res.json().success).toBe(false);

    await prisma.users.update({
      where: { id: testUserId },
      data: { failed_login_attempts: 0 },
    });
    await prisma.totp_login_tokens.deleteMany({
      where: { user_id: testUserId },
    });
  });

  it("honors remember_me with a long-lived refresh token (parity with /login/totp)", async () => {
    const loginToken = await issueLoginToken(testUserId);
    const res = await postBackupVerify({
      login_token: loginToken,
      backup_code: BACKUP_CODE_2,
      remember_me: true,
    });
    expect(res.statusCode).toBe(200);

    const raw = rawSetCookie(res, "refresh_token")!;
    expect(raw).toMatch(
      new RegExp(`Max-Age=${config.jwt.refreshTokenRememberExpiry}`, "i"),
    );

    const refreshRow = await prisma.refresh_tokens.findFirst({
      where: { user_id: testUserId },
      orderBy: { id: "desc" },
    });
    expect(refreshRow?.remember_me).toBe(true);
  });
});