Jyotisha/frontend/tests/identity-auth-integration.test.ts

import assert from "node:assert/strict";
import { spawnSync } from "node:child_process";
import { createHmac } from "node:crypto";
import { fileURLToPath } from "node:url";
import test from "node:test";
import { toNextJsHandler } from "better-auth/next-js";

import {
  createIdentityAuthServices,
  createIdentityPool,
} from "../src/modules/identity/auth.ts";
import type { SelfHostedIdentityConfig } from "../src/modules/identity/config.ts";
import { FakeEmailOtpSender } from "../src/modules/identity/email/fake-email-otp-sender.ts";
import { createHostIsolatedAuthHandlers } from "../src/modules/identity/host.ts";
import {
  GET as getPasswordStatus,
  POST as setAccountPassword,
} from "../src/app/api/account/password/route.ts";
import { startPostgresFixture } from "./helpers/postgres-fixture.ts";

const runnerPath = fileURLToPath(
  new URL("../scripts/db-migrate.mjs", import.meta.url),
);
const migrationsDirectory = fileURLToPath(
  new URL("../db/migrations", import.meta.url),
);
const userHost = "staging.jyotisha.chat";
const adminHost = "admin.staging.jyotisha.chat";

function request(
  host: string,
  path: string,
  body?: Record<string, unknown>,
  cookie?: string,
): Request {
  const headers: Record<string, string> = {
    host,
    origin: `https://${host}`,
  };
  if (body) headers["content-type"] = "application/json";
  if (cookie) headers.cookie = cookie;
  return new Request(`https://${host}${path}`, {
    method: body ? "POST" : "GET",
    headers,
    body: body ? JSON.stringify(body) : undefined,
  });
}

function sessionCookie(response: Response): string {
  return (response.headers.get("set-cookie") ?? "").split(";", 1)[0];
}

type SetCookieHeaders = Headers & { getSetCookie?: () => string[] };

function responseCookieHeader(response: Response): string {
  const headers = response.headers as SetCookieHeaders;
  const setCookies = headers.getSetCookie?.()
    ?? (headers.get("set-cookie")?.split(/,(?=\s*[^;,=\s]+=[^;,]*)/g) ?? []);
  return setCookies
    .map((value) => value.split(";", 1)[0])
    .filter((value) => value.slice(value.indexOf("=") + 1).length > 0)
    .join("; ");
}

function decodeBase32(value: string): Buffer {
  const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
  let bits = "";
  for (const character of value.replace(/=+$/g, "").toUpperCase()) {
    const index = alphabet.indexOf(character);
    if (index < 0) throw new Error("invalid base32 TOTP secret");
    bits += index.toString(2).padStart(5, "0");
  }
  const bytes: number[] = [];
  for (let offset = 0; offset + 8 <= bits.length; offset += 8) {
    bytes.push(Number.parseInt(bits.slice(offset, offset + 8), 2));
  }
  return Buffer.from(bytes);
}

function totpCode(totpUri: string, now = Date.now()): string {
  const uri = new URL(totpUri);
  const secret = uri.searchParams.get("secret");
  assert.ok(secret);
  const digits = Number(uri.searchParams.get("digits") ?? "6");
  const period = Number(uri.searchParams.get("period") ?? "30");
  const counter = Buffer.alloc(8);
  counter.writeBigUInt64BE(BigInt(Math.floor(now / (period * 1_000))));
  const digest = createHmac("sha1", decodeBase32(secret)).update(counter).digest();
  const offset = digest[digest.length - 1] & 0x0f;
  const binary = (
    ((digest[offset] & 0x7f) << 24)
    | (digest[offset + 1] << 16)
    | (digest[offset + 2] << 8)
    | digest[offset + 3]
  ) >>> 0;
  return String(binary % (10 ** digits)).padStart(digits, "0");
}

const envKeys = [
  "AUTH_PROVIDER",
  "SELF_HOSTED_IDENTITY_ENABLED",
  "IDENTITY_DATABASE_URL",
  "AUTH_USER_ORIGIN",
  "ADMIN_USER_ORIGIN",
  "BETTER_AUTH_USER_SECRET",
  "RESEND_API_KEY",
  "RESEND_FROM_EMAIL",
] as const;

test("Better Auth supports shared user OTP/password sessions for admins", async () => {
  const fixture = startPostgresFixture();
  const migration = spawnSync(process.execPath, [runnerPath], {
    encoding: "utf8",
    env: {
      ...process.env,
      MIGRATIONS_DIRECTORY: migrationsDirectory,
      SCHEMA_DATABASE_URL: fixture.connectionUrl(
        "schema_owner",
        "schema-owner-test-password",
      ),
    },
  });
  assert.equal(migration.status, 0, migration.stderr);

  const config: SelfHostedIdentityConfig = {
    provider: "self-hosted",
    databaseUrl: fixture.connectionUrl(
      "identity_runtime",
      "identity-runtime-test-password",
    ),
    userOrigin: `https://${userHost}`,
    adminOrigin: `https://${adminHost}`,
    userSecret: "user-secret-that-is-at-least-32-bytes-long",
    resendApiKey: "re_test",
    resendFrom: "Jyotisha <login@staging.jyotisha.chat>",
  };
  const previousEnv = new Map(
    envKeys.map((key) => [key, process.env[key]] as const),
  );
  Object.assign(process.env, {
    AUTH_PROVIDER: "self-hosted",
    SELF_HOSTED_IDENTITY_ENABLED: "true",
    IDENTITY_DATABASE_URL: config.databaseUrl,
    AUTH_USER_ORIGIN: config.userOrigin,
    ADMIN_USER_ORIGIN: config.adminOrigin,
    BETTER_AUTH_USER_SECRET: config.userSecret,
    RESEND_API_KEY: config.resendApiKey,
    RESEND_FROM_EMAIL: config.resendFrom,
  });

  const identityGlobal = globalThis as typeof globalThis & {
    jyotishaIdentityAuth?: ReturnType<typeof createIdentityAuthServices>;
  };
  delete identityGlobal.jyotishaIdentityAuth;

  const sender = new FakeEmailOtpSender();
  const pool = createIdentityPool(config.databaseUrl);
  const services = createIdentityAuthServices(config, {
    pool,
    emailSender: sender,
  });
  const handlers = createHostIsolatedAuthHandlers(config, {
    user: toNextJsHandler(services.user),
  });

  async function otpSignIn(email: string): Promise<string> {
    const send = await handlers.POST(
      request(userHost, "/api/auth/email-otp/send-verification-otp", {
        email,
        type: "sign-in",
      }),
    );
    assert.equal(send.status, 200);
    const message = sender.messages.at(-1);
    assert.equal(message?.email, email);
    assert.equal(message?.type, "sign-in");

    const signIn = await handlers.POST(
      request(userHost, "/api/auth/sign-in/email-otp", {
        email,
        otp: message?.otp,
      }),
    );
    assert.equal(signIn.status, 200);
    const cookie = sessionCookie(signIn);
    assert.match(cookie, /^(?:__Secure-)?jyotisha-user\.session_token=/);
    return cookie;
  }

  async function passwordSignIn(
    email: string,
    password: string,
  ): Promise<Response> {
    return handlers.POST(
      request(userHost, "/api/auth/sign-in/email", { email, password }),
    );
  }

  try {
    const unauthenticatedSet = await setAccountPassword(
      request(userHost, "/api/account/password", {
        newPassword: "not-authorized",
      }),
    );
    assert.equal(unauthenticatedSet.status, 401);

    const newEmail = "new-user@example.com";
    const firstPassword = "first-password";
    const resetPassword = "reset-password";
    const newUserOtpCookie = await otpSignIn(newEmail);

    const initialStatus = await getPasswordStatus(
      request(userHost, "/api/account/password", undefined, newUserOtpCookie),
    );
    assert.equal(initialStatus.status, 200);
    assert.deepEqual(await initialStatus.json(), { hasPassword: false });

    const firstSet = await setAccountPassword(
      request(
        userHost,
        "/api/account/password",
        { newPassword: firstPassword },
        newUserOtpCookie,
      ),
    );
    assert.equal(firstSet.status, 200);

    const secondSet = await setAccountPassword(
      request(
        userHost,
        "/api/account/password",
        { newPassword: "must-not-overwrite" },
        newUserOtpCookie,
      ),
    );
    assert.equal(secondSet.status, 409);

    const storedHash = fixture.psql(
      "select password from identity.accounts where provider_id = 'credential' and user_id = (select id from identity.users where email = 'new-user@example.com')",
    );
    assert.notEqual(storedHash, firstPassword);
    assert.match(storedHash, /^[0-9a-f]{32}:[0-9a-f]{128}$/);

    const passwordLogin = await passwordSignIn(newEmail, firstPassword);
    assert.equal(passwordLogin.status, 200);
    const passwordCookie = sessionCookie(passwordLogin);
    assert.match(passwordCookie, /^(?:__Secure-)?jyotisha-user\.session_token=/);

    const wrongPassword = await passwordSignIn(newEmail, "wrong-password");
    assert.notEqual(wrongPassword.status, 200);
    assert.equal(wrongPassword.headers.has("set-cookie"), false);

    const otpLoginCookie = await otpSignIn(newEmail);
    assert.match(otpLoginCookie, /^(?:__Secure-)?jyotisha-user\.session_token=/);

    const oldOtpEmail = "otp-only@example.com";
    const firstOldOtpCookie = await otpSignIn(oldOtpEmail);
    const signOut = await handlers.POST(
      request(
        userHost,
        "/api/auth/sign-out",
        {},
        firstOldOtpCookie,
      ),
    );
    assert.equal(signOut.status, 200);
    const returningOldOtpCookie = await otpSignIn(oldOtpEmail);
    const oldOtpStatus = await getPasswordStatus(
      request(
        userHost,
        "/api/account/password",
        undefined,
        returningOldOtpCookie,
      ),
    );
    assert.deepEqual(await oldOtpStatus.json(), { hasPassword: false });
    const oldOtpSet = await setAccountPassword(
      request(
        userHost,
        "/api/account/password",
        { newPassword: "old-user-password" },
        returningOldOtpCookie,
      ),
    );
    assert.equal(oldOtpSet.status, 200);
    assert.equal(
      (await passwordSignIn(oldOtpEmail, "old-user-password")).status,
      200,
    );

    const unknownResetMessageCount = sender.messages.length;
    const unknownReset = await handlers.POST(
      request(userHost, "/api/auth/email-otp/request-password-reset", {
        email: "missing@example.com",
      }),
    );
    assert.equal(unknownReset.status, 200);
    assert.equal(sender.messages.length, unknownResetMessageCount);

    const resetRequest = await handlers.POST(
      request(userHost, "/api/auth/email-otp/request-password-reset", {
        email: newEmail,
      }),
    );
    assert.equal(resetRequest.status, 200);
    const resetMessage = sender.messages.at(-1);
    assert.equal(resetMessage?.type, "forget-password");

    const reset = await handlers.POST(
      request(userHost, "/api/auth/email-otp/reset-password", {
        email: newEmail,
        otp: resetMessage?.otp,
        password: resetPassword,
      }),
    );
    assert.equal(reset.status, 200);

    const newUserId = fixture.psql(
      "select id from identity.users where email = 'new-user@example.com'",
    );
    assert.equal(
      fixture.psql(
        `select count(*) from identity.sessions where user_id = '${newUserId}'`,
      ),
      "0",
    );
    for (const cookie of [newUserOtpCookie, passwordCookie, otpLoginCookie]) {
      assert.equal(
        await services.user.api.getSession({
          headers: new Headers({ cookie }),
        }),
        null,
      );
    }

    const oldPasswordAfterReset = await passwordSignIn(newEmail, firstPassword);
    assert.notEqual(oldPasswordAfterReset.status, 200);
    assert.equal(oldPasswordAfterReset.headers.has("set-cookie"), false);
    const newPasswordAfterReset = await passwordSignIn(newEmail, resetPassword);
    assert.equal(newPasswordAfterReset.status, 200);
    assert.match(
      sessionCookie(newPasswordAfterReset),
      /^(?:__Secure-)?jyotisha-user\.session_token=/,
    );

    fixture.psqlAs(
      "identity_runtime",
      "identity-runtime-test-password",
      "update identity.users set role = 'user,admin' where email = 'new-user@example.com'",
    );
    const adminPasswordLogin = await handlers.POST(
      request(adminHost, "/api/auth/sign-in/email", {
        email: newEmail,
        password: resetPassword,
      }),
    );
    assert.equal(adminPasswordLogin.status, 200);
    assert.match(
      sessionCookie(adminPasswordLogin),
      /^(?:__Secure-)?jyotisha-user\.session_token=/,
    );

    const enrollmentSession = sessionCookie(adminPasswordLogin);
    const enableMfa = await handlers.POST(
      request(
        adminHost,
        "/api/auth/two-factor/enable",
        { password: resetPassword },
        enrollmentSession,
      ),
    );
    assert.equal(enableMfa.status, 200);
    const enrollment = await enableMfa.json() as {
      totpURI: string;
      backupCodes: string[];
    };
    assert.match(enrollment.totpURI, /^otpauth:\/\/totp\//);
    assert.ok(enrollment.backupCodes.length >= 1);
    assert.equal(
      fixture.psql(
        `select two_factor_enabled::text from identity.users where id = '${newUserId}'`,
      ),
      "f",
    );
    const storedMfa = fixture.psql(
      `select secret, backup_codes, verified::text from identity.two_factors where user_id = '${newUserId}'`,
    );
    const encryptedBackupCodesBeforeUse = fixture.psql(
      `select backup_codes from identity.two_factors where user_id = '${newUserId}'`,
    );
    assert.equal(
      storedMfa.includes(new URL(enrollment.totpURI).searchParams.get("secret") ?? "never"),
      false,
    );
    assert.equal(storedMfa.includes(enrollment.backupCodes[0]), false);
    assert.match(storedMfa, /\|false$/);

    const verifyEnrollment = await handlers.POST(
      request(
        adminHost,
        "/api/auth/two-factor/verify-totp",
        { code: totpCode(enrollment.totpURI), trustDevice: false },
        enrollmentSession,
      ),
    );
    assert.equal(verifyEnrollment.status, 200, await verifyEnrollment.text());
    const rotatedSession = responseCookieHeader(verifyEnrollment);
    assert.match(rotatedSession, /(?:__Secure-)?jyotisha-user\.session_token=/);
    assert.equal(
      await services.user.api.getSession({ headers: new Headers({ cookie: enrollmentSession }) }),
      null,
    );
    const enrolledSession = await services.user.api.getSession({
      headers: new Headers({ cookie: rotatedSession }),
    });
    assert.equal(
      (enrolledSession?.user as { twoFactorEnabled?: boolean } | undefined)?.twoFactorEnabled,
      true,
    );

    const mfaPasswordLogin = await handlers.POST(
      request(adminHost, "/api/auth/sign-in/email", {
        email: newEmail,
        password: resetPassword,
      }),
    );
    assert.equal(mfaPasswordLogin.status, 200);
    assert.deepEqual(await mfaPasswordLogin.json(), {
      twoFactorRedirect: true,
      twoFactorMethods: ["totp"],
    });
    const nativeChallengeCookie = responseCookieHeader(mfaPasswordLogin);
    assert.match(nativeChallengeCookie, /(?:__Secure-)?jyotisha-user\.two_factor=/);
    assert.doesNotMatch(nativeChallengeCookie, /session_token=[^;]+/);

    const verifyLoginTotp = await handlers.POST(
      request(
        adminHost,
        "/api/auth/two-factor/verify-totp",
        { code: totpCode(enrollment.totpURI), trustDevice: false },
        nativeChallengeCookie,
      ),
    );
    assert.equal(verifyLoginTotp.status, 200, await verifyLoginTotp.text());
    const totpSessionCookie = responseCookieHeader(verifyLoginTotp);
    assert.match(totpSessionCookie, /(?:__Secure-)?jyotisha-user\.session_token=/);

    const backupLogin = await handlers.POST(
      request(adminHost, "/api/auth/sign-in/email", {
        email: newEmail,
        password: resetPassword,
      }),
    );
    const backupChallengeCookie = responseCookieHeader(backupLogin);
    const verifyBackup = await handlers.POST(
      request(
        adminHost,
        "/api/auth/two-factor/verify-backup-code",
        { code: enrollment.backupCodes[0], disableSession: false },
        backupChallengeCookie,
      ),
    );
    assert.equal(verifyBackup.status, 200, await verifyBackup.text());
    const backupSessionCookie = responseCookieHeader(verifyBackup);
    assert.match(backupSessionCookie, /(?:__Secure-)?jyotisha-user\.session_token=/);
    assert.notEqual(
      fixture.psql(
        `select backup_codes from identity.two_factors where user_id = '${newUserId}'`,
      ),
      encryptedBackupCodesBeforeUse,
    );

    const regenerate = await handlers.POST(
      request(
        adminHost,
        "/api/auth/two-factor/generate-backup-codes",
        { password: resetPassword },
        backupSessionCookie,
      ),
    );
    assert.equal(regenerate.status, 200, await regenerate.clone().text());
    const replacementCodes = (await regenerate.json() as { backupCodes: string[] }).backupCodes;
    assert.ok(replacementCodes.length >= 1);
    assert.notDeepEqual(replacementCodes, enrollment.backupCodes);

    const disableMfa = await handlers.POST(
      request(
        adminHost,
        "/api/auth/two-factor/disable",
        { password: resetPassword },
        backupSessionCookie,
      ),
    );
    assert.equal(disableMfa.status, 200, await disableMfa.text());
    assert.equal(
      fixture.psql(
        `select two_factor_enabled::text || ':' || (select count(*)::text from identity.two_factors where user_id = '${newUserId}') from identity.users where id = '${newUserId}'`,
      ),
      "f:0",
    );
  } finally {
    const globalServices = (
      globalThis as typeof globalThis & {
        jyotishaIdentityAuth?: ReturnType<typeof createIdentityAuthServices>;
      }
    ).jyotishaIdentityAuth;
    if (globalServices) {
      await globalServices.pool.end();
      delete identityGlobal.jyotishaIdentityAuth;
    }
    await pool.end();
    fixture.stop();
    for (const key of envKeys) {
      const value = previousEnv.get(key);
      if (value === undefined) delete process.env[key];
      else process.env[key] = value;
    }
  }
});