From 256f9af6503257eb4ef300e6be5f8f5418e30683 Mon Sep 17 00:00:00 2001
From: Cayman <caymannava@gmail.com>
Date: Mon, 11 Sep 2023 17:04:40 -0400
Subject: [PATCH] feat: selectively use nodejs or as implementation

---
 packages/beacon-node/package.json             |  1 +
 .../beacon-node/src/network/libp2p/noise.ts   | 56 ++++++++++++++-----
 .../test/perf/network/noise/sendData.test.ts  | 12 +++-
 yarn.lock                                     |  5 ++
 4 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/packages/beacon-node/package.json b/packages/beacon-node/package.json
index dd48d49ad1a3..8239cc0bb192 100644
--- a/packages/beacon-node/package.json
+++ b/packages/beacon-node/package.json
@@ -95,6 +95,7 @@
     "check-readme": "typescript-docs-verifier"
   },
   "dependencies": {
+    "@chainsafe/as-chacha20poly1305": "^0.1.0",
     "@chainsafe/as-sha256": "^0.3.1",
     "@chainsafe/bls": "7.1.1",
     "@chainsafe/blst": "^0.2.9",
diff --git a/packages/beacon-node/src/network/libp2p/noise.ts b/packages/beacon-node/src/network/libp2p/noise.ts
index aeb017776cdd..fcb00fe41893 100644
--- a/packages/beacon-node/src/network/libp2p/noise.ts
+++ b/packages/beacon-node/src/network/libp2p/noise.ts
@@ -1,20 +1,20 @@
 import crypto from "node:crypto";
 import type {ConnectionEncrypter} from "@libp2p/interface/connection-encrypter";
 import {ICryptoInterface, noise, pureJsCrypto} from "@chainsafe/libp2p-noise";
+import {digest} from "@chainsafe/as-sha256";
+import {newInstance, ChaCha20Poly1305} from "@chainsafe/as-chacha20poly1305";
 
-type Bytes = Uint8Array;
-type Bytes32 = Uint8Array;
+const ctx = newInstance();
+const asImpl = new ChaCha20Poly1305(ctx);
 
 const CHACHA_POLY1305 = "chacha20-poly1305";
 
-// same as default, but we use node crypto chacha20poly1305 and sha256
-const lodestarCrypto: ICryptoInterface = {
-  ...pureJsCrypto,
-  hashSHA256(data: Uint8Array): Uint8Array {
+const nodeCrypto: Pick<ICryptoInterface, "hashSHA256" | "chaCha20Poly1305Encrypt" | "chaCha20Poly1305Decrypt"> = {
+  hashSHA256(data) {
     return crypto.createHash("sha256").update(data).digest();
   },
 
-  chaCha20Poly1305Encrypt(plaintext: Uint8Array, nonce: Uint8Array, ad: Uint8Array, k: Bytes32): Bytes {
+  chaCha20Poly1305Encrypt(plaintext, nonce, ad, k) {
     const cipher = crypto.createCipheriv(CHACHA_POLY1305, k, nonce, {
       authTagLength: 16,
     });
@@ -27,13 +27,7 @@ const lodestarCrypto: ICryptoInterface = {
     return encrypted;
   },
 
-  chaCha20Poly1305Decrypt(
-    ciphertext: Uint8Array,
-    nonce: Uint8Array,
-    ad: Uint8Array,
-    k: Bytes32,
-    _dst?: Uint8Array
-  ): Bytes | null {
+  chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k, _dst) {
     const authTag = ciphertext.slice(ciphertext.length - 16);
     const text = ciphertext.slice(0, ciphertext.length - 16);
     const decipher = crypto.createDecipheriv(CHACHA_POLY1305, k, nonce, {
@@ -52,6 +46,40 @@ const lodestarCrypto: ICryptoInterface = {
   },
 };
 
+const asCrypto: Pick<ICryptoInterface, "hashSHA256" | "chaCha20Poly1305Encrypt" | "chaCha20Poly1305Decrypt"> = {
+  hashSHA256(data) {
+    return digest(data);
+  },
+  chaCha20Poly1305Encrypt(plaintext, nonce, ad, k) {
+    return asImpl.seal(k, nonce, plaintext, ad);
+  },
+  chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k, dst) {
+    return asImpl.open(k, nonce, ciphertext, ad, dst);
+  },
+};
+
+// benchmarks show that for chacha20poly1305
+// the as implementation is faster for smaller payloads(<1200)
+// and the node implementation is faster for larger payloads
+const lodestarCrypto: ICryptoInterface = {
+  ...pureJsCrypto,
+  hashSHA256(data) {
+    return nodeCrypto.hashSHA256(data);
+  },
+  chaCha20Poly1305Encrypt(plaintext, nonce, ad, k) {
+    if (plaintext.length < 1200) {
+      return asCrypto.chaCha20Poly1305Encrypt(plaintext, nonce, ad, k);
+    }
+    return nodeCrypto.chaCha20Poly1305Encrypt(plaintext, nonce, ad, k);
+  },
+  chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k, dst) {
+    if (ciphertext.length < 1200) {
+      return asCrypto.chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k, dst);
+    }
+    return nodeCrypto.chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k, dst);
+  },
+};
+
 export function createNoise(): () => ConnectionEncrypter {
   return noise({crypto: lodestarCrypto});
 }
diff --git a/packages/beacon-node/test/perf/network/noise/sendData.test.ts b/packages/beacon-node/test/perf/network/noise/sendData.test.ts
index 51a64fe130f8..a6843c13ece2 100644
--- a/packages/beacon-node/test/perf/network/noise/sendData.test.ts
+++ b/packages/beacon-node/test/perf/network/noise/sendData.test.ts
@@ -8,7 +8,17 @@ import {createNoise} from "../../../../src/network/libp2p/noise.js";
 describe("network / noise / sendData", () => {
   const numberOfMessages = 1000;
 
-  for (const messageLength of [2 ** 8, 2 ** 9, 2 ** 10, 2 ** 11, 2 ** 12, 2 ** 14, 2 ** 16]) {
+  for (const messageLength of [
+    //
+    2 ** 8,
+    2 ** 9,
+    2 ** 10,
+    1200,
+    2 ** 11,
+    2 ** 12,
+    2 ** 14,
+    2 ** 16,
+  ]) {
     itBench({
       id: `send data - ${numberOfMessages} ${messageLength}B messages`,
       beforeEach: async () => {
diff --git a/yarn.lock b/yarn.lock
index ab219b6bcfda..02fe48359b25 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -436,6 +436,11 @@
   resolved "https://registry.yarnpkg.com/@bcoe/v8-coverage/-/v8-coverage-0.2.3.tgz#75a2e8b51cb758a7553d6804a5932d7aace75c39"
   integrity sha512-0hYQ8SB4Db5zvZB4axdMHGwEaQjkZzFjQiN9LVYvIFB2nSUHW9tYpxWriPrWDASIxiaXax83REcLxuSdnGPZtw==
 
+"@chainsafe/as-chacha20poly1305@^0.1.0":
+  version "0.1.0"
+  resolved "https://registry.yarnpkg.com/@chainsafe/as-chacha20poly1305/-/as-chacha20poly1305-0.1.0.tgz#7da6f8796f9b42dac6e830a086d964f1f9189e09"
+  integrity sha512-BpNcL8/lji/GM3+vZ/bgRWqJ1q5kwvTFmGPk7pxm/QQZDbaMI98waOHjEymTjq2JmdD/INdNBFOVSyJofXg7ew==
+
 "@chainsafe/as-sha256@^0.3.1":
   version "0.3.1"
   resolved "https://registry.npmjs.org/@chainsafe/as-sha256/-/as-sha256-0.3.1.tgz"