webmail/lib/store/worker/host.ts

/**
 * What runs in the owner tab's worker: one database and engine per account.
 * The list of areas is the only thing that decides what the store answers.
 */

import type { StoreArea, StoreEvent } from '../area';
import { AccountEngine } from '../engine/account-engine';
import { FetchTransport, type StoreSession } from '../engine/transport';
import type { StorageAdapter } from '../storage';
import type { OpenAccountParams, StoreRpc, StoreStatus } from './protocol';

/** Files in the pool that belong to no key. A name that isn't a store file's is left alone. */
export function storeFilesToPrune(files: readonly string[], keep: readonly string[]): string[] {
  return files.filter((file) => /^\/a[0-9a-f]+\.db$/.test(file) && !keep.includes(file));
}

/** A copy of a database under way, page by page, into a new file under a new key. */
export interface StoreCopy {
  /** Negative: everything that's left. Busy: the source is in the middle of something. */
  step(pages: number): 'more' | 'done' | 'busy';
  remaining(): number;
  pageCount(): number;
  /** Done: the file stays. */
  finish(): void;
  /** Not done: the file goes. */
  abort(): void;
}

export interface HostStorage {
  open(file: string, keyHex: string): Promise<StorageAdapter>;
  remove(file: string): boolean;
  startCopy?(source: StorageAdapter, file: string, keyHex: string): Promise<StoreCopy>;
}

const COPY_PAGES_PER_STEP = 256;
/** What's left for after the engine has stopped: a write after a copy's end isn't in it. */
const COPY_LAST_PAGES = 64;
const COPY_BUSY_WAIT_MS = 50;

interface OpenAccount {
  engine: AccountEngine;
  db: StorageAdapter;
  file: string;
  auth: string | null;
  pendingAuth: Array<(header: string | null) => void>;
  session: StoreSession;
  /** A re-encryption under way: the database can't close under it. */
  copying?: { stop: boolean; settled: Promise<void> };
}

export interface StoreHostOptions {
  areas: readonly StoreArea[];
  storage: HostStorage;
  emit(event: StoreEvent): void;
  needAuth(account: string): void;
}

/** The owner hasn't opened this account yet. The caller may try again shortly. */
function rekeyStopped(): Error {
  const error = new Error('The re-encryption was stopped');
  error.name = 'StoreRekeyStoppedError';
  return error;
}

function accountClosed(): Error {
  const error = new Error('The local store is not open for this account');
  error.name = 'StoreAccountClosedError';
  return error;
}

export class StoreHost {
  private readonly accounts = new Map<string, OpenAccount>();
  private online = true;
  private readonly synced = new Map<string, number>();

  constructor(private readonly options: StoreHostOptions) {}

  async openAccount(params: OpenAccountParams): Promise<{ created: boolean }> {
    const existing = this.accounts.get(params.account);
    if (existing) {
      if (params.session) this.setSession(params.account, params.session);
      this.setAuth(params.account, params.auth);
      if (params.scope) existing.engine.setScope(params.scope);
      if (params.invertedKeywordSort !== null) existing.engine.setInvertedKeywordSort(params.invertedKeywordSort);
      void existing.engine.requestSync();
      return { created: false };
    }
    const db = await this.options.storage.open(params.file, params.keyHex);
    const created = !db.get("SELECT 1 AS yes FROM sqlite_master WHERE name = 'schema_migration'");
    // Nothing to sync with yet. No file stays behind, or the next open couldn't tell that this one was new.
    if (created && !params.session) {
      db.close();
      this.options.storage.remove(params.file);
      throw new Error('No session for this account yet');
    }
    const session = params.session ?? AccountEngine.storedSession(db);
    if (!session) {
      db.close();
      throw new Error('No session for this account yet');
    }
    const account: OpenAccount = { db, file: params.file, auth: params.auth, pendingAuth: [], session, engine: undefined as unknown as AccountEngine };
    const transport = new FetchTransport({
      session: () => account.session,
      auth: () => account.auth,
      refreshAuth: () => new Promise((resolve) => {
        account.pendingAuth.push(resolve);
        this.options.needAuth(params.account);
      }),
    });
    account.engine = new AccountEngine({
      account: params.account,
      db,
      session,
      transport,
      areas: this.options.areas,
      scope: params.scope,
      invertedKeywordSort: params.invertedKeywordSort ?? undefined,
      emit: (event) => {
        if (event.type === 'synced') this.synced.set(event.account, event.at);
        this.options.emit(event);
      },
    });
    account.engine.setOnline(this.online);
    this.accounts.set(params.account, account);
    void account.engine.requestSync();
    return { created };
  }

  setAuth(account: string, auth: string | null): void {
    const open = this.accounts.get(account);
    if (!open) return;
    open.auth = auth;
    for (const resolve of open.pendingAuth.splice(0)) resolve(auth);
  }

  setSession(account: string, session: StoreSession): void {
    const open = this.accounts.get(account);
    if (!open) return;
    open.session = session;
    open.engine.setSession(session);
  }

  setScope(account: string, scope: Parameters<AccountEngine['setScope']>[0]): void {
    const open = this.accounts.get(account);
    if (!open) return;
    if (open.engine.setScope(scope)) void open.engine.requestSync();
  }

  setInvertedKeywordSort(account: string, inverted: boolean): void {
    this.accounts.get(account)?.engine.setInvertedKeywordSort(inverted);
  }

  async closeAccount(account: string): Promise<void> {
    const open = this.accounts.get(account);
    if (!open) return;
    if (open.copying) {
      open.copying.stop = true;
      await open.copying.settled;
    }
    this.accounts.delete(account);
    await open.engine.close();
    for (const resolve of open.pendingAuth.splice(0)) resolve(null);
    open.db.close();
  }

  /**
   * Copies an open account into a new file under a new key, while it goes on working: the source's
   * own writes reach the copy too. Then stops the engine for the last pages, and ends with the
   * account closed and nothing written since, so the old file and the new one are the same data.
   * On a failure the new file goes and the account stays open, if it can.
   */
  async rekey(account: string, file: string, keyHex: string): Promise<void> {
    const open = this.accounts.get(account);
    if (!open) throw accountClosed();
    if (!this.options.storage.startCopy) throw new Error('This storage can\'t re-encrypt');
    if (open.copying) throw new Error('A re-encryption is under way');
    let settle = () => {};
    const copying = { stop: false, settled: new Promise<void>((resolve) => { settle = resolve; }) };
    open.copying = copying;
    let copy: StoreCopy | null = null;
    const emit = (percent: number) => this.options.emit({ type: 'rekey', account, percent });
    try {
      copy = await this.options.storage.startCopy(open.db, file, keyHex);
      const yieldNow = (ms = 0) => new Promise((resolve) => setTimeout(resolve, ms));
      // One page first: only then does SQLite know how much there is.
      let answer = copy.step(1);
      for (;;) {
        if (copying.stop) throw rekeyStopped();
        if (answer === 'done') break;
        if (answer === 'busy') {
          await yieldNow(COPY_BUSY_WAIT_MS);
        } else {
          const left = copy.remaining();
          if (left <= COPY_LAST_PAGES) break;
          const total = copy.pageCount();
          emit(Math.min(99, Math.floor(((total - left) / Math.max(total, 1)) * 100)));
          await yieldNow();
        }
        if (copying.stop) throw rekeyStopped();
        answer = copy.step(Math.min(COPY_PAGES_PER_STEP, Math.max(1, copy.remaining() - COPY_LAST_PAGES)));
      }
      if (copying.stop) throw rekeyStopped();
      // Calls wait for the account to open again from here. Other tabs try again by themselves.
      this.accounts.delete(account);
      await open.engine.close();
      for (const resolve of open.pendingAuth.splice(0)) resolve(null);
      // Nothing writes now. A copy that was short enough to end above is whole only if nothing
      // wrote after, so it goes round again.
      let last = answer === 'done' ? 'again' : copy.step(-1);
      if (last === 'again') {
        copy.abort();
        copy = await this.options.storage.startCopy(open.db, file, keyHex);
        last = copy.step(-1);
      }
      for (let tries = 0; last === 'busy' && tries < 20; tries += 1) {
        await yieldNow(COPY_BUSY_WAIT_MS);
        last = copy.step(-1);
      }
      if (last !== 'done') throw new Error('The copy did not finish');
      copy.finish();
      copy = null;
      open.db.close();
      emit(100);
    } catch (error) {
      copy?.abort();
      if (!this.accounts.has(account)) {
        try {
          open.db.close();
        } catch {
          // The page opens the old file again.
        }
      }
      throw error;
    } finally {
      open.copying = undefined;
      settle();
    }
  }

  async removeAccount(account: string, file: string): Promise<boolean> {
    await this.closeAccount(account);
    this.synced.delete(account);
    const removed = this.options.storage.remove(file);
    this.options.emit({ type: 'removed', account });
    return removed;
  }

  async closeAll(): Promise<void> {
    await Promise.all([...this.accounts.keys()].map((account) => this.closeAccount(account)));
  }

  /** The calls other tabs can make. */
  async call(method: keyof StoreRpc, params: unknown): Promise<unknown> {
    switch (method) {
      case 'answer': {
        const { account, calls } = params as Parameters<StoreRpc['answer']>[0];
        const open = this.accounts.get(account);
        if (!open) throw accountClosed();
        return open.engine.answer(calls);
      }
      case 'sync': {
        const { account } = params as Parameters<StoreRpc['sync']>[0];
        await this.accounts.get(account)?.engine.requestSync();
        return undefined;
      }
      case 'noteStateChange': {
        const { account, changed } = params as Parameters<StoreRpc['noteStateChange']>[0];
        this.accounts.get(account)?.engine.noteStateChange(changed);
        return undefined;
      }
      case 'setOnline': {
        const { online } = params as Parameters<StoreRpc['setOnline']>[0];
        this.online = online;
        for (const open of this.accounts.values()) open.engine.setOnline(online);
        return undefined;
      }
      case 'status': {
        const { account } = params as Parameters<StoreRpc['status']>[0];
        return this.status(account);
      }
      case 'area': {
        const { account, name, params: callParams } = params as Parameters<StoreRpc['area']>[0];
        const open = this.accounts.get(account);
        if (!open) throw accountClosed();
        return open.engine.areaCall(name, callParams);
      }
      case 'setScope': {
        const { account, scope } = params as Parameters<StoreRpc['setScope']>[0];
        if (!this.accounts.has(account)) throw accountClosed();
        this.setScope(account, scope);
        return undefined;
      }
      case 'removeAccount': {
        const { account, file } = params as Parameters<StoreRpc['removeAccount']>[0];
        return this.removeAccount(account, file);
      }
      default:
        throw new Error(`Unknown store call ${String(method)}`);
    }
  }

  private status(account: string): StoreStatus | null {
    const open = this.accounts.get(account);
    if (!open) return null;
    return {
      account,
      open: true,
      emails: Number(open.db.get<{ n: number }>('SELECT COUNT(*) AS n FROM email')?.n ?? 0),
      bodies: Number(open.db.get<{ n: number }>('SELECT COUNT(*) AS n FROM email_body')?.n ?? 0),
      lastSyncedAt: this.synced.get(account) ?? null,
    };
  }
}