shd.build/src/lib/server/auth/authentication-transactions.ts

import { dev } from '$app/environment';
import { getDatabase } from '$lib/server/database';
import type { Database } from '@tursodatabase/database';
import type { Cookies } from '@sveltejs/kit';
import { createHash, randomBytes, randomUUID } from 'node:crypto';

export const AUTHENTICATION_TRANSACTION_COOKIE = 'shd_auth_tx';
export const AUTHENTICATION_TRANSACTION_COOKIE_PATH = '/auth/two-factor';
export const AUTHENTICATION_TRANSACTION_SECONDS = 5 * 60;
export const AUTHENTICATION_TRANSACTION_ATTEMPT_LIMIT = 5;

export const AUTHENTICATION_TRANSACTION_PROVIDERS = [
  'password',
  'discord',
  'magic_link',
  'passkey',
] as const;
export type AuthenticationTransactionProvider =
  (typeof AUTHENTICATION_TRANSACTION_PROVIDERS)[number];
export type AuthenticationTransactionKind = 'login' | 'step_up';
export type AuthenticationTransactionDatabase = Pick<
  Database,
  'prepare' | 'transactionAsync'
>;
export type AuthenticationTransactionTransaction = Pick<Database, 'prepare'>;

export type AuthenticationTransaction = {
  id: string;
  userId: string;
  provider: AuthenticationTransactionProvider | null;
  kind: AuthenticationTransactionKind;
  returnPath: string;
  requiredAal: 2;
  attemptCount: number;
  expiresAt: string;
};

type AuthenticationTransactionRow = {
  id: string;
  user_id: string;
  provider: AuthenticationTransactionProvider | null;
  kind: AuthenticationTransactionKind;
  return_path: string;
  required_aal: 2;
  attempt_count: number;
  expires_at: string;
  consumed_at: string | null;
};

export class AuthenticationTransactionError extends Error {
  constructor(message = 'Unable to create the authentication transaction.') {
    super(message);
    this.name = 'AuthenticationTransactionError';
  }
}

function hashToken(token: string): string {
  return createHash('sha256').update(token, 'utf8').digest('hex');
}

function databaseOrDefault(
  database?: AuthenticationTransactionDatabase,
): Promise<AuthenticationTransactionDatabase> {
  return database
    ? Promise.resolve(database)
    : (getDatabase() as Promise<AuthenticationTransactionDatabase>);
}

function nowDate(value: Date | string | number | undefined): Date {
  const date =
    value instanceof Date ? new Date(value) : new Date(value ?? Date.now());
  if (!Number.isFinite(date.getTime()))
    throw new AuthenticationTransactionError();
  return date;
}

/** Accept only a same-origin path, never an absolute or protocol-relative URL. */
export function validateAuthenticationReturnPath(value: string): string | null {
  if (
    typeof value !== 'string' ||
    value.length < 1 ||
    value.length > 2048 ||
    !value.startsWith('/') ||
    value.startsWith('//') ||
    value.includes('\\') ||
    value.includes('#') ||
    /[\u0000-\u001f\u007f]/u.test(value)
  )
    return null;
  try {
    const decoded = decodeURIComponent(value);
    if (
      decoded.startsWith('//') ||
      decoded.includes('\\') ||
      /[\u0000-\u001f\u007f]/u.test(decoded)
    )
      return null;
    const parsed = new URL(value, 'https://shd.invalid');
    if (
      parsed.origin !== 'https://shd.invalid' ||
      !parsed.pathname.startsWith('/')
    )
      return null;
  } catch {
    return null;
  }
  return value;
}

function cookieOptions() {
  return {
    path: AUTHENTICATION_TRANSACTION_COOKIE_PATH,
    httpOnly: true,
    sameSite: 'lax' as const,
    secure: !dev,
    maxAge: AUTHENTICATION_TRANSACTION_SECONDS,
  };
}

export function clearAuthenticationTransactionCookie(cookies: Cookies): void {
  cookies.delete(AUTHENTICATION_TRANSACTION_COOKIE, {
    path: AUTHENTICATION_TRANSACTION_COOKIE_PATH,
  });
}

function fromRow(row: AuthenticationTransactionRow): AuthenticationTransaction {
  return {
    id: String(row.id),
    userId: String(row.user_id),
    provider: row.provider ?? null,
    kind: row.kind,
    returnPath: String(row.return_path),
    requiredAal: 2,
    attemptCount: Number(row.attempt_count),
    expiresAt: String(row.expires_at),
  };
}

export async function createAuthenticationTransaction(
  cookies: Cookies,
  databaseInput: AuthenticationTransactionDatabase | undefined,
  options: {
    userId: string;
    kind: AuthenticationTransactionKind;
    provider?: AuthenticationTransactionProvider | null;
    returnPath: string;
    now?: Date | string | number;
  },
): Promise<AuthenticationTransaction> {
  const returnPath = validateAuthenticationReturnPath(options.returnPath);
  if (
    !options.userId ||
    !returnPath ||
    (options.kind === 'login' && !options.provider)
  )
    throw new AuthenticationTransactionError();
  const database = await databaseOrDefault(databaseInput);
  const now = nowDate(options.now);
  const createdAt = now.toISOString();
  const expiresAt = new Date(
    now.getTime() + AUTHENTICATION_TRANSACTION_SECONDS * 1000,
  ).toISOString();
  const id = randomUUID();
  const token = randomBytes(32).toString('base64url');
  await (
    await database.prepare(
      `INSERT INTO authentication_transactions(
        id,token_hash,user_id,provider,kind,return_path,required_aal,attempt_count,expires_at,created_at
      ) VALUES(?,?,?,?,?,?,?,?,?,?)`,
    )
  ).run(
    id,
    hashToken(token),
    options.userId,
    options.provider ?? null,
    options.kind,
    returnPath,
    2,
    0,
    expiresAt,
    createdAt,
  );
  cookies.set(AUTHENTICATION_TRANSACTION_COOKIE, token, cookieOptions());
  return {
    id,
    userId: options.userId,
    provider: options.provider ?? null,
    kind: options.kind,
    returnPath,
    requiredAal: 2,
    attemptCount: 0,
    expiresAt,
  };
}

export const beginAuthenticationTransaction = createAuthenticationTransaction;

export async function getAuthenticationTransaction(
  cookies: Cookies,
  databaseInput?: AuthenticationTransactionDatabase,
  nowValue?: Date | string | number,
): Promise<AuthenticationTransaction | null> {
  const token = cookies.get(AUTHENTICATION_TRANSACTION_COOKIE);
  if (!token) return null;
  const now = nowDate(nowValue).toISOString();
  const database = await databaseOrDefault(databaseInput);
  const row = (await (
    await database.prepare(
      `SELECT id,user_id,provider,kind,return_path,required_aal,attempt_count,expires_at,consumed_at
       FROM authentication_transactions
       WHERE token_hash=? AND consumed_at IS NULL AND expires_at>? AND attempt_count<?`,
    )
  ).get(hashToken(token), now, AUTHENTICATION_TRANSACTION_ATTEMPT_LIMIT)) as
    AuthenticationTransactionRow | undefined;
  if (!row) {
    clearAuthenticationTransactionCookie(cookies);
    return null;
  }
  return fromRow(row);
}

/** Atomically account for one failed factor attempt. */
export async function recordAuthenticationAttempt(
  cookies: Cookies,
  databaseInput?: AuthenticationTransactionDatabase,
  nowValue?: Date | string | number,
): Promise<{ accepted: boolean; attemptCount: number; locked: boolean }> {
  const token = cookies.get(AUTHENTICATION_TRANSACTION_COOKIE);
  if (!token) return { accepted: false, attemptCount: 0, locked: false };
  const database = await databaseOrDefault(databaseInput);
  const now = nowDate(nowValue).toISOString();
  const attempt = database.transactionAsync(async (transaction) => {
    const row = (await (
      await transaction.prepare(
        'SELECT id,attempt_count FROM authentication_transactions WHERE token_hash=? AND consumed_at IS NULL AND expires_at>? AND attempt_count<?',
      )
    ).get(hashToken(token), now, AUTHENTICATION_TRANSACTION_ATTEMPT_LIMIT)) as
      { id: string; attempt_count: number } | undefined;
    if (!row) return null;
    const updated = await (
      await transaction.prepare(
        'UPDATE authentication_transactions SET attempt_count=attempt_count+1 WHERE id=? AND consumed_at IS NULL AND attempt_count<?',
      )
    ).run(row.id, AUTHENTICATION_TRANSACTION_ATTEMPT_LIMIT);
    return updated.changes === 1 ? Number(row.attempt_count) + 1 : null;
  });
  const attemptCount = await attempt.immediate();
  if (attemptCount === null) {
    clearAuthenticationTransactionCookie(cookies);
    return { accepted: false, attemptCount: 0, locked: true };
  }
  return {
    accepted: true,
    attemptCount,
    locked: false,
  };
}

/**
 * Load a transaction while the caller's immediate transaction holds the
 * write claim. The conditional consume must happen before that transaction
 * commits.
 */
export async function claimAuthenticationTransactionInTransaction(
  transaction: AuthenticationTransactionTransaction,
  cookies: Cookies,
  nowValue?: Date | string | number,
): Promise<AuthenticationTransaction | null> {
  const token = cookies.get(AUTHENTICATION_TRANSACTION_COOKIE);
  if (!token) return null;
  const now = nowDate(nowValue).toISOString();
  const row = (await (
    await transaction.prepare(
      `SELECT id,user_id,provider,kind,return_path,required_aal,attempt_count,expires_at,consumed_at
       FROM authentication_transactions
       WHERE token_hash=? AND consumed_at IS NULL AND expires_at>? AND attempt_count<?`,
    )
  ).get(hashToken(token), now, AUTHENTICATION_TRANSACTION_ATTEMPT_LIMIT)) as
    AuthenticationTransactionRow | undefined;
  return row ? fromRow(row) : null;
}

/** Conditionally consume a transaction already loaded in this transaction. */
export async function consumeAuthenticationTransactionInTransaction(
  transaction: AuthenticationTransactionTransaction,
  cookies: Cookies,
  transactionId: string,
  nowValue?: Date | string | number,
): Promise<boolean> {
  const token = cookies.get(AUTHENTICATION_TRANSACTION_COOKIE);
  if (!token) return false;
  const now = nowDate(nowValue).toISOString();
  const consumed = await (
    await transaction.prepare(
      `UPDATE authentication_transactions SET consumed_at=?
       WHERE id=? AND token_hash=? AND consumed_at IS NULL AND expires_at>? AND attempt_count<?`,
    )
  ).run(
    now,
    transactionId,
    hashToken(token),
    now,
    AUTHENTICATION_TRANSACTION_ATTEMPT_LIMIT,
  );
  return consumed.changes === 1;
}

/** Consume a transaction in one conditional write; replay cannot succeed. */
export async function consumeAuthenticationTransaction(
  cookies: Cookies,
  databaseInput?: AuthenticationTransactionDatabase,
  nowValue?: Date | string | number,
): Promise<AuthenticationTransaction | null> {
  const token = cookies.get(AUTHENTICATION_TRANSACTION_COOKIE);
  if (!token) return null;
  const database = await databaseOrDefault(databaseInput);
  const now = nowDate(nowValue).toISOString();
  const consume = database.transactionAsync(async (transaction) => {
    const row = await claimAuthenticationTransactionInTransaction(
      transaction,
      cookies,
      now,
    );
    if (!row) return null;
    const consumed = await consumeAuthenticationTransactionInTransaction(
      transaction,
      cookies,
      row.id,
      now,
    );
    return consumed ? row : null;
  });
  const result = await consume.immediate();
  clearAuthenticationTransactionCookie(cookies);
  return result;
}

export const consumeAuthTransaction = consumeAuthenticationTransaction;