Refactor provisioning flow

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);
+    }
+  }
+}