feat: notifications - user storage controller

app/scripts/controllers/authentication/authentication-controller.test.ts

@@ -1,7 +1,6 @@
 import { ControllerMessenger } from '@metamask/base-controller';
-import type { HandleSnapRequest } from '@metamask/snaps-controllers';
 import AuthenticationController, {
-  AuthenticationControllerMessenger,
+  AllowedActions,
   AuthenticationControllerState,
 } from './authentication-controller';
 import {
@@ -233,11 +232,11 @@ describe('authentication/authentication-controller - getSessionProfile() tests',
 });
 
 function createAuthenticationMessenger() {
-  const messenger = new ControllerMessenger<HandleSnapRequest, never>();
+  const messenger = new ControllerMessenger<AllowedActions, never>();
   return messenger.getRestricted({
     name: 'AuthenticationController',
     allowedActions: [`SnapController:handleRequest`],
-  }) as AuthenticationControllerMessenger;
+  });
 }
 
 function createMockAuthenticationMessenger() {
@@ -247,25 +246,29 @@ function createMockAuthenticationMessenger() {
   const mockSnapSignMessage = jest
     .fn()
     .mockResolvedValue('MOCK_SIGNED_MESSAGE');
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  mockCall.mockImplementation(((actionType: any, params: any) => {
-    if (
-      actionType === 'SnapController:handleRequest' &&
-      params?.request.method === 'getPublicKey'
-    ) {
-      return mockSnapGetPublicKey();
+
+  mockCall.mockImplementation((...args) => {
+    const [actionType, params] = args;
+    if (actionType === 'SnapController:handleRequest') {
+      if (params?.request.method === 'getPublicKey') {
+        return mockSnapGetPublicKey();
+      }
+
+      if (params?.request.method === 'signMessage') {
+        return mockSnapSignMessage();
+      }
+
+      throw new Error(
+        `MOCK_FAIL - unsupported SnapController:handleRequest call: ${params?.request.method}`,
+      );
     }
 
-    if (
-      actionType === 'SnapController:handleRequest' &&
-      params?.request.method === 'signMessage'
-    ) {
-      return mockSnapSignMessage();
+    function exhaustedMessengerMocks(action: never) {
+      throw new Error(`MOCK_FAIL - unsupported messenger call: ${action}`);
     }
 
-    return '';
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  }) as any);
+    return exhaustedMessengerMocks(actionType);
+  });
 
   return { messenger, mockSnapGetPublicKey, mockSnapSignMessage };
 }

app/scripts/controllers/authentication/authentication-controller.ts

@@ -63,7 +63,11 @@ type CreateActionsObj<T extends keyof AuthenticationController> = {
   };
 };
 type ActionsObj = CreateActionsObj<
-  'performSignIn' | 'performSignOut' | 'getBearerToken' | 'getSessionProfile'
+  | 'performSignIn'
+  | 'performSignOut'
+  | 'getBearerToken'
+  | 'getSessionProfile'
+  | 'isSignedIn'
 >;
 export type Actions = ActionsObj[keyof ActionsObj];
 export type AuthenticationControllerPerformSignIn = ActionsObj['performSignIn'];
@@ -73,9 +77,10 @@ export type AuthenticationControllerGetBearerToken =
   ActionsObj['getBearerToken'];
 export type AuthenticationControllerGetSessionProfile =
   ActionsObj['getSessionProfile'];
+export type AuthenticationControllerIsSignedIn = ActionsObj['isSignedIn'];
 
 // Allowed Actions
-type AllowedActions = HandleSnapRequest;
+export type AllowedActions = HandleSnapRequest;
 
 // Messenger
 export type AuthenticationControllerMessenger = RestrictedControllerMessenger<
@@ -108,6 +113,39 @@ export default class AuthenticationController extends BaseController<
       name: controllerName,
       state: { ...defaultState, ...state },
     });
+
+    this.#registerMessageHandlers();
+  }
+
+  /**
+   * Constructor helper for registering this controller's messaging system
+   * actions.
+   */
+  #registerMessageHandlers(): void {
+    this.messagingSystem.registerActionHandler(
+      'AuthenticationController:getBearerToken',
+      this.getBearerToken.bind(this),
+    );
+
+    this.messagingSystem.registerActionHandler(
+      'AuthenticationController:getSessionProfile',
+      this.getSessionProfile.bind(this),
+    );
+
+    this.messagingSystem.registerActionHandler(
+      'AuthenticationController:isSignedIn',
+      this.isSignedIn.bind(this),
+    );
+
+    this.messagingSystem.registerActionHandler(
+      'AuthenticationController:performSignIn',
+      this.performSignIn.bind(this),
+    );
+
+    this.messagingSystem.registerActionHandler(
+      'AuthenticationController:performSignOut',
+      this.performSignOut.bind(this),
+    );
   }
 
   public async performSignIn(): Promise<string> {
@@ -152,6 +190,10 @@ export default class AuthenticationController extends BaseController<
     return profile;
   }
 
+  public isSignedIn(): boolean {
+    return this.state.isSignedIn;
+  }
+
   #assertLoggedIn(): void {
     if (!this.state.isSignedIn) {
       throw new Error(

app/scripts/controllers/user-storage/encryption.test.ts

@@ -0,0 +1,38 @@
+import encryption, { createSHA256Hash } from './encryption';
+
+describe('encryption tests', () => {
+  const PASSWORD = '123';
+  const DATA1 = 'Hello World';
+  const DATA2 = JSON.stringify({ foo: 'bar' });
+
+  it('Should encrypt and decrypt data', () => {
+    function actEncryptDecrypt(data: string) {
+      const encryptedString = encryption.encryptString(data, PASSWORD);
+      const decryptString = encryption.decryptString(encryptedString, PASSWORD);
+      return decryptString;
+    }
+
+    expect(actEncryptDecrypt(DATA1)).toBe(DATA1);
+
+    expect(actEncryptDecrypt(DATA2)).toBe(DATA2);
+  });
+
+  it('Should decrypt some existing data', () => {
+    const encryptedData = `{"v":"1","d":"R+sCbzS6clo5iLbSzBr889miNfHhCBmOCk2CFwTH55IkbOIL9f5Nm2t0nmWOVtFbjLpnj6cKyw==","iterations":900000}`;
+    const result = encryption.decryptString(encryptedData, PASSWORD);
+    expect(result).toBe(DATA1);
+  });
+
+  it('Should sha-256 hash a value and should be deterministic', () => {
+    const DATA = 'Hello World';
+    const EXPECTED_HASH =
+      'a591a6d40bf420404a011733cfb7b190d62c65bf0bcda32b57b277d9ad9f146e';
+
+    const hash1 = createSHA256Hash(DATA);
+    expect(hash1).toBe(EXPECTED_HASH);
+
+    // Hash should be deterministic (same output with same input)
+    const hash2 = createSHA256Hash(DATA);
+    expect(hash1).toBe(hash2);
+  });
+});

app/scripts/controllers/user-storage/encryption.ts

@@ -0,0 +1,138 @@
+import { pbkdf2 } from '@noble/hashes/pbkdf2';
+import { sha256 } from '@noble/hashes/sha256';
+import { utf8ToBytes, concatBytes, bytesToHex } from '@noble/hashes/utils';
+import { gcm } from '@noble/ciphers/aes';
+import { randomBytes } from '@noble/ciphers/webcrypto';
+
+export type EncryptedPayload = {
+  v: '1'; // version
+  d: string; // data
+  iterations: number;
+};
+
+function byteArrayToBase64(byteArray: Uint8Array) {
+  return Buffer.from(byteArray).toString('base64');
+}
+
+function base64ToByteArray(base64: string) {
+  return new Uint8Array(Buffer.from(base64, 'base64'));
+}
+
+function bytesToUtf8(byteArray: Uint8Array) {
+  const decoder = new TextDecoder('utf-8');
+  return decoder.decode(byteArray);
+}
+
+class EncryptorDecryptor {
+  #ALGORITHM_NONCE_SIZE: number = 12; // 12 bytes
+
+  #ALGORITHM_KEY_SIZE: number = 16; // 16 bytes
+
+  #PBKDF2_SALT_SIZE: number = 16; // 16 bytes
+
+  #PBKDF2_ITERATIONS: number = 900_000;
+
+  encryptString(plaintext: string, password: string): string {
+    try {
+      return this.#encryptStringV1(plaintext, password);
+    } catch (e) {
+      const errorMessage = e instanceof Error ? e.message : e;
+      throw new Error(`Unable to encrypt string - ${errorMessage}`);
+    }
+  }
+
+  decryptString(encryptedDataStr: string, password: string): string {
+    try {
+      const encryptedData: EncryptedPayload = JSON.parse(encryptedDataStr);
+      if (encryptedData.v === '1') {
+        return this.#decryptStringV1(encryptedData, password);
+      }
+      throw new Error(`Unsupported encrypted data payload - ${encryptedData}`);
+    } catch (e) {
+      const errorMessage = e instanceof Error ? e.message : e;
+      throw new Error(`Unable to decrypt string - ${errorMessage}`);
+    }
+  }
+
+  #encryptStringV1(plaintext: string, password: string): string {
+    const salt = randomBytes(this.#PBKDF2_SALT_SIZE);
+
+    // Derive a key using PBKDF2.
+    const key = pbkdf2(sha256, password, salt, {
+      c: this.#PBKDF2_ITERATIONS,
+      dkLen: this.#ALGORITHM_KEY_SIZE,
+    });
+
+    // Encrypt and prepend salt.
+    const plaintextRaw = utf8ToBytes(plaintext);
+    const ciphertextAndNonceAndSalt = concatBytes(
+      salt,
+      this.#encrypt(plaintextRaw, key),
+    );
+
+    // Convert to Base64
+    const encryptedData = byteArrayToBase64(ciphertextAndNonceAndSalt);
+
+    const encryptedPayload: EncryptedPayload = {
+      v: '1',
+      d: encryptedData,
+      iterations: this.#PBKDF2_ITERATIONS,
+    };
+
+    return JSON.stringify(encryptedPayload);
+  }
+
+  #decryptStringV1(data: EncryptedPayload, password: string): string {
+    const { iterations, d: base64CiphertextAndNonceAndSalt } = data;
+
+    // Decode the base64.
+    const ciphertextAndNonceAndSalt = base64ToByteArray(
+      base64CiphertextAndNonceAndSalt,
+    );
+
+    // Create buffers of salt and ciphertextAndNonce.
+    const salt = ciphertextAndNonceAndSalt.slice(0, this.#PBKDF2_SALT_SIZE);
+    const ciphertextAndNonce = ciphertextAndNonceAndSalt.slice(
+      this.#PBKDF2_SALT_SIZE,
+      ciphertextAndNonceAndSalt.length,
+    );
+
+    // Derive the key using PBKDF2.
+    const key = pbkdf2(sha256, password, salt, {
+      c: iterations,
+      dkLen: this.#ALGORITHM_KEY_SIZE,
+    });
+
+    // Decrypt and return result.
+    return bytesToUtf8(this.#decrypt(ciphertextAndNonce, key));
+  }
+
+  #encrypt(plaintext: Uint8Array, key: Uint8Array): Uint8Array {
+    const nonce = randomBytes(this.#ALGORITHM_NONCE_SIZE);
+
+    // Encrypt and prepend nonce.
+    const ciphertext = gcm(key, nonce).encrypt(plaintext);
+
+    return concatBytes(nonce, ciphertext);
+  }
+
+  #decrypt(ciphertextAndNonce: Uint8Array, key: Uint8Array): Uint8Array {
+    // Create buffers of nonce and ciphertext.
+    const nonce = ciphertextAndNonce.slice(0, this.#ALGORITHM_NONCE_SIZE);
+    const ciphertext = ciphertextAndNonce.slice(
+      this.#ALGORITHM_NONCE_SIZE,
+      ciphertextAndNonce.length,
+    );
+
+    // Decrypt and return result.
+    return gcm(key, nonce).decrypt(ciphertext);
+  }
+}
+
+const encryption = new EncryptorDecryptor();
+export default encryption;
+
+export function createSHA256Hash(data: string): string {
+  const hashedData = sha256(data);
+  return bytesToHex(hashedData);
+}

app/scripts/controllers/user-storage/mocks/mockServices.ts

@@ -0,0 +1,40 @@
+import nock from 'nock';
+import { USER_STORAGE_ENDPOINT, GetUserStorageResponse } from '../services';
+import { createEntryPath } from '../schema';
+import { MOCK_ENCRYPTED_STORAGE_DATA, MOCK_STORAGE_KEY } from './mockStorage';
+
+export const MOCK_USER_STORAGE_NOTIFICATIONS_ENDPOINT = `${USER_STORAGE_ENDPOINT}${createEntryPath(
+  'notification_settings',
+  MOCK_STORAGE_KEY,
+)}`;
+
+type MockReply = {
+  status: nock.StatusCode;
+  body?: nock.Body;
+};
+
+const MOCK_GET_USER_STORAGE_RESPONSE: GetUserStorageResponse = {
+  HashedKey: 'HASHED_KEY',
+  Data: MOCK_ENCRYPTED_STORAGE_DATA,
+};
+export function mockEndpointGetUserStorage(mockReply?: MockReply) {
+  const reply = mockReply ?? {
+    status: 200,
+    body: MOCK_GET_USER_STORAGE_RESPONSE,
+  };
+
+  const mockEndpoint = nock(MOCK_USER_STORAGE_NOTIFICATIONS_ENDPOINT)
+    .get('')
+    .reply(reply.status, reply.body);
+
+  return mockEndpoint;
+}
+
+export function mockEndpointUpsertUserStorage(
+  mockReply?: Pick<MockReply, 'status'>,
+) {
+  const mockEndpoint = nock(MOCK_USER_STORAGE_NOTIFICATIONS_ENDPOINT)
+    .put('')
+    .reply(mockReply?.status ?? 204);
+  return mockEndpoint;
+}

app/scripts/controllers/user-storage/mocks/mockStorage.ts

@@ -0,0 +1,9 @@
+import encryption, { createSHA256Hash } from '../encryption';
+
+export const MOCK_STORAGE_KEY_SIGNATURE = 'mockStorageKey';
+export const MOCK_STORAGE_KEY = createSHA256Hash(MOCK_STORAGE_KEY_SIGNATURE);
+export const MOCK_STORAGE_DATA = JSON.stringify({ hello: 'world' });
+export const MOCK_ENCRYPTED_STORAGE_DATA = encryption.encryptString(
+  MOCK_STORAGE_DATA,
+  MOCK_STORAGE_KEY,
+);

app/scripts/controllers/user-storage/schema.test.ts

@@ -0,0 +1,21 @@
+import { USER_STORAGE_ENTRIES, createEntryPath } from './schema';
+
+describe('schema.ts - createEntryPath()', () => {
+  const MOCK_STORAGE_KEY = 'MOCK_STORAGE_KEY';
+
+  test('creates a valid entry path', () => {
+    const result = createEntryPath('notification_settings', MOCK_STORAGE_KEY);
+
+    // Ensures that the path and the entry name are correct.
+    // If this differs then indicates a potential change on how this path is computed
+    const expected = `/${USER_STORAGE_ENTRIES.notification_settings.path}/50f65447980018849b991e038d7ad87de5cf07fbad9736b0280e93972e17bac8`;
+    expect(result).toBe(expected);
+  });
+
+  test('Should throw if using an entry that does not exist', () => {
+    expect(() => {
+      // @ts-expect-error mocking a fake entry for testing.
+      createEntryPath('fake_entry');
+    }).toThrow();
+  });
+});