fix: fixed invalid-revoke-token issue

Author: lwin-kyaw

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

@@ -1288,6 +1288,7 @@ describe('SeedlessOnboardingController', () => {
             pwEncKey,
             authKeyPair,
             MOCK_PASSWORD,
+            controller.state.revokeToken,
           );
 
           const expectedVaultValue = await encryptor.decrypt(
@@ -1299,7 +1300,7 @@ describe('SeedlessOnboardingController', () => {
             controller.state.vault as string,
           );
 
-          expect(expectedVaultValue).toStrictEqual(resultedVaultValue);
+          expect(resultedVaultValue).toStrictEqual(expectedVaultValue);
 
           // should be able to get the hash of the seed phrase backup from the state
           expect(
@@ -1377,6 +1378,36 @@ describe('SeedlessOnboardingController', () => {
       );
     });
 
+    it('should throw error if revokeToken is missing when creating new vault', async () => {
+      await withController(
+        {
+          state: getMockInitialControllerState({
+            withMockAuthenticatedUser: true,
+            withMockAuthPubKey: true,
+            withoutMockRevokeToken: true,
+          }),
+        },
+        async ({ controller, toprfClient }) => {
+          mockcreateLocalKey(toprfClient, MOCK_PASSWORD);
+
+          // persist the local enc key
+          jest.spyOn(toprfClient, 'persistLocalKey').mockResolvedValueOnce();
+
+          // encrypt and store the secret data
+          handleMockSecretDataAdd();
+          await expect(
+            controller.createToprfKeyAndBackupSeedPhrase(
+              MOCK_PASSWORD,
+              MOCK_SEED_PHRASE,
+              MOCK_KEYRING_ID,
+            ),
+          ).rejects.toThrow(
+            SeedlessOnboardingControllerErrorMessage.InvalidRevokeToken,
+          );
+        },
+      );
+    });
+
     it('should throw an error if create encryption key fails', async () => {
       await withController(
         {
@@ -2084,27 +2115,6 @@ describe('SeedlessOnboardingController', () => {
       );
     });
 
-    it('should throw an error if the user is not authenticated', async () => {
-      await withController(
-        {
-          state: {
-            userId,
-            authConnectionId,
-            groupedAuthConnectionId,
-            metadataAccessToken,
-            refreshToken,
-            revokeToken,
-            nodeAuthTokens: MOCK_NODE_AUTH_TOKENS,
-          },
-        },
-        async ({ controller }) => {
-          await expect(controller.fetchAllSecretData()).rejects.toThrow(
-            SeedlessOnboardingControllerErrorMessage.InvalidAccessToken,
-          );
-        },
-      );
-    });
-
     it('should be able to fetch seed phrases with cached encryption key without providing password', async () => {
       const mockToprfEncryptor = createMockToprfEncryptor();
 
@@ -5331,6 +5341,24 @@ describe('SeedlessOnboardingController', () => {
       });
     });
 
+    it('should return true if access token is missing', async () => {
+      await withController(
+        {
+          state: getMockInitialControllerState({
+            withMockAuthenticatedUser: true,
+            withoutMockAccessToken: true,
+          }),
+        },
+        async ({ controller }) => {
+          // Restore the original implementation to test the real logic
+          jest.spyOn(controller, 'checkAccessTokenExpired').mockRestore();
+
+          const result = controller.checkAccessTokenExpired();
+          expect(result).toBe(true);
+        },
+      );
+    });
+
     it('should return true if token has invalid format', async () => {
       await withController(
         {
@@ -5350,6 +5378,113 @@ describe('SeedlessOnboardingController', () => {
     });
   });
 
+  describe('#getAccessToken', () => {
+    const MOCK_PASSWORD = 'mock-password';
+
+    it('should retrieve the access token from the vault if it is not available in the state', async () => {
+      const mockToprfEncryptor = createMockToprfEncryptor();
+      const MOCK_ENCRYPTION_KEY =
+        mockToprfEncryptor.deriveEncKey(MOCK_PASSWORD);
+      const MOCK_PASSWORD_ENCRYPTION_KEY =
+        mockToprfEncryptor.derivePwEncKey(MOCK_PASSWORD);
+      const MOCK_AUTH_KEY_PAIR =
+        mockToprfEncryptor.deriveAuthKeyPair(MOCK_PASSWORD);
+
+      const mockResult = await createMockVault(
+        MOCK_ENCRYPTION_KEY,
+        MOCK_PASSWORD_ENCRYPTION_KEY,
+        MOCK_AUTH_KEY_PAIR,
+        MOCK_PASSWORD,
+      );
+
+      const MOCK_VAULT = mockResult.encryptedMockVault;
+      const MOCK_VAULT_ENCRYPTION_KEY = mockResult.vaultEncryptionKey;
+      const MOCK_VAULT_ENCRYPTION_SALT = mockResult.vaultEncryptionSalt;
+
+      await withController(
+        {
+          state: getMockInitialControllerState({
+            withMockAuthenticatedUser: true,
+            withoutMockAccessToken: true,
+            vault: MOCK_VAULT,
+            vaultEncryptionKey: MOCK_VAULT_ENCRYPTION_KEY,
+            vaultEncryptionSalt: MOCK_VAULT_ENCRYPTION_SALT,
+          }),
+        },
+        async ({ controller, toprfClient }) => {
+          // fetch and decrypt the secret data
+          mockRecoverEncKey(toprfClient, MOCK_PASSWORD);
+
+          // mock the secret data get
+          jest
+            .spyOn(toprfClient, 'fetchAllSecretDataItems')
+            .mockResolvedValueOnce([
+              stringToBytes(
+                JSON.stringify({
+                  data: bytesToBase64(MOCK_SEED_PHRASE),
+                  timestamp: 1234567890,
+                  type: SecretType.Mnemonic,
+                  version: 'v1',
+                }),
+              ),
+              stringToBytes(
+                JSON.stringify({
+                  data: bytesToBase64(MOCK_PRIVATE_KEY),
+                  timestamp: 1234567890,
+                  type: SecretType.PrivateKey,
+                  version: 'v1',
+                }),
+              ),
+            ]);
+
+          const secretData = await controller.fetchAllSecretData(MOCK_PASSWORD);
+          expect(secretData).toBeDefined();
+          expect(secretData).toHaveLength(2);
+          expect(secretData[0].type).toStrictEqual(SecretType.Mnemonic);
+          expect(secretData[0].data).toStrictEqual(MOCK_SEED_PHRASE);
+          expect(secretData[1].type).toStrictEqual(SecretType.PrivateKey);
+          expect(secretData[1].data).toStrictEqual(MOCK_PRIVATE_KEY);
+
+          // expect(mockSecretDataGet.isDone()).toBe(true);
+        },
+      );
+    });
+
+    it('should throw error if access token is not available either in the state or the vault', async () => {
+      await withController(
+        {
+          state: getMockInitialControllerState({
+            withMockAuthenticatedUser: true,
+            withoutMockAccessToken: true,
+          }),
+        },
+        async ({ controller, toprfClient }) => {
+          // fetch and decrypt the secret data
+          mockRecoverEncKey(toprfClient, MOCK_PASSWORD);
+          // mock the incorrect data shape
+          jest
+            .spyOn(toprfClient, 'fetchAllSecretDataItems')
+            .mockResolvedValueOnce([
+              stringToBytes(
+                JSON.stringify({
+                  data: 'value',
+                  timestamp: 1234567890,
+                  type: 'mnemonic',
+                  version: 'v1',
+                }),
+              ),
+            ]);
+
+          await expect(
+            controller.fetchAllSecretData(MOCK_PASSWORD),
+          ).rejects.toThrow(
+            SeedlessOnboardingControllerErrorMessage.InvalidAccessToken,
+          );
+        },
+      );
+    });
+  });
+
   describe('renewRefreshToken', () => {
     const MOCK_PASSWORD = 'mock-password';
     const MOCK_REVOKE_TOKEN = 'newRevokeToken';

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

@@ -22,7 +22,6 @@ import { secp256k1 } from '@noble/curves/secp256k1';
 import { Mutex } from 'async-mutex';
 
 import {
-  assertIsAuthUserInfoValid,
   assertIsPasswordOutdatedCacheValid,
   assertIsSeedlessOnboardingUserAuthenticated,
   assertIsValidVaultData,
@@ -365,7 +364,7 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     groupedAuthConnectionId?: string;
     socialLoginEmail?: string;
     refreshToken: string;
-    revokeToken: string;
+    revokeToken?: string;
     skipLock?: boolean;
   }) {
     const doAuthenticateWithNodes = async () => {
@@ -399,8 +398,10 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
           state.socialLoginEmail = socialLoginEmail;
           state.metadataAccessToken = metadataAccessToken;
           state.refreshToken = refreshToken;
-          // Temporarily store revoke token & access token in state for later vault creation
-          state.revokeToken = revokeToken;
+          if (revokeToken) {
+            // Temporarily store revoke token & access token in state for later vault creation
+            state.revokeToken = revokeToken;
+          }
           state.accessToken = accessToken;
 
           // we will check if the controller state is properly set with the authenticated user info
@@ -892,7 +893,7 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
         }
       }
 
-      assertIsAuthUserInfoValid(this.state);
+      this.#assertIsAuthenticatedUser(this.state);
       const {
         nodeAuthTokens,
         authConnectionId,
@@ -948,8 +949,9 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
   async checkIsSeedlessOnboardingUserAuthenticated(): Promise<boolean> {
     let isAuthenticated = false;
     try {
-      assertIsSeedlessOnboardingUserAuthenticated(this.state);
-      isAuthenticated = true;
+      this.#assertIsAuthenticatedUser(this.state);
+      isAuthenticated =
+        Boolean(this.state.accessToken) && Boolean(this.state.refreshToken);
     } catch {
       isAuthenticated = false;
     }
@@ -959,6 +961,35 @@ 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;
   }
@@ -1116,7 +1147,7 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
   async #recoverEncKey(
     password: string,
   ): Promise<Omit<RecoverEncryptionKeyResult, 'rateLimitResetResult'>> {
-    assertIsAuthUserInfoValid(this.state);
+    this.#assertIsAuthenticatedUser(this.state);
     const {
       nodeAuthTokens,
       authConnectionId,
@@ -1529,7 +1560,13 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
   }): Promise<void> {
     this.#assertIsAuthenticatedUser(this.state);
 
-    const { revokeToken, accessToken } = this.state;
+    const { revokeToken } = this.state;
+    if (!revokeToken) {
+      throw new Error(
+        SeedlessOnboardingControllerErrorMessage.InvalidRevokeToken,
+      );
+    }
+    const accessToken = await this.#getAccessToken(password);
 
     const vaultData: DeserializedVaultData = {
       toprfAuthKeyPair: rawToprfAuthKeyPair,
@@ -1768,7 +1805,7 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
    */
   async refreshAuthTokens(): Promise<void> {
     this.#assertIsAuthenticatedUser(this.state);
-    const { refreshToken, revokeToken } = this.state;
+    const { refreshToken } = this.state;
 
     const res = await this.#refreshJWTToken({
       connection: this.state.authConnection,
@@ -1792,7 +1829,6 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
         groupedAuthConnectionId: this.state.groupedAuthConnectionId,
         userId: this.state.userId,
         refreshToken,
-        revokeToken,
         skipLock: true,
       });
     } catch (error) {
@@ -1990,7 +2026,6 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
       const isNodeAuthTokenExpired = this.checkNodeAuthTokenExpired();
       const isMetadataAccessTokenExpired =
         this.checkMetadataAccessTokenExpired();
-
       // access token is only accessible when the vault is unlocked
       // so skip the check if the vault is locked
       let isAccessTokenExpired = false;
@@ -2078,6 +2113,9 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     try {
       this.#assertIsAuthenticatedUser(this.state);
       const { accessToken } = this.state;
+      if (!accessToken) {
+        return true; // Consider missing token as expired
+      }
       const decodedToken = decodeJWTToken(accessToken);
       return decodedToken.exp < Math.floor(Date.now() / 1000);
     } catch {