feat: add encryption capabilities

.gitignore

@@ -76,3 +76,6 @@ node_modules/
 !.yarn/releases
 !.yarn/sdks
 !.yarn/versions
+
+# ignore intellij local files
+.idea/

snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/message-signing-snap.git"
   },
   "source": {
-    "shasum": "poMEiKK8bywGEFXmN99xRFTlHWZ0Ea/0G5Bu2zhR6g0=",
+    "shasum": "a/rW2l5FJjl1LpdYd94q8lzu7emRJNfWX1JDmEQjO/A=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

src/entropy-keys.test.ts

@@ -4,6 +4,8 @@ import { sha256 } from '@noble/hashes/sha256';
 
 import {
   getAllPublicEntropyKeys,
+  decryptMessage,
+  getEncryptionPublicKey,
   getPublicEntropyKey,
   signMessageWithEntropyKey,
 } from './entropy-keys';
@@ -12,6 +14,8 @@ const MOCK_PRIVATE_KEY =
   '0xec180de430cef919666c2009b91ca3d3b7f6c471136abc9937fa40b89357bbb9';
 const MOCK_PUBLIC_KEY =
   '0x02c291ee55d10abcc46de22b775cb0782b06f386ced8b0d0fccb8007a686bbddad';
+const MOCK_ENCRYPTION_PUBLIC_KEY =
+  '0xcd1f3cd0cf8cdc6437b36224968c56b661b9c3765fc0f49dd75051c341705b07';
 
 const MOCK_PRIVATE_KEY_SRP2 =
   '0x1111111130cef919666c2009b91ca3d3b7f6c471136abc9937fa40b811111111';
@@ -29,7 +33,7 @@ describe('getPublicEntropyKey() tests', () => {
   it('should return a public key from a known private key with a source ID', async () => {
     mockSnapGetEntropy();
 
-    const address = await getPublicEntropyKey('mockEntropySourceId');
+    const address = await getPublicEntropyKey('mock_id_2');
     expect(address).toBe(MOCK_PUBLIC_KEY_SRP2);
   });
 });
@@ -56,10 +60,7 @@ describe('signMessageWithEntropyKey() tests', () => {
     mockSnapGetEntropy();
 
     const message = 'hello world';
-    const signature = await signMessageWithEntropyKey(
-      message,
-      'mockEntropySourceId',
-    );
+    const signature = await signMessageWithEntropyKey(message, 'mock_id_2');
     const EXPECTED_SIGNATURE =
       '0xc7cd8af7ddd59287eee7e99f111e637d3e16add417edab1efd388e2688db77dd6e36f1189e47600eeab49c750d4247c5300dbdfbf1d1fad2b6a970070e5148c7';
     expect(signature).toBe(EXPECTED_SIGNATURE);
@@ -75,47 +76,99 @@ describe('signMessageWithEntropyKey() tests', () => {
 
 describe('getAllPublicEntropyKeys() tests', () => {
   it('should get entropy source IDs and SRP IDs relationship map', async () => {
-    const mockEntropySources = [
-      { name: 'source1', id: 'id1', type: 'mnemonic', primary: true },
-      { name: 'source2', id: 'id2', type: 'mnemonic', primary: false },
-    ] as ListEntropySourcesResult;
-
-    const mockSnapRequest = jest
-      .fn()
-      .mockImplementation(async (r: { method: string; params: any }) => {
-        if (r.method === 'snap_listEntropySources') {
-          return mockEntropySources;
-        } else if (r.method === 'snap_getEntropy') {
-          if (r.params.source === 'id2') {
-            return MOCK_PRIVATE_KEY_SRP2;
-          }
-          return MOCK_PRIVATE_KEY;
-        }
-
-        throw new Error(`TEST ENV - Snap Request was not mocked: ${r.method}`);
-      });
-
-    (global as any).snap = {
-      request: mockSnapRequest,
-    };
+    mockSnapGetEntropy();
 
     const relationshipMap = await getAllPublicEntropyKeys();
     expect(relationshipMap).toStrictEqual([
-      ['id1', MOCK_PUBLIC_KEY],
-      ['id2', MOCK_PUBLIC_KEY_SRP2],
+      ['mock_id_1', MOCK_PUBLIC_KEY],
+      ['mock_id_2', MOCK_PUBLIC_KEY_SRP2],
     ]);
   });
 });
 
+describe('encryption tests', () => {
+  it('gets the expected encryption key', async () => {
+    mockSnapGetEntropy();
+    const publicEncryptionKey = await getEncryptionPublicKey();
+    expect(publicEncryptionKey).toBe(MOCK_ENCRYPTION_PUBLIC_KEY);
+  });
+
+  it('decrypts a message intended for this public encryption key', async () => {
+    mockSnapGetEntropy();
+
+    const encrypted = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: '8jFlr2cYdkyBIookw6akE8lA0f+odanN',
+      ephemPublicKey: 'UOPWFRaPdY6kKDEoM/ovCBCT2p4PSx6MMdOgNzA2gC0=',
+      ciphertext: '2UNvHmxnmOHtI9vhf7gfeD/J7Q/q6vqjEQqY',
+    };
+
+    const decrypted = await decryptMessage(encrypted);
+    expect(decrypted).toBe('hello world');
+  });
+
+  it('fails tp decrypt a message intended for someone else', async () => {
+    mockSnapGetEntropy();
+
+    const encrypted = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: 'h63LvxvCOBP3x3Oou2n5JYgCM1p4p+DF',
+      ephemPublicKey: 'lmIBlLKUuSBIRjlo+/hL7ngWYpMWQ7biqk7Y6pDsaXY=',
+      ciphertext: 'some/ONE/else/SHOULD/read/this/COSJY',
+    };
+
+    await expect(decryptMessage(encrypted)).rejects.toThrow(/invalid tag/u);
+  });
+
+  it('decrypts a message intended for a specific public encryption key', async () => {
+    mockSnapGetEntropy();
+
+    const encrypted = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: '8jFlr2cYdkyBIookw6akE8lA0f+odanN',
+      ephemPublicKey: 'UOPWFRaPdY6kKDEoM/ovCBCT2p4PSx6MMdOgNzA2gC0=',
+      ciphertext: '2UNvHmxnmOHtI9vhf7gfeD/J7Q/q6vqjEQqY',
+    };
+
+    const decrypted = await decryptMessage(encrypted, 'mock_id_1');
+    expect(decrypted).toBe('hello world');
+  });
+
+  it('fails to decrypt a message with the wrong version', async () => {
+    mockSnapGetEntropy();
+
+    const encrypted = {
+      version: 'x25519-something-else-entirely',
+      nonce: 'dontcare',
+      ephemPublicKey: 'dontcare',
+      ciphertext: 'dontcare',
+    };
+
+    await expect(decryptMessage(encrypted)).rejects.toThrow(
+      /Encryption type\/version not supported \(x25519-something-else-entirely\)./u,
+    );
+  });
+});
+
 function mockSnapGetEntropy() {
+  const mockEntropySources = [
+    { name: 'source1', id: 'mock_id_1', type: 'mnemonic', primary: true },
+    { name: 'source2', id: 'mock_id_2', type: 'mnemonic', primary: false },
+  ] as ListEntropySourcesResult;
+
   const mockSnapRequest = jest
     .fn()
     .mockImplementation(async (r: { method: string; params: any }) => {
-      if (r.method === 'snap_getEntropy') {
-        if (r.params.source === 'mockEntropySourceId') {
-          return MOCK_PRIVATE_KEY_SRP2;
+      if (r.method === 'snap_listEntropySources') {
+        return mockEntropySources;
+      } else if (r.method === 'snap_getEntropy') {
+        switch (r.params.source) {
+          case 'mock_id_2':
+            return MOCK_PRIVATE_KEY_SRP2;
+          case 'mock_id_1':
+          default:
+            return MOCK_PRIVATE_KEY;
         }
-        return MOCK_PRIVATE_KEY;
       }
 
       throw new Error(`TEST ENV - Snap Request was not mocked: ${r.method}`);

src/entropy-keys.ts

@@ -1,9 +1,13 @@
 import type { ListEntropySourcesResult } from '@metamask/snaps-sdk';
+import type { Eip1024EncryptedData } from '@metamask/utils';
+import { bytesToHex } from '@metamask/utils';
+import { x25519 } from '@noble/curves/ed25519';
 import { secp256k1 } from '@noble/curves/secp256k1';
 import { sha256 } from '@noble/hashes/sha256';
 
 import type { EntropySourceId, EntropySourceIdSrpIdMap } from './types';
 import { addressToBytes, bytesToAddress } from './utils/address-conversion';
+import { ERC1024 } from './utils/ERC1024';
 
 /**
  * Lists the entropy sources available for the snap.
@@ -18,8 +22,9 @@ async function listEntropySources(): Promise<ListEntropySourcesResult> {
 /**
  * Retrieve the snap entropy private key.
  * @param entropySourceId - Optional entropy Source ID following SIP-30.
- * @param salt - The salt used to obtain a domain specific entropy. Metamask internal origins should use `undefined`.
+ * @param salt - Optional salt to use for the entropy derivation. Useful for generating keys for different purposes and/or to obtain a domain specific entropy.
  * @returns Entropy Private Key Hex.
+ * @see https://metamask.github.io/SIPs/SIPS/sip-6
  */
 async function getEntropy(
   entropySourceId?: EntropySourceId,
@@ -88,6 +93,81 @@ export async function getAllPublicEntropyKeys(
   return entropySourceIdsAndSrpIdsMap;
 }
 
+// This is used to derive an encryption key from the entropy, to avoid key reuse.
+const KEY_PURPOSE_ENCRYPTION = 'metamask:snaps:encryption';
+
+/**
+ * Retrieve the secret encryption key for this snap.
+ * @param entropySourceId - Optional entropy Source ID following SIP-30.
+ * @param extraSalt - The extraSalt used to obtain a domain specific entropy. Metamask internal origins should use `undefined`.
+ * @returns Encryption Secret Key Bytes.
+ * @see https://metamask.github.io/SIPs/SIPS/sip-6 for more information about how the derivation works.
+ */
+async function getEncryptionSecretKey(
+  entropySourceId?: EntropySourceId,
+  extraSalt?: string,
+): Promise<`0x${string}`> {
+  const salt = extraSalt
+    ? `${KEY_PURPOSE_ENCRYPTION}${extraSalt}`
+    : KEY_PURPOSE_ENCRYPTION;
+  return await getEntropy(entropySourceId, salt);
+}
+
+/**
+ * Retrieve the public encryption key for this snap.
+ * @param entropySourceId - Optional entropy Source ID following SIP-30.
+ * @param salt - The salt used to obtain a domain specific entropy. Metamask internal origins should use `undefined`.
+ * @returns Public Key Hex.
+ */
+export async function getEncryptionPublicKey(
+  entropySourceId?: EntropySourceId,
+  salt?: string,
+): Promise<`0x${string}`> {
+  const secretKeyHex = await getEncryptionSecretKey(entropySourceId, salt);
+  return bytesToHex(x25519.getPublicKey(secretKeyHex.slice(2)));
+}
+
+/**
+ * Error message for decrypting with the wrong private key.
+ * This is thrown by @noble/ciphers, so we will need to match.
+ */
+const INVALID_TAG_ERROR = 'invalid tag';
+
+/**
+ * Decrypt an encrypted message using the snap specific encryption key.
+ * In case there are multiple possible private keys, the entropy source ID can be used to specify which one to use.
+ * For privacy reasons, it may be impossible to know which entropy source ID to use, so all entropy sources will be tried if this parameter is missing.
+ * @param encryptedMessage - The encrypted message, encoded as a `Eip1024EncryptedData` object.
+ * @param entropySourceId - Optional entropy Source ID following SIP-30. If this is missing, all available entropy sources will be tried.
+ * @param salt - The salt used to obtain a domain specific entropy. Metamask internal origins should use `undefined`.
+ * @returns The decrypted message (string).
+ */
+export async function decryptMessage(
+  encryptedMessage: Eip1024EncryptedData,
+  entropySourceId?: EntropySourceId,
+  salt?: string,
+): Promise<string> {
+  if (entropySourceId) {
+    const secretKeyHex = await getEncryptionSecretKey(entropySourceId, salt);
+    return ERC1024.decrypt(encryptedMessage, secretKeyHex);
+  }
+  const entropySources = await listEntropySources();
+  let decryptionError = null;
+  for (const source of entropySources) {
+    const secretKeyHex = await getEncryptionSecretKey(source.id, salt);
+    try {
+      return ERC1024.decrypt(encryptedMessage, secretKeyHex);
+    } catch (error: any) {
+      if (error.message !== INVALID_TAG_ERROR) {
+        // If decryption fails because of the key, try the next entropy source.
+        // Otherwise, it's likely we matched the correct key so remember the error.
+        decryptionError = decryptionError ?? error;
+      }
+    }
+  }
+  throw decryptionError ?? new Error(INVALID_TAG_ERROR);
+}
+
 /**
  * Signs a message and returns the signature.
  * @param message - Message to sign.