feat: reuse encrypted data to avoid memory allocation (#242)

feat: reused ciphertext allocation in decrypt function
ChainSafe · Nov 22, 2022 · 0b4b7f6 · 0b4b7f6
1 parent 9137bf6
commit 0b4b7f6
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 11 deletions.
diff --git a/src/@types/handshake-interface.ts b/src/@types/handshake-interface.ts
@@ -8,5 +8,5 @@ export interface IHandshake {
   remotePeer: PeerId
   remoteExtensions: NoiseExtensions
   encrypt: (plaintext: bytes, session: NoiseSession) => bytes
-  decrypt: (ciphertext: bytes, session: NoiseSession) => { plaintext: bytes, valid: boolean }
+  decrypt: (ciphertext: bytes, session: NoiseSession, dst?: Uint8Array) => { plaintext: bytes, valid: boolean }
 }
diff --git a/src/crypto.ts b/src/crypto.ts
@@ -12,5 +12,5 @@ export interface ICryptoInterface {
   generateX25519SharedKey: (privateKey: Uint8Array, publicKey: Uint8Array) => Uint8Array
 
   chaCha20Poly1305Encrypt: (plaintext: Uint8Array, nonce: Uint8Array, ad: Uint8Array, k: bytes32) => bytes
-  chaCha20Poly1305Decrypt: (ciphertext: Uint8Array, nonce: Uint8Array, ad: Uint8Array, k: bytes32) => bytes | null
+  chaCha20Poly1305Decrypt: (ciphertext: Uint8Array, nonce: Uint8Array, ad: Uint8Array, k: bytes32, dst?: Uint8Array) => bytes | null
 }
diff --git a/src/crypto/stablelib.ts b/src/crypto/stablelib.ts
@@ -52,9 +52,9 @@ export const stablelib: ICryptoInterface = {
     return ctx.seal(nonce, plaintext, ad)
   },
 
-  chaCha20Poly1305Decrypt (ciphertext: Uint8Array, nonce: Uint8Array, ad: Uint8Array, k: bytes32): bytes | null {
+  chaCha20Poly1305Decrypt (ciphertext: Uint8Array, nonce: Uint8Array, ad: Uint8Array, k: bytes32, dst?: Uint8Array): bytes | null {
     const ctx = new ChaCha20Poly1305(k)
 
-    return ctx.open(nonce, ciphertext, ad)
+    return ctx.open(nonce, ciphertext, ad, dst)
   }
 }
diff --git a/src/crypto/streaming.ts b/src/crypto/streaming.ts
@@ -1,3 +1,4 @@
+import { TAG_LENGTH } from '@stablelib/chacha20poly1305'
 import type { Transform } from 'it-stream-types'
 import type { Uint8ArrayList } from 'uint8arraylist'
 import type { IHandshake } from '../@types/handshake-interface.js'
@@ -35,7 +36,16 @@ export function decryptStream (handshake: IHandshake, metrics?: MetricsRegistry)
           end = chunk.length
         }
 
-        const { plaintext: decrypted, valid } = handshake.decrypt(chunk.subarray(i, end), handshake.session)
+        if (end - TAG_LENGTH < i) {
+          throw new Error('Invalid chunk')
+        }
+        const encrypted = chunk.subarray(i, end)
+        // memory allocation is not cheap so reuse the encrypted Uint8Array
+        // see https://github.com/ChainSafe/js-libp2p-noise/pull/242#issue-1422126164
+        // this is ok because chacha20 reads bytes one by one and don't reread after that
+        // it's also tested in https://github.com/ChainSafe/as-chacha20poly1305/pull/1/files#diff-25252846b58979dcaf4e41d47b3eadd7e4f335e7fb98da6c049b1f9cd011f381R48
+        const dst = chunk.subarray(i, end - TAG_LENGTH)
+        const { plaintext: decrypted, valid } = handshake.decrypt(encrypted, handshake.session, dst)
         if (!valid) {
           metrics?.decryptErrors.increment()
           throw new Error('Failed to validate decrypted chunk')

diff --git a/src/handshake-xx.ts b/src/handshake-xx.ts
@@ -147,10 +147,10 @@ export class XXHandshake implements IHandshake {
     return this.xx.encryptWithAd(cs, new Uint8Array(0), plaintext)
   }
 
-  public decrypt (ciphertext: Uint8Array, session: NoiseSession): { plaintext: bytes, valid: boolean } {
+  public decrypt (ciphertext: Uint8Array, session: NoiseSession, dst?: Uint8Array): { plaintext: bytes, valid: boolean } {
     const cs = this.getCS(session, false)
 
-    return this.xx.decryptWithAd(cs, new Uint8Array(0), ciphertext)
+    return this.xx.decryptWithAd(cs, new Uint8Array(0), ciphertext, dst)
   }
 
   public getRemoteStaticKey (): bytes {

diff --git a/src/handshakes/abstract-handshake.ts b/src/handshakes/abstract-handshake.ts
@@ -21,8 +21,8 @@ export abstract class AbstractHandshake {
     return e
   }
 
-  public decryptWithAd (cs: CipherState, ad: Uint8Array, ciphertext: Uint8Array): {plaintext: bytes, valid: boolean} {
-    const { plaintext, valid } = this.decrypt(cs.k, cs.n, ad, ciphertext)
+  public decryptWithAd (cs: CipherState, ad: Uint8Array, ciphertext: Uint8Array, dst?: Uint8Array): {plaintext: bytes, valid: boolean} {
+    const { plaintext, valid } = this.decrypt(cs.k, cs.n, ad, ciphertext, dst)
     if (valid) cs.n.increment()
 
     return { plaintext, valid }
@@ -60,10 +60,10 @@ export abstract class AbstractHandshake {
     return ciphertext
   }
 
-  protected decrypt (k: bytes32, n: Nonce, ad: bytes, ciphertext: bytes): {plaintext: bytes, valid: boolean} {
+  protected decrypt (k: bytes32, n: Nonce, ad: bytes, ciphertext: bytes, dst?: Uint8Array): {plaintext: bytes, valid: boolean} {
     n.assertValue()
 
-    const encryptedMessage = this.crypto.chaCha20Poly1305Decrypt(ciphertext, n.getBytes(), ad, k)
+    const encryptedMessage = this.crypto.chaCha20Poly1305Decrypt(ciphertext, n.getBytes(), ad, k, dst)
 
     if (encryptedMessage) {
       return {