fix(content/KeyAutoAdd): get output of any step, not just of previous

src/content/keyAutoAdd/lib/keyAutoAdd.ts

@@ -15,22 +15,22 @@ import type {
   KeyAutoAddToBackgroundMessage,
   KeyAutoAddToBackgroundMessagesMap,
   Step,
-  StepRunParams,
+  StepRun,
+  StepRunHelpers,
   StepWithStatus,
 } from './types';
 
 export type { StepRun } from './types';
 
 export const LOGIN_WAIT_TIMEOUT = 10 * 60 * 1000;
 
-const SYMBOL_SKIP = Symbol.for('skip');
-
 export class KeyAutoAdd {
   private port: Runtime.Port;
   private ui: HTMLIFrameElement;
 
   private stepsInput: Map<string, Step>;
   private steps: StepWithStatus[];
+  private outputs = new Map<StepRun, unknown>();
 
   constructor(steps: Step[]) {
     this.stepsInput = new Map(steps.map((step) => [step.name, step]));
@@ -43,7 +43,7 @@ export class KeyAutoAdd {
     this.port.onMessage.addListener(
       (message: BackgroundToKeyAutoAddMessage) => {
         if (message.action === 'BEGIN') {
-          this.run(message.payload);
+          this.runAll(message.payload);
         }
       },
     );
@@ -121,24 +121,19 @@ export class KeyAutoAdd {
     return promise;
   }
 
-  private async run({
-    walletAddressUrl,
-    publicKey,
-    nickName,
-    keyId,
-    keyAddUrl,
-  }: BeginPayload) {
-    const params: StepRunParams = {
-      walletAddressUrl,
-      publicKey,
-      nickName,
-      keyId,
-      keyAddUrl,
+  private async runAll(payload: BeginPayload) {
+    const helpers: StepRunHelpers = {
+      output: <T extends StepRun>(fn: T) => {
+        if (!this.outputs.has(fn)) {
+          // Was never run? Was skipped?
+          throw new Error('Given step has no output');
+        }
+        return this.outputs.get(fn) as Awaited<ReturnType<T>>;
+      },
       skip: (details) => {
-        throw {
-          type: SYMBOL_SKIP,
-          details: typeof details === 'string' ? new Error(details) : details,
-        };
+        throw new SkipError(
+          typeof details === 'string' ? { message: details } : details,
+        );
       },
       setNotificationSize: (size: 'notification' | 'fullscreen') => {
         this.setNotificationSize(size);
@@ -148,8 +143,6 @@ export class KeyAutoAdd {
     await this.addNotification();
     this.postMessage('PROGRESS', { steps: this.steps });
 
-    let prevStepId = '';
-    let prevStepResult: unknown = undefined;
     for (let stepIdx = 0; stepIdx < this.steps.length; stepIdx++) {
       const step = this.steps[stepIdx];
       const stepInfo = this.stepsInput.get(step.name)!;
@@ -159,20 +152,17 @@ export class KeyAutoAdd {
           : undefined,
       });
       try {
-        prevStepResult = await this.stepsInput
-          .get(step.name)!
-          .run(params, prevStepId ? prevStepResult : null);
+        const run = this.stepsInput.get(step.name)!.run;
+        const res = await run(payload, helpers);
+        this.outputs.set(run, res);
         this.setStatus(stepIdx, 'success', {});
-        prevStepId = step.name;
       } catch (error) {
-        if (this.isSkip(error)) {
-          const details = this.errorToDetails(
-            error.details.error || error.details,
-          );
+        if (error instanceof SkipError) {
+          const details = error.toJSON();
           this.setStatus(stepIdx, 'skipped', { details });
           continue;
         }
-        const details = this.errorToDetails(error);
+        const details = errorToDetails(error);
         this.setStatus(stepIdx, 'error', { details: details });
         this.postMessage('ERROR', { details, stepName: step.name, stepIdx });
         this.port.disconnect();
@@ -205,15 +195,23 @@ export class KeyAutoAdd {
     };
     this.postMessage('PROGRESS', { steps: this.steps });
   }
+}
 
-  private isSkip(err: unknown): err is { type: symbol; details: Details } {
-    if (!err || typeof err !== 'object') return false;
-    return 'type' in err && err.type === SYMBOL_SKIP;
+class SkipError extends Error {
+  public readonly error?: ErrorWithKeyLike;
+  constructor(err: ErrorWithKeyLike | { message: string }) {
+    const { message, error } = errorToDetails(err);
+    super(message);
+    this.error = error;
   }
 
-  private errorToDetails(err: { message: string } | ErrorWithKeyLike) {
-    return isErrorWithKey(err)
-      ? { error: errorWithKeyToJSON(err), message: err.key }
-      : { message: err.message as string };
+  toJSON(): Details {
+    return { message: this.message, error: this.error };
   }
 }
+
+function errorToDetails(err: { message: string } | ErrorWithKeyLike): Details {
+  return isErrorWithKey(err)
+    ? { error: errorWithKeyToJSON(err), message: err.key }
+    : { message: err.message as string };
+}

src/content/keyAutoAdd/lib/types.ts

@@ -1,21 +1,20 @@
 import type { ErrorWithKeyLike } from '@/shared/helpers';
 import type { ErrorResponse } from '@/shared/messages';
 
-export interface StepRunParams extends BeginPayload {
+export interface StepRunHelpers {
   skip: (message: string | Error | ErrorWithKeyLike) => never;
   setNotificationSize: (size: 'notification' | 'fullscreen') => void;
+  output: <T extends StepRun>(fn: T) => Awaited<ReturnType<T>>;
 }
 
-export type StepRun<T = unknown, R = void> = (
-  params: StepRunParams,
-  prevStepResult: T extends (...args: any[]) => PromiseLike<any>
-    ? Exclude<Awaited<ReturnType<T>>, void | { type: symbol }>
-    : T,
-) => Promise<R | void>;
+export type StepRun<TOutput = unknown> = (
+  payload: BeginPayload,
+  helpers: StepRunHelpers,
+) => Promise<TOutput>;
 
-export interface Step<T = unknown, R = unknown> {
+export interface Step<TOutput = unknown> {
   name: string;
-  run: StepRun<T, R>;
+  run: StepRun<TOutput>;
   maxDuration?: number;
 }
 

src/content/keyAutoAdd/testWallet.ts

@@ -3,7 +3,7 @@ import { errorWithKey, ErrorWithKey, sleep } from '@/shared/helpers';
 import {
   KeyAutoAdd,
   LOGIN_WAIT_TIMEOUT,
-  type StepRun as Step,
+  type StepRun as Run,
 } from './lib/keyAutoAdd';
 import { isTimedOut, waitForURL } from './lib/helpers';
 import { toWalletAddressUrl } from '@/popup/lib/utils';
@@ -26,11 +26,10 @@ type AccountDetails = {
   };
 };
 
-const waitForLogin: Step<never> = async ({
-  skip,
-  setNotificationSize,
-  keyAddUrl,
-}) => {
+const waitForLogin: Run<void> = async (
+  { keyAddUrl },
+  { skip, setNotificationSize },
+) => {
   let alreadyLoggedIn = window.location.href.startsWith(keyAddUrl);
   if (!alreadyLoggedIn) setNotificationSize('notification');
   try {
@@ -51,9 +50,7 @@ const waitForLogin: Step<never> = async ({
   }
 };
 
-const getAccountDetails: Step<never, Account[]> = async ({
-  setNotificationSize,
-}) => {
+const getAccounts: Run<Account[]> = async (_, { setNotificationSize }) => {
   setNotificationSize('fullscreen');
   await sleep(1000);
 
@@ -87,17 +84,22 @@ const getAccountDetails: Step<never, Account[]> = async ({
  * The test wallet associates key with an account. If the same key is associated
  * with a different account (user disconnected and changed account), revoke from
  * there first.
+ *
+ * Why? Say, user connected once to `USD -> Addr#1`. Then disconnected. The key
+ * is still there in wallet added to `USD -> Addr#1` account. If now user wants
+ * to connect `EUR -> Addr#2` account, the extension still has the same key. So
+ * adding it again will throw an `internal server error`. But we'll continue
+ * getting `invalid_client` if we try to connect without the key added to new
+ * address. That's why we first revoke existing key (by ID) if any (from any
+ * existing account/address). It's a test-wallet specific thing.
  */
-const revokeExistingKey: Step<typeof getAccountDetails, Account[]> = async (
-  { keyId, skip },
-  accounts,
-) => {
+const revokeExistingKey: Run<void> = async ({ keyId }, { skip, output }) => {
+  const accounts = output(getAccounts);
   for (const account of accounts) {
     for (const wallet of account.walletAddresses) {
       for (const key of wallet.keys) {
         if (key.id === keyId) {
           await revokeKey(account.id, wallet.id, key.id);
-          return accounts;
         }
       }
     }
@@ -106,10 +108,11 @@ const revokeExistingKey: Step<typeof getAccountDetails, Account[]> = async (
   skip('No existing keys that need to be revoked');
 };
 
-const findWallet: Step<
-  typeof revokeExistingKey,
-  { accountId: string; walletId: string }
-> = async ({ walletAddressUrl }, accounts) => {
+const findWallet: Run<{ accountId: string; walletId: string }> = async (
+  { walletAddressUrl },
+  { output },
+) => {
+  const accounts = output(getAccounts);
   for (const account of accounts) {
     for (const wallet of account.walletAddresses) {
       if (toWalletAddressUrl(wallet.url) === walletAddressUrl) {
@@ -120,10 +123,8 @@ const findWallet: Step<
   throw new ErrorWithKey('connectWalletKeyService_error_accountNotFound');
 };
 
-const addKey: Step<typeof findWallet> = async (
-  { publicKey, nickName },
-  { accountId, walletId },
-) => {
+const addKey: Run<void> = async ({ publicKey, nickName }, { output }) => {
+  const { accountId, walletId } = output(findWallet);
   const url = `https://api.rafiki.money/accounts/${accountId}/wallet-addresses/${walletId}/upload-key`;
   const res = await fetch(url, {
     method: 'POST',
@@ -169,7 +170,7 @@ new KeyAutoAdd([
     run: waitForLogin,
     maxDuration: LOGIN_WAIT_TIMEOUT,
   },
-  { name: 'Getting account details', run: getAccountDetails },
+  { name: 'Getting account details', run: getAccounts },
   { name: 'Revoking existing key', run: revokeExistingKey },
   { name: 'Finding wallet', run: findWallet },
   { name: 'Adding key', run: addKey },