feat: add encryption capabilities

fixes #101

.gitignore

@@ -75,3 +75,5 @@ node_modules/
 !.yarn/releases
 !.yarn/sdks
 !.yarn/versions
+
+.idea/

snap.manifest.json

@@ -7,7 +7,7 @@
     "url": "https://github.com/MetaMask/message-signing-snap.git"
   },
   "source": {
-    "shasum": "MVxPWtAZ1Sdw63kwfZswKlYQ6dddnr7vvbj4V7j/Fc4=",
+    "shasum": "9DZJX/w4gCd9AF+cPGI+7HRVVxcbmfeTBVl1dIBvlxI=",
     "location": {
       "npm": {
         "filePath": "dist/bundle.js",

src/entropy-keys.test.ts

@@ -1,4 +1,9 @@
-import { getPublicEntropyKey, signMessageWithEntropyKey } from './entropy-keys';
+import {
+  decryptMessage,
+  getEncryptionPublicKey,
+  getPublicEntropyKey,
+  signMessageWithEntropyKey,
+} from './entropy-keys';
 
 const MOCK_PRIVATE_KEY =
   '0xec180de430cef919666c2009b91ca3d3b7f6c471136abc9937fa40b89357bbb9';
@@ -26,6 +31,31 @@ describe('signMessageWithEntropyKey() tests', () => {
   });
 });
 
+describe('getEncryptionPublicKey() tests', () => {
+  it('gets the expected encryption key', async () => {
+    mockSnapGetEntropy();
+
+    const publicEncryptionKey = await getEncryptionPublicKey();
+    const EXPECTED_KEY =
+      '0x50cbcf3915730e501b7476e92157307f6e9aade2a2798cf3832f73cd4990281b';
+    expect(publicEncryptionKey).toBe(EXPECTED_KEY);
+  });
+
+  it('decrypts a message intended for this public encryption key', async () => {
+    mockSnapGetEntropy();
+
+    const encrypted = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: 'h63LvxvCOBP3x3Oou2n5JYgCM1p4p+DF',
+      ephemPublicKey: 'lmIBlLKUuSBIRjlo+/hL7ngWYpMWQ7biqk7Y6pDsaXY=',
+      ciphertext: 'g+TpY8OlU0AS9VPvaTIIqpFnWNKvWw2COSJY',
+    };
+
+    const decrypted = await decryptMessage(encrypted);
+    expect(decrypted).toBe('hello world');
+  });
+});
+
 function mockSnapGetEntropy() {
   const mockSnapRequest = jest
     .fn()

src/entropy-keys.ts

@@ -1,7 +1,12 @@
+import type { Eip1024EncryptedData, Hex } from '@metamask/utils';
+import { bytesToHex, concatBytes } from '@metamask/utils';
+import { utf8ToBytes } from '@noble/ciphers/utils';
+import { x25519 } from '@noble/curves/ed25519';
 import { secp256k1 } from '@noble/curves/secp256k1';
 import { sha256 } from '@noble/hashes/sha256';
 
 import { addressToBytes, bytesToAddress } from './utils/address-conversion';
+import { ERC1024 } from './utils/ERC1024';
 
 /**
  * Retrieve the snap entropy private key.
@@ -37,6 +42,40 @@ export async function getPublicEntropyKey(): Promise<string> {
   return bytesToAddress(secp256k1.getPublicKey(privateKey));
 }
 
+// This is used to derive an encryption key from the entropy, to avoid key reuse.
+const staticSalt = 'metamask:snaps:encryption';
+
+/**
+ * Retrieve the secret encryption key for this snap.
+ * The key is derived from the entropy key and a static salt as sha256(entropy | staticSalt).
+ * @returns Secret Key Bytes.
+ */
+async function getEncryptionSecretKey(): Promise<Uint8Array> {
+  const privateEntropy = await getPrivateEntropyKey();
+  return sha256(concatBytes([privateEntropy, utf8ToBytes(staticSalt)]));
+}
+
+/**
+ * Retrieve the public encryption key for this snap.
+ * @returns Public Key Hex.
+ */
+export async function getEncryptionPublicKey(): Promise<Hex> {
+  const secretKeyBytes = await getEncryptionSecretKey();
+  return bytesToHex(x25519.getPublicKey(secretKeyBytes));
+}
+
+/**
+ * Decrypt an encrypted message using the snap specific encryption key.
+ * @param encryptedMessage - The encrypted message, encoded as a `Eip1024EncryptedData` object.
+ * @returns The decrypted message (string).
+ */
+export async function decryptMessage(
+  encryptedMessage: Eip1024EncryptedData,
+): Promise<string> {
+  const secretKeyBytes = await getEncryptionSecretKey();
+  return ERC1024.decrypt(encryptedMessage, secretKeyBytes);
+}
+
 /**
  * Signs a message and returns the signature.
  * @param message - Message to sign.

src/index.test.ts

@@ -1,8 +1,11 @@
 import { installSnap } from '@metamask/snaps-jest';
+import type { Hex } from '@metamask/utils';
 import { hexToBytes } from '@noble/ciphers/utils';
 import { secp256k1 } from '@noble/curves/secp256k1';
 import { sha256 } from '@noble/hashes/sha256';
 
+import { ERC1024 } from './utils/ERC1024';
+
 describe('onRpcRequest - getPublicKey', () => {
   it('should return this snaps public key', async () => {
     const snap = await installSnap();
@@ -19,6 +22,146 @@ describe('onRpcRequest - getPublicKey', () => {
   });
 });
 
+describe('onRpcRequest - getEncryptionPublicKey', () => {
+  it('should return this snaps encryption public key', async () => {
+    const snap = await installSnap();
+    const response = await snap.request({
+      method: 'getEncryptionPublicKey',
+    });
+
+    // E.g. length = 66 chars
+    // E.g. starts with '0x'
+    const result = 'result' in response.response && response.response.result;
+    expect(result?.toString().length).toBe(66);
+    expect(result?.toString().startsWith('0x')).toBe(true);
+  });
+});
+
+describe('onRpcRequest - decryptMessage', () => {
+  it('should decrypt a message intended for the snaps public key', async () => {
+    const snap = await installSnap();
+    const pkResponse = await snap.request({
+      method: 'getEncryptionPublicKey',
+    });
+    const publicKey = (
+      'result' in pkResponse.response && pkResponse.response.result
+    )?.toString() as Hex;
+    const message = 'hello world';
+    const encryptedMessage = ERC1024.encrypt(publicKey, message);
+    const response = await snap.request({
+      method: 'decryptMessage',
+      params: { data: encryptedMessage },
+    });
+
+    const result = 'result' in response.response && response.response.result;
+    expect(result?.toString()).toBe('hello world');
+  });
+
+  it('should fail to decrypt a message intended for a different recipient', async () => {
+    const snap = await installSnap();
+    const encryptedMessage = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: 'h63LvxvCOBP3x3Oou2n5JYgCM1p4p+DF',
+      ephemPublicKey: 'lmIBlLKUuSBIRjlo+/hL7ngWYpMWQ7biqk7Y6pDsaXY=',
+      ciphertext: 'g+TpY8OlU0AS9VPvaTIIqpFnWNKvWw2COSJY',
+    };
+    const response = await snap.request({
+      method: 'decryptMessage',
+      params: { data: encryptedMessage },
+    });
+
+    expect(response).toRespondWithError({
+      code: -32603,
+      message: 'invalid tag',
+      stack: expect.any(String),
+    });
+  });
+
+  it('should reject a message with invalid version', async () => {
+    const snap = await installSnap();
+    const encryptedMessage = {
+      version: '1', // invalid version
+      nonce: 'h63LvxvCOBP3x3Oou2n5JYgCM1p4p+DF',
+      ephemPublicKey: 'lmIBlLKUuSBIRjlo+/hL7ngWYpMWQ7biqk7Y6pDsaXY=',
+      ciphertext: 'g+TpY8OlU0AS9VPvaTIIqpFnWNKvWw2COSJY',
+    };
+    const response = await snap.request({
+      method: 'decryptMessage',
+      params: { data: encryptedMessage },
+    });
+
+    expect(response).toRespondWithError({
+      code: -32602,
+      message:
+        '`decryptMessage`, must take a `data` parameter that must match the Eip1024EncryptedData schema',
+      stack: expect.any(String),
+    });
+  });
+
+  it('should reject a message with invalid nonce', async () => {
+    const snap = await installSnap();
+    const encryptedMessage = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: 'tooshort',
+      ephemPublicKey: 'lmIBlLKUuSBIRjlo+/hL7ngWYpMWQ7biqk7Y6pDsaXY=',
+      ciphertext: 'g+TpY8OlU0AS9VPvaTIIqpFnWNKvWw2COSJY',
+    };
+    const response = await snap.request({
+      method: 'decryptMessage',
+      params: { data: encryptedMessage },
+    });
+
+    expect(response).toRespondWithError({
+      code: -32602,
+      message:
+        '`decryptMessage`, must take a `data` parameter that must match the Eip1024EncryptedData schema',
+      stack: expect.any(String),
+    });
+  });
+
+  it('should reject a message with invalid ephemPublicKey', async () => {
+    const snap = await installSnap();
+    const encryptedMessage = {
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: 'h63LvxvCOBP3x3Oou2n5JYgCM1p4p+DF',
+      ephemPublicKey: 'invalid base 64',
+      ciphertext: 'g+TpY8OlU0AS9VPvaTIIqpFnWNKvWw2COSJY',
+    };
+    const response = await snap.request({
+      method: 'decryptMessage',
+      params: { data: encryptedMessage },
+    });
+
+    expect(response).toRespondWithError({
+      code: -32602,
+      message:
+        '`decryptMessage`, must take a `data` parameter that must match the Eip1024EncryptedData schema',
+      stack: expect.any(String),
+    });
+  });
+
+  it('should reject a message with invalid params', async () => {
+    const snap = await installSnap();
+    const encryptedMessage = JSON.stringify({
+      version: 'x25519-xsalsa20-poly1305',
+      nonce: 'h63LvxvCOBP3x3Oou2n5JYgCM1p4p+DF',
+      ephemPublicKey: 'lmIBlLKUuSBIRjlo+/hL7ngWYpMWQ7biqk7Y6pDsaXY=',
+      ciphertext: 'g+TpY8OlU0AS9VPvaTIIqpFnWNKvWw2COSJY',
+    });
+    const response = await snap.request({
+      method: 'decryptMessage',
+      params: { data: encryptedMessage },
+    });
+
+    expect(response).toRespondWithError({
+      code: -32602,
+      message:
+        '`decryptMessage`, must take a `data` parameter that must match the Eip1024EncryptedData schema',
+      stack: expect.any(String),
+    });
+  });
+});
+
 describe('onRpcRequest - signMessage', () => {
   it('should return a signature that can be verified', async () => {
     const snap = await installSnap();

src/index.ts

@@ -2,12 +2,26 @@ import { rpcErrors } from '@metamask/rpc-errors';
 import type { OnRpcRequestHandler } from '@metamask/snaps-sdk';
 import { z } from 'zod';
 
-import { getPublicEntropyKey, signMessageWithEntropyKey } from './entropy-keys';
+import {
+  decryptMessage,
+  getEncryptionPublicKey,
+  getPublicEntropyKey,
+  signMessageWithEntropyKey,
+} from './entropy-keys';
 
 const SignMessageParamsSchema = z.object({
   message: z.string().startsWith('metamask:'),
 });
 
+const DecryptMessageParamsSchema = z.object({
+  data: z.object({
+    version: z.string().regex(/^x25519-xsalsa20-poly1305$/u),
+    nonce: z.string().length(32).base64(), // 24 bytes, base64 encoded
+    ephemPublicKey: z.string().length(44).base64(), // 32 bytes, base64 encoded
+    ciphertext: z.string().base64(),
+  }),
+});
+
 /**
  * Asserts the shape of the `signMessage` request.
  * @param params - Any method params to assert.
@@ -26,6 +40,24 @@ function assertSignMessageParams(
   }
 }
 
+/**
+ * Asserts the shape of the `decryptMessage` request matches the expected {data: Eip1024EncryptedData}.
+ * @param params - The input params to assert.
+ * @returns {never} Returns nothing, but will throw error if params don't match what is required.
+ */
+function assertDecryptMessageParams(
+  params: unknown,
+): asserts params is z.infer<typeof DecryptMessageParamsSchema> {
+  try {
+    DecryptMessageParamsSchema.parse(params);
+  } catch (error: any) {
+    throw rpcErrors.invalidParams({
+      message:
+        '`decryptMessage`, must take a `data` parameter that must match the Eip1024EncryptedData schema',
+    });
+  }
+}
+
 export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
   switch (request.method) {
     case 'getPublicKey': {
@@ -37,6 +69,15 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => {
       const { message } = params;
       return await signMessageWithEntropyKey(message);
     }
+    case 'getEncryptionPublicKey': {
+      return getEncryptionPublicKey();
+    }
+    case 'decryptMessage': {
+      const { params } = request;
+      assertDecryptMessageParams(params);
+      const { data } = params;
+      return await decryptMessage(data);
+    }
     default:
       throw rpcErrors.methodNotFound({
         data: { method: request.method },