feat: notifications - user storage controller

app/scripts/controllers/user-storage/auth-snap-requests.ts

@@ -0,0 +1,21 @@
+import type { SnapId } from '@metamask/snaps-sdk';
+import type { HandleSnapRequest } from '@metamask/snaps-controllers';
+import { HandlerType } from '@metamask/snaps-utils';
+
+type SnapRPCRequest = Parameters<HandleSnapRequest['handler']>[0];
+
+const snapId = 'npm:@metamask/message-signing-snap' as SnapId;
+
+export function createSnapSignMessageRequest(
+  message: `metamask:${string}`,
+): SnapRPCRequest {
+  return {
+    snapId,
+    origin: '',
+    handler: HandlerType.OnRpcRequest,
+    request: {
+      method: 'signMessage',
+      params: { message },
+    },
+  };
+}

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.ts

@@ -0,0 +1,38 @@
+import { createSHA256Hash } from './encryption';
+
+type UserStorageEntry = { path: string; entryName: string };
+
+/**
+ * The User Storage Endpoint requires a path and an entry name.
+ * Developers can provide additional paths by extending this variable below
+ */
+export const USER_STORAGE_ENTRIES = {
+  notification_settings: {
+    path: 'notifications',
+    entryName: 'notification_settings',
+  },
+} satisfies Record<string, UserStorageEntry>;
+
+export type UserStorageEntryKeys = keyof typeof USER_STORAGE_ENTRIES;
+
+/**
+ * Constructs a unique entry path for a user.
+ * This can be done due to the uniqueness of the storage key (no users will share the same storage key).
+ * The users entry is a unique hash that cannot be reversed.
+ *
+ * @param entryKey
+ * @param storageKey
+ * @returns
+ */
+export function createEntryPath(
+  entryKey: UserStorageEntryKeys,
+  storageKey: string,
+): string {
+  const entry = USER_STORAGE_ENTRIES[entryKey];
+  if (!entry) {
+    throw new Error(`user-storage - invalid entry provided: ${entryKey}`);
+  }
+
+  const hashedKey = createSHA256Hash(entry.entryName + storageKey);
+  return `/${entry.path}/${hashedKey}`;
+}

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

@@ -0,0 +1,89 @@
+import {
+  mockEndpointGetUserStorage,
+  mockEndpointUpsertUserStorage,
+} from './mocks/mockServices';
+import {
+  MOCK_ENCRYPTED_STORAGE_DATA,
+  MOCK_STORAGE_DATA,
+  MOCK_STORAGE_KEY,
+} from './mocks/mockStorage';
+import {
+  GetUserStorageResponse,
+  getUserStorage,
+  upsertUserStorage,
+} from './services';
+
+describe('user-storage/services.ts - getUserStorage() tests', () => {
+  test('returns user storage data', async () => {
+    const mockGetUserStorage = mockEndpointGetUserStorage();
+    const result = await actCallGetUserStorage();
+
+    mockGetUserStorage.done();
+    expect(result).toBe(MOCK_STORAGE_DATA);
+  });
+
+  test('returns null if endpoint does not have entry', async () => {
+    const mockGetUserStorage = mockEndpointGetUserStorage({ status: 404 });
+    const result = await actCallGetUserStorage();
+
+    mockGetUserStorage.done();
+    expect(result).toBe(null);
+  });
+
+  test('returns null if endpoint fails', async () => {
+    const mockGetUserStorage = mockEndpointGetUserStorage({ status: 500 });
+    const result = await actCallGetUserStorage();
+
+    mockGetUserStorage.done();
+    expect(result).toBe(null);
+  });
+
+  test('returns null if unable to decrypt data', async () => {
+    const badResponseData: GetUserStorageResponse = {
+      HashedKey: 'MOCK_HASH',
+      Data: 'Bad Encrypted Data',
+    };
+    const mockGetUserStorage = mockEndpointGetUserStorage({
+      status: 200,
+      body: badResponseData,
+    });
+    const result = await actCallGetUserStorage();
+
+    mockGetUserStorage.done();
+    expect(result).toBe(null);
+  });
+
+  function actCallGetUserStorage() {
+    return getUserStorage({
+      bearerToken: 'MOCK_BEARER_TOKEN',
+      entryKey: 'notification_settings',
+      storageKey: MOCK_STORAGE_KEY,
+    });
+  }
+});
+
+describe('user-storage/services.ts - upsertUserStorage() tests', () => {
+  test('invokes upsert endpoint with no errors', async () => {
+    const mockUpsertUserStorage = mockEndpointUpsertUserStorage();
+    await actCallUpsertUserStorage();
+
+    expect(mockUpsertUserStorage.isDone()).toBe(true);
+  });
+
+  test('throws error if unable to upsert user storage', async () => {
+    const mockUpsertUserStorage = mockEndpointUpsertUserStorage({
+      status: 500,
+    });
+
+    await expect(actCallUpsertUserStorage()).rejects.toThrow();
+    mockUpsertUserStorage.done();
+  });
+
+  function actCallUpsertUserStorage() {
+    return upsertUserStorage(MOCK_ENCRYPTED_STORAGE_DATA, {
+      bearerToken: 'MOCK_BEARER_TOKEN',
+      entryKey: 'notification_settings',
+      storageKey: MOCK_STORAGE_KEY,
+    });
+  }
+});