fix: improve nonce performance

benchmarks/nonce.js

@@ -0,0 +1,45 @@
+/* eslint-disable */
+import benchmark from 'benchmark'
+import { Nonce } from '../dist/src/@types/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.increase()
+      }
+    }
+  })
+    .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/@types/nonce.ts

@@ -0,0 +1,32 @@
+import type { bytes, uint64 } from './basic'
+
+/**
+ * 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: uint64) {
+    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)
+  }
+
+  increase (): 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
+  }
+}

src/handshakes/abstract-handshake.ts

@@ -1,10 +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'
+import { Nonce } from '../@types/nonce.js'
 
 export const MIN_NONCE = 0
 // For performance reasons, the nonce is represented as a JS `number`
@@ -24,15 +25,15 @@ 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))
+    const e = this.encrypt(cs, ad, plaintext)
+    cs.n.increase()
 
     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))
+    const { plaintext, valid } = this.decrypt(cs, ad, ciphertext)
+    cs.n.increase()
 
     return { plaintext, valid }
   }
@@ -42,10 +43,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 +52,13 @@ 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)
-
-    return nonce
-  }
-
-  protected encrypt (k: bytes32, n: uint64, ad: Uint8Array, plaintext: Uint8Array): bytes {
-    if (n > MAX_NONCE) {
+  protected encrypt (cs: CipherState, ad: Uint8Array, plaintext: Uint8Array): bytes {
+    const nonce = cs.n
+    if (nonce.getUint64() > 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, nonce.getBytes(), ad, cs.k)
   }
 
   protected encryptAndHash (ss: SymmetricState, plaintext: bytes): bytes {
@@ -89,13 +73,13 @@ export abstract class AbstractHandshake {
     return ciphertext
   }
 
-  protected decrypt (k: bytes32, n: uint64, ad: bytes, ciphertext: bytes): {plaintext: bytes, valid: boolean} {
-    if (n > MAX_NONCE) {
+  protected decrypt (cs: CipherState, ad: bytes, ciphertext: bytes): {plaintext: bytes, valid: boolean} {
+    const nonce = cs.n
+    if (nonce.getUint64() > MAX_NONCE) {
       throw new Error(ERR_MAX_NONCE)
     }
 
-    const nonce = this.nonceToBytes(n)
-    const encryptedMessage = this.crypto.chaCha20Poly1305Decrypt(ciphertext, nonce, ad, k)
+    const encryptedMessage = this.crypto.chaCha20Poly1305Decrypt(ciphertext, nonce.getBytes(), ad, cs.k)
 
     if (encryptedMessage) {
       return {
@@ -154,8 +138,7 @@ export abstract class AbstractHandshake {
   }
 
   protected initializeKey (k: bytes32): CipherState {
-    const n = MIN_NONCE
-    return { k, n }
+    return { k, n: new Nonce(MIN_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.')
   }