feat: persist UserStorage e2e content keys using an encrypted keyStore

fixes #5128

Author: mirceanis

.gitignore

@@ -35,3 +35,6 @@ scripts/coverage
 
 # typescript
 packages/*/*.tsbuildinfo
+
+# jetbrains IDE local files
+/.idea

packages/profile-sync-controller/package.json

@@ -106,7 +106,9 @@
     "@metamask/network-controller": "^22.1.1",
     "@metamask/snaps-sdk": "^6.7.0",
     "@metamask/snaps-utils": "^8.3.0",
+    "@metamask/utils": "^11.0.1",
     "@noble/ciphers": "^0.5.2",
+    "@noble/curves": "^1.7.0",
     "@noble/hashes": "^1.4.0",
     "immer": "^9.0.6",
     "loglevel": "^1.8.1",

packages/profile-sync-controller/src/controllers/authentication/auth-snap-requests.ts

@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import type { HandleSnapRequest } from '@metamask/snaps-controllers';
 import type { SnapId } from '@metamask/snaps-sdk';
+import type { Eip1024EncryptedData } from '@metamask/utils';
 
 type SnapRPCRequest = Parameters<HandleSnapRequest['handler']>[0];
 
@@ -22,6 +23,22 @@ export function createSnapPublicKeyRequest(): SnapRPCRequest {
   };
 }
 
+/**
+ * Constructs Request to Message Signing Snap to get the Encryption Public Key
+ *
+ * @returns Snap Encryption Public Key Request
+ */
+export function createSnapEncryptionPublicKeyRequest(): SnapRPCRequest {
+  return {
+    snapId,
+    origin: '',
+    handler: 'onRpcRequest' as any,
+    request: {
+      method: 'getEncryptionPublicKey',
+    },
+  };
+}
+
 /**
  * Constructs Request to get Message Signing Snap to sign a message.
  *
@@ -41,3 +58,23 @@ export function createSnapSignMessageRequest(
     },
   };
 }
+
+/**
+ * Constructs Request to get Message Signing Snap to decrypt a message.
+ *
+ * @param data - message to decrypt
+ * @returns Snap Sign Message Request
+ */
+export function createSnapDecryptMessageRequest(
+  data: Eip1024EncryptedData,
+): SnapRPCRequest {
+  return {
+    snapId,
+    origin: '',
+    handler: 'onRpcRequest' as any,
+    request: {
+      method: 'decryptMessage',
+      params: { data },
+    },
+  };
+}

packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.test.ts

@@ -109,6 +109,7 @@ describe('user-storage/user-storage-controller - performGetStorage() tests', ()
         isAccountSyncingInProgress: false,
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -191,6 +192,7 @@ describe('user-storage/user-storage-controller - performGetStorageAllFeatureEntr
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -273,6 +275,7 @@ describe('user-storage/user-storage-controller - performSetStorage() tests', ()
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -379,6 +382,7 @@ describe('user-storage/user-storage-controller - performBatchSetStorage() tests'
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -482,6 +486,7 @@ describe('user-storage/user-storage-controller - performBatchDeleteStorage() tes
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -586,6 +591,7 @@ describe('user-storage/user-storage-controller - performDeleteStorage() tests',
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -687,6 +693,7 @@ describe('user-storage/user-storage-controller - performDeleteStorageAllFeatureE
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -780,6 +787,7 @@ describe('user-storage/user-storage-controller - getStorageKey() tests', () => {
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 
@@ -827,6 +835,7 @@ describe('user-storage/user-storage-controller - enableProfileSyncing() tests',
         hasAccountSyncingSyncedAtLeastOnce: false,
         isAccountSyncingReadyToBeDispatched: false,
         isAccountSyncingInProgress: false,
+        encryptedContentKeys: {},
       },
     });
 

packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts

@@ -1,8 +1,8 @@
 import type {
+  AccountsControllerAccountAddedEvent,
+  AccountsControllerAccountRenamedEvent,
   AccountsControllerListAccountsAction,
   AccountsControllerUpdateAccountMetadataAction,
-  AccountsControllerAccountRenamedEvent,
-  AccountsControllerAccountAddedEvent,
 } from '@metamask/accounts-controller';
 import type {
   ControllerGetStateAction,
@@ -12,10 +12,10 @@ import type {
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
 import {
+  type KeyringControllerAddNewAccountAction,
   type KeyringControllerGetStateAction,
   type KeyringControllerLockEvent,
   type KeyringControllerUnlockEvent,
-  type KeyringControllerAddNewAccountAction,
 } from '@metamask/keyring-controller';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import type {
@@ -26,22 +26,29 @@ import type {
   NetworkControllerUpdateNetworkAction,
 } from '@metamask/network-controller';
 import type { HandleSnapRequest } from '@metamask/snaps-controllers';
+import type { Eip1024EncryptedData } from '@metamask/utils';
 
 import { createSHA256Hash } from '../../shared/encryption';
+import type { KeyStore } from '../../shared/encryption/key-storage';
+import { ERC1024WrappedKeyStore } from '../../shared/encryption/key-storage';
 import type { UserStorageFeatureKeys } from '../../shared/storage-schema';
 import {
   type UserStoragePathWithFeatureAndKey,
   type UserStoragePathWithFeatureOnly,
 } from '../../shared/storage-schema';
 import type { NativeScrypt } from '../../shared/types/encryption';
-import { createSnapSignMessageRequest } from '../authentication/auth-snap-requests';
+import {
+  createSnapDecryptMessageRequest,
+  createSnapEncryptionPublicKeyRequest,
+  createSnapSignMessageRequest,
+} from '../authentication/auth-snap-requests';
 import type {
   AuthenticationControllerGetBearerToken,
   AuthenticationControllerGetSessionProfile,
   AuthenticationControllerIsSignedIn,
   AuthenticationControllerPerformSignIn,
   AuthenticationControllerPerformSignOut,
-} from '../authentication/AuthenticationController';
+} from '../authentication';
 import {
   saveInternalAccountToUserStorage,
   syncInternalAccountsWithUserStorage,
@@ -102,6 +109,10 @@ export type UserStorageControllerState = {
    * Condition used to ensure that we do not perform any network sync mutations until we have synced at least once
    */
   hasNetworkSyncingSyncedAtLeastOnce?: boolean;
+  /**
+   * Content keys used to encrypt/decrypt user storage content. These are wrapped while at rest.
+   */
+  encryptedContentKeys: Record<string, string>;
 };
 
 export const defaultState: UserStorageControllerState = {
@@ -110,6 +121,7 @@ export const defaultState: UserStorageControllerState = {
   hasAccountSyncingSyncedAtLeastOnce: false,
   isAccountSyncingReadyToBeDispatched: false,
   isAccountSyncingInProgress: false,
+  encryptedContentKeys: {},
 };
 
 const metadata: StateMetadata<UserStorageControllerState> = {
@@ -137,6 +149,10 @@ const metadata: StateMetadata<UserStorageControllerState> = {
     persist: true,
     anonymous: false,
   },
+  encryptedContentKeys: {
+    persist: true,
+    anonymous: false,
+  },
 };
 
 type ControllerConfig = {
@@ -313,6 +329,79 @@ export default class UserStorageController extends BaseController<
     isNetworkSyncingEnabled: false,
   };
 
+  #_snapPublicKeyCache: string | null = null;
+
+  #keyWrapping = {
+    /**
+     * Returns the snap Encryption public key.
+     *
+     * @returns The snap Encryption public key.
+     */
+    snapGetEncryptionPublicKey: async (): Promise<string> => {
+      if (this.#_snapPublicKeyCache) {
+        return this.#_snapPublicKeyCache;
+      }
+
+      if (!this.#isUnlocked) {
+        throw new Error(
+          '#snapGetEncryptionPublicKey - unable to call snap, wallet is locked',
+        );
+      }
+
+      const result = (await this.messagingSystem.call(
+        'SnapController:handleRequest',
+        createSnapEncryptionPublicKeyRequest(),
+      )) as string;
+
+      this.#_snapPublicKeyCache = result;
+
+      return result;
+    },
+
+    /**
+     * Decrypts a message using the message signing snap.
+     *
+     * @param data - Eip1024EncryptedData - The encrypted data.
+     * @returns The decrypted message, if it was intended for this wallet, null otherwise. TODO: check error scenarios
+     */
+    snapDecryptMessage: async (data: Eip1024EncryptedData): Promise<string> => {
+      if (!this.#isUnlocked) {
+        throw new Error(
+          '#snapDecryptMessage - unable to call snap, wallet is locked',
+        );
+      }
+
+      const result = (await this.messagingSystem.call(
+        'SnapController:handleRequest',
+        createSnapDecryptMessageRequest(data),
+      )) as string;
+
+      return result;
+    },
+
+    loadWrappedKey: async (keyRef: string): Promise<string | null> => {
+      return this.state.encryptedContentKeys[keyRef] ?? null;
+    },
+
+    storeWrappedKey: (keyRef: string, wrappedKey: string): Promise<void> => {
+      return new Promise((resolve) => {
+        this.update((state) => {
+          state.encryptedContentKeys[keyRef] = wrappedKey;
+          resolve();
+        });
+      });
+    },
+
+    getWrappedKeyStore: (): KeyStore => {
+      return new ERC1024WrappedKeyStore({
+        decryptMessage: this.#keyWrapping.snapDecryptMessage,
+        getPublicKey: this.#keyWrapping.snapGetEncryptionPublicKey,
+        getItem: this.#keyWrapping.loadWrappedKey,
+        setItem: this.#keyWrapping.storeWrappedKey,
+      });
+    },
+  };
+
   #auth = {
     getBearerToken: async () => {
       return await this.messagingSystem.call(
@@ -374,7 +463,9 @@ export default class UserStorageController extends BaseController<
     },
   };
 
-  #nativeScryptCrypto: NativeScrypt | undefined = undefined;
+  #nativeScryptCrypto: NativeScrypt | undefined;
+
+  #keyStore: KeyStore | undefined;
 
   getMetaMetricsState: () => boolean;
 
@@ -385,6 +476,7 @@ export default class UserStorageController extends BaseController<
     config,
     getMetaMetricsState,
     nativeScryptCrypto,
+    keyStore,
   }: {
     messenger: UserStorageControllerMessenger;
     state?: UserStorageControllerState;
@@ -395,6 +487,7 @@ export default class UserStorageController extends BaseController<
     };
     getMetaMetricsState: () => boolean;
     nativeScryptCrypto?: NativeScrypt;
+    keyStore?: KeyStore;
   }) {
     super({
       messenger,
@@ -432,6 +525,8 @@ export default class UserStorageController extends BaseController<
           !this.state.hasNetworkSyncingSyncedAtLeastOnce,
       });
     }
+
+    this.#keyStore = keyStore ?? this.#keyWrapping.getWrappedKeyStore();
   }
 
   /**
@@ -491,6 +586,7 @@ export default class UserStorageController extends BaseController<
       storageKey,
       bearerToken,
       nativeScryptCrypto: this.#nativeScryptCrypto,
+      keyStore: this.#keyStore,
     };
   }
 
@@ -581,6 +677,7 @@ export default class UserStorageController extends BaseController<
       bearerToken,
       storageKey,
       nativeScryptCrypto: this.#nativeScryptCrypto,
+      keyStore: this.#keyStore,
     });
 
     return result;
@@ -606,6 +703,7 @@ export default class UserStorageController extends BaseController<
       bearerToken,
       storageKey,
       nativeScryptCrypto: this.#nativeScryptCrypto,
+      keyStore: this.#keyStore,
     });
 
     return result;
@@ -633,6 +731,7 @@ export default class UserStorageController extends BaseController<
       bearerToken,
       storageKey,
       nativeScryptCrypto: this.#nativeScryptCrypto,
+      keyStore: this.#keyStore,
     });
   }
 
@@ -660,6 +759,7 @@ export default class UserStorageController extends BaseController<
       bearerToken,
       storageKey,
       nativeScryptCrypto: this.#nativeScryptCrypto,
+      keyStore: this.#keyStore,
     });
   }
 
@@ -730,6 +830,7 @@ export default class UserStorageController extends BaseController<
       bearerToken,
       storageKey,
       nativeScryptCrypto: this.#nativeScryptCrypto,
+      keyStore: this.#keyStore,
     });
   }