From ae6ed8278e20a1462aac42bdd6acd441c5e17948 Mon Sep 17 00:00:00 2001
From: Nicholas Molnar <65710+neekolas@users.noreply.github.com>
Date: Fri, 6 Oct 2023 09:10:54 -0700
Subject: [PATCH] feat: add self-encryption support

---
 package-lock.json                |  8 ++++++-
 package.json                     |  3 ++-
 src/crypto/SelfEncryption.ts     | 32 ++++++++++++++++++++++++++++
 test/crypto/P4Encryption.test.ts | 36 ++++++++++++++++++++++++++++++++
 4 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 src/crypto/SelfEncryption.ts
 create mode 100644 test/crypto/P4Encryption.test.ts

diff --git a/package-lock.json b/package-lock.json
index 4f00959e3..49612d1ce 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -10,7 +10,8 @@
       "license": "MIT",
       "dependencies": {
         "@noble/secp256k1": "^1.5.2",
-        "@xmtp/proto": "^3.29.0",
+        "@xmtp/ecies-bindings-wasm": "^0.1.5",
+        "@xmtp/proto": "^3.28.0-beta.1",
         "async-mutex": "^0.4.0",
         "elliptic": "^6.5.4",
         "ethers": "^5.5.3",
@@ -4819,6 +4820,11 @@
         "url": "https://opencollective.com/eslint"
       }
     },
+    "node_modules/@xmtp/ecies-bindings-wasm": {
+      "version": "0.1.5",
+      "resolved": "https://registry.npmjs.org/@xmtp/ecies-bindings-wasm/-/ecies-bindings-wasm-0.1.5.tgz",
+      "integrity": "sha512-SBxBxQCX0weIIHgrVxZ0QmAX7bqjhBsheXtmcgiciC/ZyAB7UC0KsNx2xqYrIDE0YQm6WXRg90qsX7a30SrLfQ=="
+    },
     "node_modules/@xmtp/proto": {
       "version": "3.29.0",
       "resolved": "https://registry.npmjs.org/@xmtp/proto/-/proto-3.29.0.tgz",
diff --git a/package.json b/package.json
index 1485823f8..dcf5a5e05 100644
--- a/package.json
+++ b/package.json
@@ -84,7 +84,8 @@
   },
   "dependencies": {
     "@noble/secp256k1": "^1.5.2",
-    "@xmtp/proto": "^3.29.0",
+    "@xmtp/ecies-bindings-wasm": "^0.1.5",
+    "@xmtp/proto": "^3.28.0-beta.1",
     "async-mutex": "^0.4.0",
     "elliptic": "^6.5.4",
     "ethers": "^5.5.3",
diff --git a/src/crypto/SelfEncryption.ts b/src/crypto/SelfEncryption.ts
new file mode 100644
index 000000000..e30d30a0d
--- /dev/null
+++ b/src/crypto/SelfEncryption.ts
@@ -0,0 +1,32 @@
+import {
+  // eslint-disable-next-line camelcase
+  ecies_decrypt_k256_sha3_256,
+  // eslint-disable-next-line camelcase
+  ecies_encrypt_k256_sha3_256,
+} from '@xmtp/ecies-bindings-wasm'
+import { PrivateKey } from '.'
+
+// Uses ECIES to encrypt messages where the sender and recipient are the same
+export default class SelfEncryption {
+  privateKey: PrivateKey
+
+  constructor(identityKey: PrivateKey) {
+    this.privateKey = identityKey
+  }
+
+  encrypt(data: Uint8Array): Uint8Array {
+    return ecies_encrypt_k256_sha3_256(
+      this.privateKey.publicKey.secp256k1Uncompressed.bytes,
+      this.privateKey.secp256k1.bytes,
+      data
+    )
+  }
+
+  decrypt(message: Uint8Array): Uint8Array {
+    return ecies_decrypt_k256_sha3_256(
+      this.privateKey.publicKey.secp256k1Uncompressed.bytes,
+      this.privateKey.secp256k1.bytes,
+      message
+    )
+  }
+}
diff --git a/test/crypto/P4Encryption.test.ts b/test/crypto/P4Encryption.test.ts
new file mode 100644
index 000000000..17d16e015
--- /dev/null
+++ b/test/crypto/P4Encryption.test.ts
@@ -0,0 +1,36 @@
+import { PrivateKeyBundleV1 } from '../../src/crypto/PrivateKeyBundle'
+import SelfEncryption from '../../src/crypto/SelfEncryption'
+import { newWallet } from '../helpers'
+import { equalBytes } from '../../src/crypto/utils'
+
+describe('SelfEncryption', () => {
+  let bundle: PrivateKeyBundleV1
+
+  beforeEach(async () => {
+    bundle = await PrivateKeyBundleV1.generate(newWallet())
+  })
+
+  it('round trips data', async () => {
+    const message = new TextEncoder().encode('hello world')
+    const encryptor = new SelfEncryption(bundle.identityKey)
+
+    const ciphertext = encryptor.encrypt(message)
+    expect(ciphertext).toBeDefined()
+
+    const decrypted = encryptor.decrypt(ciphertext)
+    expect(equalBytes(decrypted, message)).toBeTruthy()
+  })
+
+  it('throws on decryption failure', async () => {
+    const message = new TextEncoder().encode('hello world')
+    const encryptor = new SelfEncryption(bundle.identityKey)
+
+    const ciphertext = encryptor.encrypt(message)
+    expect(ciphertext).toBeDefined()
+
+    const differentEncryptor = new SelfEncryption(
+      (await PrivateKeyBundleV1.generate(newWallet())).identityKey
+    )
+    expect(() => differentEncryptor.decrypt(ciphertext)).toThrow()
+  })
+})