fix: improve nonce performance (#162)

* feat: improve nonce performance

Introduce Nonce class to maintain both Uint8Array and DataView and uint64 so that we can remove nonceToBytes function

* chore: keep the encrypt() and decrypt() interfaces

* chore: move nonce contants to nonce.ts

* chore: fix MAX_NONCE

Author: twoeths

benchmarks/nonce.js

@@ -0,0 +1,45 @@
+/* eslint-disable */
+import benchmark from 'benchmark'
+import { Nonce } from '../dist/src/nonce.js'
+
+/**
+ * Using Nonce class is 150x faster than nonceToBytes
+ * nonceToBytes x 2.25 ops/sec ±1.41% (10 runs sampled)
+ * Nonce class x 341 ops/sec ±0.71% (87 runs sampled)
+ */
+function nonceToBytes (n) {
+  // Even though we're treating the nonce as 8 bytes, RFC7539 specifies 12 bytes for a nonce.
+  const nonce = new Uint8Array(12)
+  new DataView(nonce.buffer, nonce.byteOffset, nonce.byteLength).setUint32(4, n, true)
+  return nonce
+}
+const main = function () {
+  const bench1 = new benchmark('nonceToBytes', {
+    fn: function () {
+      for (let i = 1e6; i < 2 * 1e6; i++) {
+        nonceToBytes(i)
+      }
+    }
+  })
+    .on('complete', function (stats) {
+      console.log(String(stats.currentTarget))
+    })
+
+  bench1.run()
+
+  const bench2 = new benchmark('Nonce class', {
+    fn: function () {
+      const nonce = new Nonce(1e6)
+      for (let i = 1e6; i < 2 * 1e6; i++) {
+        nonce.increment()
+      }
+    }
+  })
+    .on('complete', function (stats) {
+      console.log(String(stats.currentTarget))
+    })
+
+  bench2.run()
+}
+
+main()

src/@types/handshake.ts

@@ -1,5 +1,6 @@
 import type { bytes, bytes32, uint64 } from './basic.js'
 import type { KeyPair } from './libp2p.js'
+import type { Nonce } from '../nonce.js'
 
 export type Hkdf = [bytes, bytes, bytes]
 
@@ -11,9 +12,9 @@ export interface MessageBuffer {
 
 export interface CipherState {
   k: bytes32
-  // For performance reasons, the nonce is represented as a JS `number`
+  // For performance reasons, the nonce is represented as a Nonce object
   // The nonce is treated as a uint64, even though the underlying `number` only has 52 safely-available bits.
-  n: uint64
+  n: Nonce
 }
 
 export interface SymmetricState {

src/handshakes/abstract-handshake.ts

@@ -1,20 +1,11 @@
 import { equals as uint8ArrayEquals } from 'uint8arrays/equals'
 import { concat as uint8ArrayConcat } from 'uint8arrays/concat'
 import { fromString as uint8ArrayFromString } from 'uint8arrays'
-import type { bytes, bytes32, uint64 } from '../@types/basic.js'
+import type { bytes, bytes32 } from '../@types/basic.js'
 import type { CipherState, MessageBuffer, SymmetricState } from '../@types/handshake.js'
 import type { ICryptoInterface } from '../crypto.js'
 import { logger } from '../logger.js'
-
-export const MIN_NONCE = 0
-// For performance reasons, the nonce is represented as a JS `number`
-// JS `number` can only safely represent integers up to 2 ** 53 - 1
-// This is a slight deviation from the noise spec, which describes the max nonce as 2 ** 64 - 2
-// The effect is that this implementation will need a new handshake to be performed after fewer messages are exchanged than other implementations with full uint64 nonces.
-// 2 ** 53 - 1 is still a large number of messages, so the practical effect of this is negligible.
-export const MAX_NONCE = Number.MAX_SAFE_INTEGER
-
-const ERR_MAX_NONCE = 'Cipherstate has reached maximum n, a new handshake must be performed'
+import { Nonce } from '../nonce.js'
 
 export abstract class AbstractHandshake {
   public crypto: ICryptoInterface
@@ -25,14 +16,14 @@ export abstract class AbstractHandshake {
 
   public encryptWithAd (cs: CipherState, ad: Uint8Array, plaintext: Uint8Array): bytes {
     const e = this.encrypt(cs.k, cs.n, ad, plaintext)
-    this.setNonce(cs, this.incrementNonce(cs.n))
+    cs.n.increment()
 
     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)
-    this.setNonce(cs, this.incrementNonce(cs.n))
+    cs.n.increment()
 
     return { plaintext, valid }
   }
@@ -42,10 +33,6 @@ export abstract class AbstractHandshake {
     return !this.isEmptyKey(cs.k)
   }
 
-  protected setNonce (cs: CipherState, nonce: uint64): void {
-    cs.n = nonce
-  }
-
   protected createEmptyKey (): bytes32 {
     return new Uint8Array(32)
   }
@@ -55,26 +42,10 @@ export abstract class AbstractHandshake {
     return uint8ArrayEquals(emptyKey, k)
   }
 
-  protected incrementNonce (n: uint64): uint64 {
-    return n + 1
-  }
-
-  protected nonceToBytes (n: uint64): bytes {
-    // Even though we're treating the nonce as 8 bytes, RFC7539 specifies 12 bytes for a nonce.
-    const nonce = new Uint8Array(12)
-    new DataView(nonce.buffer, nonce.byteOffset, nonce.byteLength).setUint32(4, n, true)
+  protected encrypt (k: bytes32, n: Nonce, ad: Uint8Array, plaintext: Uint8Array): bytes {
+    n.assertValue()
 
-    return nonce
-  }
-
-  protected encrypt (k: bytes32, n: uint64, ad: Uint8Array, plaintext: Uint8Array): bytes {
-    if (n > MAX_NONCE) {
-      throw new Error(ERR_MAX_NONCE)
-    }
-
-    const nonce = this.nonceToBytes(n)
-
-    return this.crypto.chaCha20Poly1305Encrypt(plaintext, nonce, ad, k)
+    return this.crypto.chaCha20Poly1305Encrypt(plaintext, n.getBytes(), ad, k)
   }
 
   protected encryptAndHash (ss: SymmetricState, plaintext: bytes): bytes {
@@ -89,13 +60,10 @@ export abstract class AbstractHandshake {
     return ciphertext
   }
 
-  protected decrypt (k: bytes32, n: uint64, ad: bytes, ciphertext: bytes): {plaintext: bytes, valid: boolean} {
-    if (n > MAX_NONCE) {
-      throw new Error(ERR_MAX_NONCE)
-    }
+  protected decrypt (k: bytes32, n: Nonce, ad: bytes, ciphertext: bytes): {plaintext: bytes, valid: boolean} {
+    n.assertValue()
 
-    const nonce = this.nonceToBytes(n)
-    const encryptedMessage = this.crypto.chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k)
+    const encryptedMessage = this.crypto.chaCha20Poly1305Decrypt(ciphertext, n.getBytes(), ad, k)
 
     if (encryptedMessage) {
       return {
@@ -154,8 +122,7 @@ export abstract class AbstractHandshake {
   }
 
   protected initializeKey (k: bytes32): CipherState {
-    const n = MIN_NONCE
-    return { k, n }
+    return { k, n: new Nonce() }
   }
 
   // Symmetric state related

src/logger.ts

@@ -43,8 +43,8 @@ export function logRemoteEphemeralKey (re: Uint8Array): void {
 
 export function logCipherState (session: NoiseSession): void {
   if (session.cs1 && session.cs2) {
-    keyLogger(`CIPHER_STATE_1 ${session.cs1.n} ${uint8ArrayToString(session.cs1.k, 'hex')}`)
-    keyLogger(`CIPHER_STATE_2 ${session.cs2.n} ${uint8ArrayToString(session.cs2.k, 'hex')}`)
+    keyLogger(`CIPHER_STATE_1 ${session.cs1.n.getUint64()} ${uint8ArrayToString(session.cs1.k, 'hex')}`)
+    keyLogger(`CIPHER_STATE_2 ${session.cs2.n.getUint64()} ${uint8ArrayToString(session.cs2.k, 'hex')}`)
   } else {
     keyLogger('Missing cipher state.')
   }

src/nonce.ts

@@ -0,0 +1,49 @@
+import type { bytes, uint64 } from './@types/basic'
+
+export const MIN_NONCE = 0
+// For performance reasons, the nonce is represented as a JS `number`
+// Although JS `number` can safely represent integers up to 2 ** 53 - 1, we choose to only use
+// 4 bytes to store the data for performance reason.
+// This is a slight deviation from the noise spec, which describes the max nonce as 2 ** 64 - 2
+// The effect is that this implementation will need a new handshake to be performed after fewer messages are exchanged than other implementations with full uint64 nonces.
+// this MAX_NONCE is still a large number of messages, so the practical effect of this is negligible.
+export const MAX_NONCE = 0xffffffff
+
+const ERR_MAX_NONCE = 'Cipherstate has reached maximum n, a new handshake must be performed'
+
+/**
+ * The nonce is an uint that's increased over time.
+ * Maintaining different representations help improve performance.
+ */
+export class Nonce {
+  private n: uint64
+  private readonly bytes: bytes
+  private readonly view: DataView
+
+  constructor (n = MIN_NONCE) {
+    this.n = n
+    this.bytes = new Uint8Array(12)
+    this.view = new DataView(this.bytes.buffer, this.bytes.byteOffset, this.bytes.byteLength)
+    this.view.setUint32(4, n, true)
+  }
+
+  increment (): void {
+    this.n++
+    // Even though we're treating the nonce as 8 bytes, RFC7539 specifies 12 bytes for a nonce.
+    this.view.setUint32(4, this.n, true)
+  }
+
+  getBytes (): bytes {
+    return this.bytes
+  }
+
+  getUint64 (): uint64 {
+    return this.n
+  }
+
+  assertValue (): void {
+    if (this.n > MAX_NONCE) {
+      throw new Error(ERR_MAX_NONCE)
+    }
+  }
+}