fix: fixed revokeToken missing in vault

Author: lwin-kyaw

packages/seedless-onboarding-controller/src/SeedlessOnboardingController.test.ts

@@ -1385,14 +1385,28 @@ describe('SeedlessOnboardingController', () => {
             withMockAuthenticatedUser: true,
             withMockAuthPubKey: true,
             withoutMockRevokeToken: true,
+            vault: 'MOCK_VAULT',
           }),
         },
-        async ({ controller, toprfClient }) => {
+        async ({ controller, toprfClient, encryptor }) => {
           mockcreateLocalKey(toprfClient, MOCK_PASSWORD);
 
           // persist the local enc key
           jest.spyOn(toprfClient, 'persistLocalKey').mockResolvedValueOnce();
 
+          // since revoke token is not in the state, it should be retrieved from the vault
+          // mock the decryptWithDetail to return the mock vault data without revoke token
+          jest.spyOn(encryptor, 'decryptWithDetail').mockResolvedValue({
+            vault: JSON.stringify({
+              toprfEncryptionKey: '',
+              toprfPwEncryptionKey: '',
+              toprfAuthKeyPair: '',
+              accessToken,
+            }),
+            exportedKeyString: '',
+            salt: '',
+          });
+
           // encrypt and store the secret data
           handleMockSecretDataAdd();
           await expect(
@@ -3849,6 +3863,7 @@ describe('SeedlessOnboardingController', () => {
         initialPwEncKey,
         initialAuthKeyPair,
         OLD_PASSWORD,
+        revokeToken,
       );
 
       MOCK_VAULT = mockResult.encryptedMockVault;
@@ -5378,10 +5393,10 @@ describe('SeedlessOnboardingController', () => {
     });
   });
 
-  describe('#getAccessToken', () => {
+  describe('#getAccessTokenAndRevokeToken', () => {
     const MOCK_PASSWORD = 'mock-password';
 
-    it('should retrieve the access token from the vault if it is not available in the state', async () => {
+    it('should retrieve the access token and revoke token from the vault if it is not available in the state', async () => {
       const mockToprfEncryptor = createMockToprfEncryptor();
       const MOCK_ENCRYPTION_KEY =
         mockToprfEncryptor.deriveEncKey(MOCK_PASSWORD);
@@ -5456,12 +5471,28 @@ describe('SeedlessOnboardingController', () => {
           state: getMockInitialControllerState({
             withMockAuthenticatedUser: true,
             withoutMockAccessToken: true,
+            vault: 'MOCK_VAULT',
           }),
         },
-        async ({ controller, toprfClient }) => {
+        async ({ controller, toprfClient, encryptor }) => {
           // fetch and decrypt the secret data
           mockRecoverEncKey(toprfClient, MOCK_PASSWORD);
-          // mock the incorrect data shape
+
+          // since access token is not in the state, it should be retrieved from the vault
+          // mock the decryptWithDetail to return the mock vault data without access token
+          const decryptWithDetailMock = jest
+            .spyOn(encryptor, 'decryptWithDetail')
+            .mockResolvedValue({
+              vault: JSON.stringify({
+                toprfEncryptionKey: '',
+                toprfPwEncryptionKey: '',
+                toprfAuthKeyPair: '',
+                revokeToken: 'MOCK_REVOKE_TOKEN',
+              }),
+              exportedKeyString: '',
+              salt: '',
+            });
+
           jest
             .spyOn(toprfClient, 'fetchAllSecretDataItems')
             .mockResolvedValueOnce([
@@ -5480,6 +5511,10 @@ describe('SeedlessOnboardingController', () => {
           ).rejects.toThrow(
             SeedlessOnboardingControllerErrorMessage.InvalidAccessToken,
           );
+          expect(decryptWithDetailMock).toHaveBeenCalledWith(
+            MOCK_PASSWORD,
+            'MOCK_VAULT',
+          );
         },
       );
     });

packages/seedless-onboarding-controller/src/SeedlessOnboardingController.ts

@@ -961,35 +961,6 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     return isAuthenticated;
   }
 
-  /**
-   * Get the access token from the state or the vault.
-   * If the access token is not in the state, it will be retrieved from the vault by decrypting it with the password.
-   *
-   * If both the access token and the vault are not available, an error will be thrown.
-   *
-   * @param password - The optional password to unlock the vault. If not provided, the access token will be retrieved from the vault.
-   * @returns The access token.
-   */
-  async #getAccessToken(password: string): Promise<string> {
-    const { accessToken, vault } = this.state;
-    if (accessToken) {
-      // if the access token is in the state, return it
-      return accessToken;
-    }
-
-    // otherwise, check the vault availability and decrypt the access token from the vault
-    if (!vault) {
-      throw new Error(
-        SeedlessOnboardingControllerErrorMessage.InvalidAccessToken,
-      );
-    }
-
-    const { vaultData } = await this.#decryptAndParseVaultData({
-      password,
-    });
-    return vaultData.accessToken;
-  }
-
   #setUnlocked(): void {
     this.#isUnlocked = true;
   }
@@ -1558,36 +1529,46 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     rawToprfPwEncryptionKey: Uint8Array;
     rawToprfAuthKeyPair: KeyPair;
   }): Promise<void> {
-    this.#assertIsAuthenticatedUser(this.state);
+    try {
+      this.#assertIsAuthenticatedUser(this.state);
 
-    const { revokeToken } = this.state;
-    if (!revokeToken) {
-      throw new Error(
-        SeedlessOnboardingControllerErrorMessage.InvalidRevokeToken,
-      );
-    }
-    const accessToken = await this.#getAccessToken(password);
-
-    const vaultData: DeserializedVaultData = {
-      toprfAuthKeyPair: rawToprfAuthKeyPair,
-      toprfEncryptionKey: rawToprfEncryptionKey,
-      toprfPwEncryptionKey: rawToprfPwEncryptionKey,
-      revokeToken,
-      accessToken,
-    };
+      let { accessToken, revokeToken } = this.state;
+      if (!accessToken || !revokeToken) {
+        ({
+          vaultData: { accessToken, revokeToken },
+        } = await this.#decryptAndParseVaultData({
+          password,
+        }));
+      }
 
-    await this.#updateVault({
-      password,
-      vaultData,
-      pwEncKey: rawToprfPwEncryptionKey,
-    });
+      const vaultData: DeserializedVaultData = {
+        toprfAuthKeyPair: rawToprfAuthKeyPair,
+        toprfEncryptionKey: rawToprfEncryptionKey,
+        toprfPwEncryptionKey: rawToprfPwEncryptionKey,
+        revokeToken,
+        accessToken,
+      };
 
-    // update the authPubKey in the state
-    this.#persistAuthPubKey({
-      authPubKey: rawToprfAuthKeyPair.pk,
-    });
+      await this.#updateVault({
+        password,
+        vaultData,
+        pwEncKey: rawToprfPwEncryptionKey,
+      });
 
-    this.#setUnlocked();
+      // update the authPubKey in the state
+      this.#persistAuthPubKey({
+        authPubKey: rawToprfAuthKeyPair.pk,
+      });
+
+      this.#setUnlocked();
+    } catch (error) {
+      log(
+        'Error creating new vault with auth data',
+        error,
+        JSON.stringify(this.state),
+      );
+      throw error;
+    }
   }
 
   /**
@@ -1860,11 +1841,6 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
         password,
         encryptionKey: vaultEncryptionKey,
       });
-      if (!revokeToken) {
-        throw new Error(
-          SeedlessOnboardingControllerErrorMessage.InvalidRevokeToken,
-        );
-      }
 
       const { newRevokeToken, newRefreshToken } = await this.#renewRefreshToken(
         {

packages/seedless-onboarding-controller/src/assertions.test.ts

@@ -81,51 +81,23 @@ describe('assertIsValidVaultData', () => {
       }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
     });
 
-    it('should throw when revokeToken exists but is not a string or undefined', () => {
-      const invalidData = {
-        ...createValidVaultData(),
-        revokeToken: 789,
-      };
-
-      expect(() => {
-        assertIsValidVaultData(invalidData);
-      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
-
-      const invalidData2 = {
-        ...createValidVaultData(),
-        revokeToken: null,
-      };
-
-      expect(() => {
-        assertIsValidVaultData(invalidData2);
-      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
-
-      const invalidData3 = {
-        ...createValidVaultData(),
-        revokeToken: {},
-      };
-
-      expect(() => {
-        assertIsValidVaultData(invalidData3);
-      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
-    });
-
     it('should throw when accessToken is missing or not a string', () => {
       const invalidData = createValidVaultData();
       delete (invalidData as Record<string, unknown>).accessToken;
 
       expect(() => {
         assertIsValidVaultData(invalidData);
-      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
+      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidAccessToken);
 
       const invalidData2 = {
         ...createValidVaultData(),
+        revokeToken: 'MOCK_REVOKE_TOKEN',
         accessToken: 999,
       };
 
       expect(() => {
         assertIsValidVaultData(invalidData2);
-      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
+      }).toThrow(SeedlessOnboardingControllerErrorMessage.InvalidAccessToken);
     });
   });
 
@@ -138,50 +110,6 @@ describe('assertIsValidVaultData', () => {
       }).not.toThrow();
     });
 
-    it('should not throw when revokeToken is undefined', () => {
-      const validData = {
-        ...createValidVaultData(),
-        revokeToken: undefined,
-      };
-
-      expect(() => {
-        assertIsValidVaultData(validData);
-      }).not.toThrow();
-    });
-
-    it('should not throw when revokeToken is a valid string', () => {
-      const validData = {
-        ...createValidVaultData(),
-        revokeToken: 'valid_revoke_token',
-      };
-
-      expect(() => {
-        assertIsValidVaultData(validData);
-      }).not.toThrow();
-    });
-
-    it('should not throw when revokeToken property is missing entirely', () => {
-      const validData = createValidVaultData();
-      delete (validData as Record<string, unknown>).revokeToken;
-
-      expect(() => {
-        assertIsValidVaultData(validData);
-      }).not.toThrow();
-    });
-
-    it('should not throw with minimal valid vault data', () => {
-      const minimalValidData = {
-        toprfEncryptionKey: 'key1',
-        toprfPwEncryptionKey: 'key2',
-        toprfAuthKeyPair: 'keyPair',
-        accessToken: 'token',
-      };
-
-      expect(() => {
-        assertIsValidVaultData(minimalValidData);
-      }).not.toThrow();
-    });
-
     it('should not throw with extra properties in valid vault data', () => {
       const validDataWithExtras = {
         ...createValidVaultData(),

packages/seedless-onboarding-controller/src/assertions.ts

@@ -89,14 +89,20 @@ export function assertIsValidVaultData(
     !('toprfPwEncryptionKey' in value) || // toprfPwEncryptionKey is not defined
     typeof value.toprfPwEncryptionKey !== 'string' || // toprfPwEncryptionKey is not a string
     !('toprfAuthKeyPair' in value) || // toprfAuthKeyPair is not defined
-    typeof value.toprfAuthKeyPair !== 'string' || // toprfAuthKeyPair is not a string
-    // revoke token exists but is not a string and is not undefined
-    ('revokeToken' in value &&
-      typeof value.revokeToken !== 'string' &&
-      value.revokeToken !== undefined) ||
-    !('accessToken' in value) || // accessToken is not defined
-    typeof value.accessToken !== 'string' // accessToken is not a string
+    typeof value.toprfAuthKeyPair !== 'string' // toprfAuthKeyPair is not a string
   ) {
     throw new Error(SeedlessOnboardingControllerErrorMessage.InvalidVaultData);
   }
+
+  if (!('revokeToken' in value) || typeof value.revokeToken !== 'string') {
+    throw new Error(
+      SeedlessOnboardingControllerErrorMessage.InvalidRevokeToken,
+    );
+  }
+
+  if (!('accessToken' in value) || typeof value.accessToken !== 'string') {
+    throw new Error(
+      SeedlessOnboardingControllerErrorMessage.InvalidAccessToken,
+    );
+  }
 }

packages/seedless-onboarding-controller/src/types.ts

@@ -359,7 +359,7 @@ export type VaultData = {
    * The revoke token to revoke refresh token and get new refresh token and new revoke token.
    * The revoke token may no longer be available after a large number of password changes. In this case, re-authentication is advised.
    */
-  revokeToken?: string;
+  revokeToken: string;
   /**
    * The access token used for pairing with profile sync auth service and to access other services.
    */