Warp/apps/web/src/matrix/crypto.ts

import { MatrixError, type AuthDict, type MatrixClient } from 'matrix-js-sdk'
import {
  decodeRecoveryKey,
  deriveRecoveryKeyFromPassphrase,
  ImportRoomKeyStage,
  type CryptoApi,
  type CryptoCallbacks,
  type KeyBackupRestoreResult,
  type VerificationRequest,
} from 'matrix-js-sdk/lib/crypto-api'
import type { SecretStorageKeyDescription } from 'matrix-js-sdk/lib/secret-storage'

// End-to-end encryption, on top of the SDK's Rust crypto stack: whether
// this session is verified, the recovery key that unlocks secret storage
// (cross-signing keys and the key backup's key), setting it all up, and
// verification requests. The stores and dialogs call these; nothing else
// touches the crypto API.

let current: MatrixClient | null = null

/** Called by client.ts once the crypto stack is up, and with null when it stops. */
export function bindCrypto(c: MatrixClient | null) {
  current = c
  cachedKey = null
}

function crypto(): CryptoApi {
  const api = current?.getCrypto()
  if (!api) throw new Error('End-to-end encryption isn’t running in this session.')
  return api
}

/** The IndexedDB name of a session's crypto store. */
export function cryptoDatabasePrefix(userId: string, deviceId: string): string {
  return `warp-web-crypto:${userId}:${deviceId}`
}

// The secret storage key. The SDK asks for it, often several times in a
// row, while it reads or writes secret storage; it is kept in memory only
// for as long as this page runs.
let cachedKey: { keyId: string; key: Uint8Array<ArrayBuffer> } | null = null

export class WrongRecoveryKeyError extends Error {
  constructor() {
    super('That isn’t the recovery key for this account. Check it and try again.')
  }
}

export const cryptoCallbacks: CryptoCallbacks = {
  async getSecretStorageKey({ keys }) {
    // Only a key the person entered (or one just created) is ever returned:
    // unlockWithRecoveryKey checks it before it is cached.
    if (cachedKey && keys[cachedKey.keyId]) return [cachedKey.keyId, cachedKey.key]
    return null
  },
  cacheSecretStorageKey(keyId, _info, key) {
    cachedKey = { keyId, key }
  },
}

export interface EncryptionStatus {
  /** The account has cross-signing keys: sessions can be verified. */
  crossSigning: boolean
  /** This session is signed by the account's own identity. */
  sessionVerified: boolean
  /** Secret storage holds the cross-signing keys: a recovery key unlocks them. */
  secretStorage: boolean
  /** The key backup on the homeserver, if there is one. */
  backupVersion: string | null
  /** This session backs its keys up to it. */
  backupActive: boolean
  /** This session holds the backup's key, so it can restore from it. */
  backupKey: boolean
}

export async function encryptionStatus(): Promise<EncryptionStatus> {
  const api = crypto()
  const c = current!
  const userId = c.getUserId()!
  const deviceId = c.getDeviceId()!
  const [crossSigning, device, secretStorage, backupInfo, backupActive, backupKey] = await Promise.all([
    api.userHasCrossSigningKeys(userId, true),
    api.getDeviceVerificationStatus(userId, deviceId),
    api.isSecretStorageReady(),
    api.getKeyBackupInfo().catch(() => null),
    api.getActiveSessionBackupVersion(),
    api.getSessionBackupPrivateKey(),
  ])
  return {
    crossSigning,
    sessionVerified: Boolean(device?.crossSigningVerified),
    secretStorage,
    backupVersion: backupInfo?.version ?? null,
    backupActive: backupActive !== null,
    backupKey: backupKey !== null,
  }
}

/** Raised when the homeserver wants the account password to go on. */
export class PasswordNeededError extends Error {
  constructor() {
    super('Enter your account password to continue.')
  }
}

function uiaSession(e: unknown): string | undefined | null {
  if (e instanceof MatrixError && e.httpStatus === 401 && Array.isArray(e.data?.flows)) return (e.data.session as string) ?? undefined
  return null
}

/**
 * Sets up encryption for the account: cross-signing keys (this session
 * signs them, so it is verified), secret storage under a new recovery key,
 * and a new key backup. Returns the recovery key to show the person once.
 *
 * With `reset`, replaces cross-signing keys that already exist: anyone who
 * verified this account has to verify it again.
 */
export async function setUpEncryption(opts: { password?: string; reset?: boolean }): Promise<string> {
  const api = crypto()
  const c = current!
  await api.bootstrapCrossSigning({
    setupNewCrossSigning: opts.reset ?? false,
    authUploadDeviceSigningKeys: async (makeRequest) => {
      try {
        await makeRequest(null)
      } catch (e) {
        const session = uiaSession(e)
        if (session === null) throw e
        if (!opts.password) throw new PasswordNeededError()
        const auth: AuthDict = {
          type: 'm.login.password',
          identifier: { type: 'm.id.user', user: c.getUserId()! },
          password: opts.password,
          session,
        }
        try {
          await makeRequest(auth)
        } catch (e2) {
          if (uiaSession(e2) !== null || (e2 instanceof MatrixError && e2.errcode === 'M_FORBIDDEN')) {
            throw new Error('That password isn’t right.')
          }
          throw e2
        }
      }
    },
  })
  let recoveryKey = ''
  await api.bootstrapSecretStorage({
    setupNewSecretStorage: true,
    setupNewKeyBackup: true,
    createSecretStorageKey: async () => {
      const key = await api.createRecoveryKeyFromPassphrase()
      recoveryKey = key.encodedPrivateKey ?? ''
      return key
    },
  })
  return recoveryKey
}

async function keyFromInput(input: string, info: SecretStorageKeyDescription): Promise<Uint8Array<ArrayBuffer>> {
  const text = input.trim()
  try {
    return decodeRecoveryKey(text)
  } catch {
    // Not a recovery key: it may be the passphrase the key was made from.
  }
  const pp = info.passphrase
  if (!pp?.salt || !pp.iterations) throw new WrongRecoveryKeyError()
  return deriveRecoveryKeyFromPassphrase(text, pp.salt, pp.iterations, pp.bits)
}

/**
 * Unlocks secret storage with the recovery key (or its passphrase): this
 * session gets the cross-signing keys and signs itself, which verifies
 * it, and takes the key backup's key so older messages can be restored.
 */
export async function unlockWithRecoveryKey(input: string): Promise<void> {
  const api = crypto()
  const c = current!
  const keyId = await c.secretStorage.getDefaultKeyId()
  const tuple = keyId ? await c.secretStorage.getKey(keyId) : null
  if (!tuple) throw new Error('This account has no secret storage. Set up encryption instead.')
  const key = await keyFromInput(input, tuple[1])
  if (!(await c.secretStorage.checkKey(key, tuple[1]))) throw new WrongRecoveryKeyError()
  cachedKey = { keyId: tuple[0], key }
  await api.bootstrapCrossSigning({})
  await api.loadSessionBackupPrivateKeyFromSecretStorage().catch(() => undefined)
  await api.checkKeyBackupAndEnable()
}

/** Fetches and imports every room key in the backup this session can read. */
export async function restoreKeyBackup(onProgress?: (done: number, total: number) => void): Promise<KeyBackupRestoreResult> {
  const api = crypto()
  await api.checkKeyBackupAndEnable()
  return api.restoreKeyBackup({
    progressCallback: (p) => {
      if (p.stage === ImportRoomKeyStage.LoadKeys) onProgress?.(p.successes + p.failures, p.total)
    },
  })
}

/** After this session was verified by another, it may have been sent the backup key. */
export async function enableBackupIfPossible() {
  const api = current?.getCrypto()
  if (!api) return
  await api.checkKeyBackupAndEnable().catch(() => undefined)
}

export function verifyOwnSession(): Promise<VerificationRequest> {
  return crypto().requestOwnUserVerification()
}

export function verifyDevice(userId: string, deviceId: string): Promise<VerificationRequest> {
  return crypto().requestDeviceVerification(userId, deviceId)
}

export interface SessionInfo {
  deviceId: string
  name?: string
  verified: boolean
  /** This browser's own session. */
  own: boolean
}

/** Someone's sessions as their keys say, with whether this session trusts each. */
export async function sessionsOf(userId: string): Promise<SessionInfo[]> {
  const api = crypto()
  const c = current!
  const map = await api.getUserDeviceInfo([userId], true)
  const devices = [...(map.get(userId)?.values() ?? [])].filter((d) => !d.dehydrated)
  const out = await Promise.all(
    devices.map(async (d) => {
      const status = await api.getDeviceVerificationStatus(userId, d.deviceId)
      return {
        deviceId: d.deviceId,
        name: d.displayName,
        verified: Boolean(status?.isVerified()),
        own: userId === c.getUserId() && d.deviceId === c.getDeviceId(),
      }
    }),
  )
  return out.sort((a, b) => Number(b.own) - Number(a.own) || (a.name ?? a.deviceId).localeCompare(b.name ?? b.deviceId))
}

/** The recovery key was entered (or created) in this page, so secret storage is open. */
export function recoveryKeyKnown(): boolean {
  return cachedKey !== null
}

/** Starts a new key backup, its key kept in secret storage. Needs the recovery key first. */
export async function createKeyBackup(): Promise<void> {
  await crypto().bootstrapSecretStorage({ setupNewKeyBackup: true })
}