Refactor provisioning flow

ts/state/smart/InstallScreen.tsx

@@ -12,7 +12,7 @@ import { hasExpired as hasExpiredSelector } from '../selectors/expiration';
 import * as log from '../../logging/log';
 import type { Loadable } from '../../util/loadable';
 import { LoadingState } from '../../util/loadable';
-import { assertDev } from '../../util/assert';
+import { assertDev, strictAssert } from '../../util/assert';
 import { explodePromise } from '../../util/explodePromise';
 import { missingCaseError } from '../../util/missingCaseError';
 import * as Registration from '../../util/registration';
@@ -26,7 +26,7 @@ import { MAX_DEVICE_NAME_LENGTH } from '../../components/installScreen/InstallSc
 import { WidthBreakpoint } from '../../components/_util';
 import { HTTPError } from '../../textsecure/Errors';
 import { isRecord } from '../../util/isRecord';
-import type { ConfirmNumberResultType } from '../../textsecure/AccountManager';
+import { Provisioner } from '../../textsecure/Provisioner';
 import * as Errors from '../../types/errors';
 import { normalizeDeviceName } from '../../util/normalizeDeviceName';
 import OS from '../../util/os/osMain';
@@ -89,8 +89,8 @@ function classifyError(
         return { loadError: LoadError.Unknown };
     }
   }
-  // AccountManager.registerSecondDevice uses this specific "websocket closed" error
-  //   message.
+  // AccountManager.registerSecondDevice uses this specific "websocket closed"
+  // error message.
   if (isRecord(err) && err.message === 'websocket closed') {
     return { installError: InstallError.ConnectionFailed };
   }
@@ -197,82 +197,80 @@ export const SmartInstallScreen = memo(function SmartInstallScreen() {
 
   useEffect(() => {
     let hasCleanedUp = false;
-    const qrCodeResolution = explodePromise<void>();
 
+    const { server } = window.textsecure;
+    strictAssert(server, 'Expected a server');
+
+    let provisioner = new Provisioner(server);
     const accountManager = window.getAccountManager();
-    assertDev(accountManager, 'Expected an account manager');
-
-    const updateProvisioningUrl = (value: string): void => {
-      if (hasCleanedUp) {
-        return;
-      }
-      qrCodeResolution.resolve();
-      setProvisioningUrl(value);
-    };
-
-    const confirmNumber = async (): Promise<ConfirmNumberResultType> => {
-      if (hasCleanedUp) {
-        throw new Error('Cannot confirm number; the component was unmounted');
-      }
-      onQrCodeScanned();
-
-      let deviceName: string;
-      let backupFileData: Uint8Array | undefined;
-      if (window.SignalCI) {
-        ({ deviceName, backupData: backupFileData } = window.SignalCI);
-      } else {
-        deviceName = await chooseDeviceNamePromiseWrapperRef.current.promise;
-        const backupFile =
-          await chooseBackupFilePromiseWrapperRef.current.promise;
-
-        backupFileData = backupFile ? await fileToBytes(backupFile) : undefined;
-      }
-
-      if (hasCleanedUp) {
-        throw new Error('Cannot confirm number; the component was unmounted');
-      }
-
-      // Delete all data from the database unless we're in the middle of a re-link.
-      //   Without this, the app restarts at certain times and can cause weird things to
-      //   happen, like data from a previous light import showing up after a new install.
-      const shouldRetainData = Registration.everDone();
-      if (!shouldRetainData) {
-        try {
-          await window.textsecure.storage.protocol.removeAllData();
-        } catch (error) {
-          log.error(
-            'confirmNumber: error clearing database',
-            Errors.toLogFormat(error)
-          );
-        }
-      }
-
-      if (hasCleanedUp) {
-        throw new Error('Cannot confirm number; the component was unmounted');
-      }
-
-      return { deviceName, backupFile: backupFileData };
-    };
+    strictAssert(accountManager, 'Expected an account manager');
 
     async function getQRCode(): Promise<void> {
       const sleepError = new TimeoutError();
       try {
-        const qrCodePromise = accountManager.registerSecondDevice(
-          updateProvisioningUrl,
-          confirmNumber
-        );
+        const qrCodePromise = provisioner.getURL();
         const sleepMs = qrCodeBackOff.getAndIncrement();
         log.info(`InstallScreen/getQRCode: race to ${sleepMs}ms`);
-        await Promise.all([
-          pTimeout(qrCodeResolution.promise, sleepMs, sleepError),
 
-          // Note that `registerSecondDevice` resolves once the registration
-          // is fully complete and thus should not be subjected to a timeout.
-          qrCodePromise,
-        ]);
+        const url = await pTimeout(qrCodePromise, sleepMs, sleepError);
+        if (hasCleanedUp) {
+          return;
+        }
 
         window.IPC.removeSetupMenuItems();
+        setProvisioningUrl(url);
+
+        await provisioner.waitForEnvelope();
+        onQrCodeScanned();
+
+        let deviceName: string;
+        let backupFileData: Uint8Array | undefined;
+        if (window.SignalCI) {
+          ({ deviceName, backupData: backupFileData } = window.SignalCI);
+        } else {
+          deviceName = await chooseDeviceNamePromiseWrapperRef.current.promise;
+          const backupFile =
+            await chooseBackupFilePromiseWrapperRef.current.promise;
+
+          backupFileData = backupFile
+            ? await fileToBytes(backupFile)
+            : undefined;
+        }
+
+        if (hasCleanedUp) {
+          throw new Error('Cannot confirm number; the component was unmounted');
+        }
+
+        // Delete all data from the database unless we're in the middle of a
+        //   re-link. Without this, the app restarts at certain times and can
+        //   cause weird things to happen, like data from a previous light
+        //   import showing up after a new install.
+        const shouldRetainData = Registration.everDone();
+        if (!shouldRetainData) {
+          try {
+            await window.textsecure.storage.protocol.removeAllData();
+          } catch (error) {
+            log.error(
+              'confirmNumber: error clearing database',
+              Errors.toLogFormat(error)
+            );
+          }
+        }
+
+        if (hasCleanedUp) {
+          throw new Error('Cannot confirm number; the component was unmounted');
+        }
+
+        const data = provisioner.prepareLinkData({
+          deviceName,
+          backupFile: backupFileData,
+        });
+        await accountManager.registerSecondDevice(data);
       } catch (error) {
+        provisioner.close();
+        strictAssert(server, 'Expected a server');
+        provisioner = new Provisioner(server);
+
         log.error(
           'account.registerSecondDevice: got an error',
           Errors.toLogFormat(error)

ts/textsecure/AccountManager.ts

@@ -21,9 +21,6 @@ import type {
   KyberPreKeyType,
   PniKeyMaterialType,
 } from './Types.d';
-import ProvisioningCipher from './ProvisioningCipher';
-import type { IncomingWebSocketRequest } from './WebsocketResources';
-import { ServerRequestType } from './WebsocketResources';
 import createTaskWithTimeout from './TaskWithTimeout';
 import * as Bytes from '../Bytes';
 import * as Errors from '../types/errors';
@@ -46,7 +43,6 @@ import {
 import type { AciString, PniString, ServiceIdString } from '../types/ServiceId';
 import {
   isUntaggedPniString,
-  normalizePni,
   ServiceIdKind,
   toTaggedPni,
 } from '../types/ServiceId';
@@ -61,7 +57,6 @@ import { missingCaseError } from '../util/missingCaseError';
 import { SignalService as Proto } from '../protobuf';
 import * as log from '../logging/log';
 import type { StorageAccessType } from '../types/Storage';
-import { linkDeviceRoute } from '../util/signalRoutes';
 import { getRelativePath, createName } from '../util/attachmentPath';
 import { isBackupEnabled } from '../util/isBackupEnabled';
 
@@ -116,7 +111,7 @@ const SIGNED_PRE_KEY_UPDATE_TIME_KEY: StorageKeyByServiceIdKind = {
   [ServiceIdKind.PNI]: 'signedKeyUpdateTimePNI',
 };
 
-enum AccountType {
+export enum AccountType {
   Primary = 'Primary',
   Linked = 'Linked',
 }
@@ -146,7 +141,7 @@ type CreatePrimaryDeviceOptionsType = Readonly<{
 }> &
   CreateAccountSharedOptionsType;
 
-type CreateLinkedDeviceOptionsType = Readonly<{
+export type CreateLinkedDeviceOptionsType = Readonly<{
   type: AccountType.Linked;
 
   deviceName: string;
@@ -327,146 +322,26 @@ export default class AccountManager extends EventTarget {
       const accessKey = deriveAccessKey(profileKey);
       const masterKey = getRandomBytes(MASTER_KEY_LENGTH);
 
-      const registrationBaton = this.server.startRegistration();
-      try {
-        await this.createAccount({
-          type: AccountType.Primary,
-          number,
-          verificationCode,
-          sessionId,
-          aciKeyPair,
-          pniKeyPair,
-          profileKey,
-          accessKey,
-          masterKey,
-          readReceipts: true,
-        });
-      } finally {
-        this.server.finishRegistration(registrationBaton);
-      }
-      await this.registrationDone();
+      await this.createAccount({
+        type: AccountType.Primary,
+        number,
+        verificationCode,
+        sessionId,
+        aciKeyPair,
+        pniKeyPair,
+        profileKey,
+        accessKey,
+        masterKey,
+        readReceipts: true,
+      });
     });
   }
 
   async registerSecondDevice(
-    setProvisioningUrl: (url: string) => void,
-    confirmNumber: (number?: string) => Promise<ConfirmNumberResultType>
+    options: CreateLinkedDeviceOptionsType
   ): Promise<void> {
-    const provisioningCipher = new ProvisioningCipher();
-    const pubKey = await provisioningCipher.getPublicKey();
-
-    let envelopeCallbacks:
-      | {
-          resolve(data: Proto.ProvisionEnvelope): void;
-          reject(error: Error): void;
-        }
-      | undefined;
-    const envelopePromise = new Promise<Proto.ProvisionEnvelope>(
-      (resolve, reject) => {
-        envelopeCallbacks = { resolve, reject };
-      }
-    );
-
-    const wsr = await this.server.getProvisioningResource({
-      handleRequest(request: IncomingWebSocketRequest) {
-        if (
-          request.requestType === ServerRequestType.ProvisioningAddress &&
-          request.body
-        ) {
-          const proto = Proto.ProvisioningUuid.decode(request.body);
-          const { uuid } = proto;
-          if (!uuid) {
-            throw new Error('registerSecondDevice: expected a UUID');
-          }
-          const url = linkDeviceRoute
-            .toAppUrl({
-              uuid,
-              pubKey: Bytes.toBase64(pubKey),
-            })
-            .toString();
-
-          window.SignalCI?.setProvisioningURL(url);
-
-          setProvisioningUrl(url);
-          request.respond(200, 'OK');
-        } else if (
-          request.requestType === ServerRequestType.ProvisioningMessage &&
-          request.body
-        ) {
-          const envelope = Proto.ProvisionEnvelope.decode(request.body);
-          request.respond(200, 'OK');
-          wsr.close();
-          envelopeCallbacks?.resolve(envelope);
-        } else {
-          log.error('Unknown websocket message', request.requestType);
-        }
-      },
-    });
-
-    log.info('provisioning socket open');
-
-    wsr.addEventListener('close', ({ code, reason }) => {
-      log.info(`provisioning socket closed. Code: ${code} Reason: ${reason}`);
-
-      // Note: if we have resolved the envelope already - this has no effect
-      envelopeCallbacks?.reject(new Error('websocket closed'));
-    });
-
-    const envelope = await envelopePromise;
-    const provisionMessage = await provisioningCipher.decrypt(envelope);
-
     await this.queueTask(async () => {
-      const { deviceName, backupFile } = await confirmNumber(
-        provisionMessage.number
-      );
-      if (typeof deviceName !== 'string' || deviceName.length === 0) {
-        throw new Error(
-          'AccountManager.registerSecondDevice: Invalid device name'
-        );
-      }
-      if (
-        !provisionMessage.number ||
-        !provisionMessage.provisioningCode ||
-        !provisionMessage.aciKeyPair ||
-        !provisionMessage.pniKeyPair ||
-        !provisionMessage.aci ||
-        !Bytes.isNotEmpty(provisionMessage.profileKey) ||
-        !Bytes.isNotEmpty(provisionMessage.masterKey) ||
-        !isUntaggedPniString(provisionMessage.untaggedPni)
-      ) {
-        throw new Error(
-          'AccountManager.registerSecondDevice: Provision message was missing key data'
-        );
-      }
-
-      const ourAci = normalizeAci(provisionMessage.aci, 'provisionMessage.aci');
-      const ourPni = normalizePni(
-        toTaggedPni(provisionMessage.untaggedPni),
-        'provisionMessage.pni'
-      );
-
-      const registrationBaton = this.server.startRegistration();
-      try {
-        await this.createAccount({
-          type: AccountType.Linked,
-          number: provisionMessage.number,
-          verificationCode: provisionMessage.provisioningCode,
-          aciKeyPair: provisionMessage.aciKeyPair,
-          pniKeyPair: provisionMessage.pniKeyPair,
-          profileKey: provisionMessage.profileKey,
-          deviceName,
-          backupFile,
-          userAgent: provisionMessage.userAgent,
-          ourAci,
-          ourPni,
-          readReceipts: Boolean(provisionMessage.readReceipts),
-          masterKey: provisionMessage.masterKey,
-        });
-      } finally {
-        this.server.finishRegistration(registrationBaton);
-      }
-
-      await this.registrationDone();
+      await this.createAccount(options);
     });
   }
 
@@ -1021,6 +896,18 @@ export default class AccountManager extends EventTarget {
 
   private async createAccount(
     options: CreateAccountOptionsType
+  ): Promise<void> {
+    const registrationBaton = this.server.startRegistration();
+    try {
+      await this.doCreateAccount(options);
+    } finally {
+      this.server.finishRegistration(registrationBaton);
+    }
+    await this.registrationDone();
+  }
+
+  private async doCreateAccount(
+    options: CreateAccountOptionsType
   ): Promise<void> {
     const {
       number,

ts/textsecure/Provisioner.ts

@@ -0,0 +1,266 @@
+// Copyright 2024 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import {
+  type ExplodePromiseResultType,
+  explodePromise,
+} from '../util/explodePromise';
+import { linkDeviceRoute } from '../util/signalRoutes';
+import { strictAssert } from '../util/assert';
+import { normalizeAci } from '../util/normalizeAci';
+import {
+  isUntaggedPniString,
+  normalizePni,
+  toTaggedPni,
+} from '../types/ServiceId';
+import * as Errors from '../types/errors';
+import { SignalService as Proto } from '../protobuf';
+import * as Bytes from '../Bytes';
+import * as log from '../logging/log';
+import { type WebAPIType } from './WebAPI';
+import ProvisioningCipher, {
+  type ProvisionDecryptResult,
+} from './ProvisioningCipher';
+import {
+  type CreateLinkedDeviceOptionsType,
+  AccountType,
+} from './AccountManager';
+import {
+  type IWebSocketResource,
+  type IncomingWebSocketRequest,
+  ServerRequestType,
+} from './WebsocketResources';
+
+enum Step {
+  Idle = 'Idle',
+  Connecting = 'Connecting',
+  WaitingForURL = 'WaitingForURL',
+  WaitingForEnvelope = 'WaitingForEnvelope',
+  ReadyToLink = 'ReadyToLink',
+  Done = 'Done',
+}
+
+type StateType = Readonly<
+  | {
+      step: Step.Idle;
+    }
+  | {
+      step: Step.Connecting;
+    }
+  | {
+      step: Step.WaitingForURL;
+      url: ExplodePromiseResultType<string>;
+    }
+  | {
+      step: Step.WaitingForEnvelope;
+      done: ExplodePromiseResultType<void>;
+    }
+  | {
+      step: Step.ReadyToLink;
+      envelope: ProvisionDecryptResult;
+    }
+  | {
+      step: Step.Done;
+    }
+>;
+
+export type PrepareLinkDataOptionsType = Readonly<{
+  deviceName: string;
+  backupFile?: Uint8Array;
+}>;
+
+export class Provisioner {
+  private readonly cipher = new ProvisioningCipher();
+
+  private state: StateType = { step: Step.Idle };
+  private wsr: IWebSocketResource | undefined;
+
+  constructor(private readonly server: WebAPIType) {}
+
+  public close(error = new Error('Provisioner closed')): void {
+    try {
+      this.wsr?.close();
+    } catch {
+      // Best effort
+    }
+
+    const prevState = this.state;
+    this.state = { step: Step.Done };
+
+    if (prevState.step === Step.WaitingForURL) {
+      prevState.url.reject(error);
+    } else if (prevState.step === Step.WaitingForEnvelope) {
+      prevState.done.reject(error);
+    }
+  }
+
+  public async getURL(): Promise<string> {
+    strictAssert(
+      this.state.step === Step.Idle,
+      `Invalid state for getURL: ${this.state.step}`
+    );
+    this.state = { step: Step.Connecting };
+
+    const wsr = await this.server.getProvisioningResource({
+      handleRequest: (request: IncomingWebSocketRequest) => {
+        try {
+          this.handleRequest(request);
+        } catch (error) {
+          log.error(
+            'Provisioner.handleRequest: failure',
+            Errors.toLogFormat(error)
+          );
+          this.close();
+        }
+      },
+    });
+    this.wsr = wsr;
+
+    if (this.state.step !== Step.Connecting) {
+      this.close();
+      throw new Error('Provisioner closed early');
+    }
+
+    this.state = {
+      step: Step.WaitingForURL,
+      url: explodePromise(),
+    };
+
+    wsr.addEventListener('close', ({ code, reason }) => {
+      if (this.state.step === Step.ReadyToLink) {
+        // WebSocket close is not an issue since we no longer need it
+        return;
+      }
+
+      log.info(`provisioning socket closed. Code: ${code} Reason: ${reason}`);
+      this.close(new Error('websocket closed'));
+    });
+
+    return this.state.url.promise;
+  }
+
+  public async waitForEnvelope(): Promise<void> {
+    strictAssert(
+      this.state.step === Step.WaitingForEnvelope,
+      `Invalid state for waitForEnvelope: ${this.state.step}`
+    );
+    await this.state.done.promise;
+  }
+
+  public prepareLinkData({
+    deviceName,
+    backupFile,
+  }: PrepareLinkDataOptionsType): CreateLinkedDeviceOptionsType {
+    strictAssert(
+      this.state.step === Step.ReadyToLink,
+      `Invalid state for prepareLinkData: ${this.state.step}`
+    );
+    const { envelope } = this.state;
+    this.state = { step: Step.Done };
+
+    const {
+      number,
+      provisioningCode,
+      aciKeyPair,
+      pniKeyPair,
+      aci,
+      profileKey,
+      masterKey,
+      untaggedPni,
+      userAgent,
+      readReceipts,
+    } = envelope;
+
+    strictAssert(number, 'prepareLinkData: missing number');
+    strictAssert(provisioningCode, 'prepareLinkData: missing provisioningCode');
+    strictAssert(aciKeyPair, 'prepareLinkData: missing aciKeyPair');
+    strictAssert(pniKeyPair, 'prepareLinkData: missing pniKeyPair');
+    strictAssert(aci, 'prepareLinkData: missing aci');
+    strictAssert(
+      Bytes.isNotEmpty(profileKey),
+      'prepareLinkData: missing profileKey'
+    );
+    strictAssert(
+      Bytes.isNotEmpty(masterKey),
+      'prepareLinkData: missing masterKey'
+    );
+    strictAssert(
+      isUntaggedPniString(untaggedPni),
+      'prepareLinkData: invalid untaggedPni'
+    );
+
+    const ourAci = normalizeAci(aci, 'provisionMessage.aci');
+    const ourPni = normalizePni(
+      toTaggedPni(untaggedPni),
+      'provisionMessage.pni'
+    );
+
+    return {
+      type: AccountType.Linked,
+      number,
+      verificationCode: provisioningCode,
+      aciKeyPair,
+      pniKeyPair,
+      profileKey,
+      deviceName,
+      backupFile,
+      userAgent,
+      ourAci,
+      ourPni,
+      readReceipts: Boolean(readReceipts),
+      masterKey,
+    };
+  }
+
+  private handleRequest(request: IncomingWebSocketRequest): void {
+    const pubKey = this.cipher.getPublicKey();
+
+    if (
+      request.requestType === ServerRequestType.ProvisioningAddress &&
+      request.body
+    ) {
+      strictAssert(
+        this.state.step === Step.WaitingForURL,
+        `Unexpected provisioning address, state: ${this.state}`
+      );
+      const prevState = this.state;
+      this.state = { step: Step.WaitingForEnvelope, done: explodePromise() };
+
+      const proto = Proto.ProvisioningUuid.decode(request.body);
+      const { uuid } = proto;
+      strictAssert(uuid, 'Provisioner.getURL: expected a UUID');
+
+      const url = linkDeviceRoute
+        .toAppUrl({
+          uuid,
+          pubKey: Bytes.toBase64(pubKey),
+        })
+        .toString();
+
+      window.SignalCI?.setProvisioningURL(url);
+      prevState.url.resolve(url);
+
+      request.respond(200, 'OK');
+    } else if (
+      request.requestType === ServerRequestType.ProvisioningMessage &&
+      request.body
+    ) {
+      strictAssert(
+        this.state.step === Step.WaitingForEnvelope,
+        `Unexpected provisioning address, state: ${this.state}`
+      );
+      const prevState = this.state;
+
+      const ciphertext = Proto.ProvisionEnvelope.decode(request.body);
+      const message = this.cipher.decrypt(ciphertext);
+
+      this.state = { step: Step.ReadyToLink, envelope: message };
+      request.respond(200, 'OK');
+      this.wsr?.close();
+
+      prevState.done.resolve();
+    } else {
+      log.error('Unknown websocket message', request.requestType);
+    }
+  }
+}

ts/textsecure/ProvisioningCipher.ts

@@ -15,7 +15,7 @@ import { SignalService as Proto } from '../protobuf';
 import { strictAssert } from '../util/assert';
 import { dropNull } from '../util/dropNull';
 
-type ProvisionDecryptResult = {
+export type ProvisionDecryptResult = Readonly<{
   aciKeyPair: KeyPairType;
   pniKeyPair?: KeyPairType;
   number?: string;
@@ -26,14 +26,12 @@ type ProvisionDecryptResult = {
   readReceipts?: boolean;
   profileKey?: Uint8Array;
   masterKey?: Uint8Array;
-};
+}>;
 
 class ProvisioningCipherInner {
   keyPair?: KeyPairType;
 
-  async decrypt(
-    provisionEnvelope: Proto.ProvisionEnvelope
-  ): Promise<ProvisionDecryptResult> {
+  decrypt(provisionEnvelope: Proto.ProvisionEnvelope): ProvisionDecryptResult {
     strictAssert(
       provisionEnvelope.publicKey,
       'Missing publicKey in ProvisionEnvelope'
@@ -77,7 +75,7 @@ class ProvisioningCipherInner {
     strictAssert(aci, 'Missing aci in provisioning message');
     strictAssert(pni, 'Missing pni in provisioning message');
 
-    const ret: ProvisionDecryptResult = {
+    return {
       aciKeyPair,
       pniKeyPair,
       number: dropNull(provisionMessage.number),
@@ -86,17 +84,16 @@ class ProvisioningCipherInner {
       provisioningCode: dropNull(provisionMessage.provisioningCode),
       userAgent: dropNull(provisionMessage.userAgent),
       readReceipts: provisionMessage.readReceipts ?? false,
+      profileKey: Bytes.isNotEmpty(provisionMessage.profileKey)
+        ? provisionMessage.profileKey
+        : undefined,
+      masterKey: Bytes.isNotEmpty(provisionMessage.masterKey)
+        ? provisionMessage.masterKey
+        : undefined,
     };
-    if (Bytes.isNotEmpty(provisionMessage.profileKey)) {
-      ret.profileKey = provisionMessage.profileKey;
-    }
-    if (Bytes.isNotEmpty(provisionMessage.masterKey)) {
-      ret.masterKey = provisionMessage.masterKey;
-    }
-    return ret;
   }
 
-  async getPublicKey(): Promise<Uint8Array> {
+  getPublicKey(): Uint8Array {
     if (!this.keyPair) {
       this.keyPair = generateKeyPair();
     }
@@ -119,7 +116,7 @@ export default class ProvisioningCipher {
 
   decrypt: (
     provisionEnvelope: Proto.ProvisionEnvelope
-  ) => Promise<ProvisionDecryptResult>;
+  ) => ProvisionDecryptResult;
 
-  getPublicKey: () => Promise<Uint8Array>;
+  getPublicKey: () => Uint8Array;
 }