feat!: refactor codebase (#407)

* chore: cleanup comments

* fix: simplify uint16 encode/decode

* fix: fix node crypto

* chore: consolidate types

* fix: avoid reallocations of empty key

* fix: simplify XX initialization

* chore: add more comments to types

* fix: remove unused psk state variable

* feat: refactor the codebase

Author: wemeetagain

src/@types/basic.ts

@@ -1,17 +1,27 @@
-import { type Uint8ArrayList } from 'uint8arraylist'
-import type { bytes32 } from './@types/basic.js'
-import type { Hkdf } from './@types/handshake.js'
-import type { KeyPair } from './@types/libp2p.js'
+import type { ICrypto, KeyPair } from './types.js'
+import type { Uint8ArrayList } from 'uint8arraylist'
 
+/** Underlying crypto implementation, meant to be overridable */
 export interface ICryptoInterface {
   hashSHA256(data: Uint8Array | Uint8ArrayList): Uint8Array
 
-  getHKDF(ck: bytes32, ikm: Uint8Array): Hkdf
+  getHKDF(ck: Uint8Array, ikm: Uint8Array): [Uint8Array, Uint8Array, Uint8Array]
 
   generateX25519KeyPair(): KeyPair
   generateX25519KeyPairFromSeed(seed: Uint8Array): KeyPair
   generateX25519SharedKey(privateKey: Uint8Array | Uint8ArrayList, publicKey: Uint8Array | Uint8ArrayList): Uint8Array
 
-  chaCha20Poly1305Encrypt(plaintext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: bytes32): Uint8ArrayList | Uint8Array
-  chaCha20Poly1305Decrypt(ciphertext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: bytes32, dst?: Uint8Array): Uint8ArrayList | Uint8Array | null
+  chaCha20Poly1305Encrypt(plaintext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: Uint8Array): Uint8ArrayList | Uint8Array
+  chaCha20Poly1305Decrypt(ciphertext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: Uint8Array, dst?: Uint8Array): Uint8ArrayList | Uint8Array
+}
+
+export function wrapCrypto (crypto: ICryptoInterface): ICrypto {
+  return {
+    generateKeypair: crypto.generateX25519KeyPair,
+    dh: (keypair, publicKey) => crypto.generateX25519SharedKey(keypair.privateKey, publicKey).subarray(0, 32),
+    encrypt: crypto.chaCha20Poly1305Encrypt,
+    decrypt: crypto.chaCha20Poly1305Decrypt,
+    hash: crypto.hashSHA256,
+    hkdf: crypto.getHKDF
+  }
 }

src/@types/handshake-interface.ts

@@ -4,8 +4,8 @@ import { digest } from '@chainsafe/as-sha256'
 import { Uint8ArrayList } from 'uint8arraylist'
 import { isElectronMain } from 'wherearewe'
 import { pureJsCrypto } from './js.js'
-import type { KeyPair } from '../@types/libp2p.js'
 import type { ICryptoInterface } from '../crypto.js'
+import type { KeyPair } from '../types.js'
 
 const ctx = newInstance()
 const asImpl = new ChaCha20Poly1305(ctx)
@@ -38,7 +38,7 @@ const nodeCrypto: Pick<ICryptoInterface, 'hashSHA256' | 'chaCha20Poly1305Encrypt
       const final = cipher.final()
       const tag = cipher.getAuthTag()
 
-      return Buffer.concat([updated, tag, final], updated.byteLength + tag.byteLength + final.byteLength)
+      return Buffer.concat([updated, final, tag], updated.byteLength + final.byteLength + tag.byteLength)
     }
 
     const output = new Uint8ArrayList()
@@ -112,7 +112,11 @@ const asCrypto: Pick<ICryptoInterface, 'hashSHA256' | 'chaCha20Poly1305Encrypt'
     return asImpl.seal(k, nonce, plaintext.subarray(), ad)
   },
   chaCha20Poly1305Decrypt (ciphertext, nonce, ad, k, dst) {
-    return asImpl.open(k, nonce, ciphertext.subarray(), ad, dst)
+    const plaintext = asImpl.open(k, nonce, ciphertext.subarray(), ad, dst)
+    if (!plaintext) {
+      throw new Error('Invalid chacha20poly1305 decryption')
+    }
+    return plaintext
   }
 }
 
@@ -181,8 +185,7 @@ export const defaultCrypto: ICryptoInterface = {
         publicKey
       ], X25519_PREFIX.byteLength + publicKey.byteLength)
     } else {
-      publicKey.prepend(X25519_PREFIX)
-      publicKey = publicKey.subarray()
+      publicKey = new Uint8ArrayList(X25519_PREFIX, publicKey).subarray()
     }
 
     if (privateKey instanceof Uint8Array) {
@@ -191,8 +194,7 @@ export const defaultCrypto: ICryptoInterface = {
         privateKey
       ], PKCS8_PREFIX.byteLength + privateKey.byteLength)
     } else {
-      privateKey.prepend(PKCS8_PREFIX)
-      privateKey = privateKey.subarray()
+      privateKey = new Uint8ArrayList(PKCS8_PREFIX, privateKey).subarray()
     }
 
     return crypto.diffieHellman({

src/@types/handshake.ts

@@ -2,18 +2,16 @@ import { chacha20poly1305 } from '@noble/ciphers/chacha'
 import { x25519 } from '@noble/curves/ed25519'
 import { extract, expand } from '@noble/hashes/hkdf'
 import { sha256 } from '@noble/hashes/sha256'
-import type { bytes32 } from '../@types/basic.js'
-import type { Hkdf } from '../@types/handshake.js'
-import type { KeyPair } from '../@types/libp2p.js'
 import type { ICryptoInterface } from '../crypto.js'
+import type { KeyPair } from '../types.js'
 import type { Uint8ArrayList } from 'uint8arraylist'
 
 export const pureJsCrypto: ICryptoInterface = {
   hashSHA256 (data: Uint8Array | Uint8ArrayList): Uint8Array {
     return sha256(data.subarray())
   },
 
-  getHKDF (ck: bytes32, ikm: Uint8Array): Hkdf {
+  getHKDF (ck: Uint8Array, ikm: Uint8Array): [Uint8Array, Uint8Array, Uint8Array] {
     const prk = extract(sha256, ikm, ck)
     const okmU8Array = expand(sha256, prk, undefined, 96)
     const okm = okmU8Array
@@ -48,11 +46,11 @@ export const pureJsCrypto: ICryptoInterface = {
     return x25519.getSharedSecret(privateKey.subarray(), publicKey.subarray())
   },
 
-  chaCha20Poly1305Encrypt (plaintext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: bytes32): Uint8Array {
+  chaCha20Poly1305Encrypt (plaintext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: Uint8Array): Uint8Array {
     return chacha20poly1305(k, nonce, ad).encrypt(plaintext.subarray())
   },
 
-  chaCha20Poly1305Decrypt (ciphertext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: bytes32, dst?: Uint8Array): Uint8Array | null {
+  chaCha20Poly1305Decrypt (ciphertext: Uint8Array | Uint8ArrayList, nonce: Uint8Array, ad: Uint8Array, k: Uint8Array, dst?: Uint8Array): Uint8Array {
     return chacha20poly1305(k, nonce, ad).decrypt(ciphertext.subarray(), dst)
   }
 }

src/@types/libp2p.ts

@@ -1,12 +1,11 @@
-import { Uint8ArrayList } from 'uint8arraylist'
-import { alloc as uint8ArrayAlloc, allocUnsafe as uint8ArrayAllocUnsafe } from 'uint8arrays/alloc'
-import type { bytes } from './@types/basic.js'
-import type { MessageBuffer } from './@types/handshake.js'
+import { type Uint8ArrayList } from 'uint8arraylist'
+import { allocUnsafe as uint8ArrayAllocUnsafe } from 'uint8arrays/alloc'
 import type { LengthDecoderFunction } from 'it-length-prefixed'
 
 export const uint16BEEncode = (value: number): Uint8Array => {
   const target = uint8ArrayAllocUnsafe(2)
-  new DataView(target.buffer, target.byteOffset, target.byteLength).setUint16(0, value, false)
+  target[0] = value >> 8
+  target[1] = value
   return target
 }
 uint16BEEncode.bytes = 2
@@ -15,59 +14,12 @@ export const uint16BEDecode: LengthDecoderFunction = (data: Uint8Array | Uint8Ar
   if (data.length < 2) throw RangeError('Could not decode int16BE')
 
   if (data instanceof Uint8Array) {
-    return new DataView(data.buffer, data.byteOffset, data.byteLength).getUint16(0, false)
+    let value = 0
+    value += data[0] << 8
+    value += data[1]
+    return value
   }
 
   return data.getUint16(0)
 }
 uint16BEDecode.bytes = 2
-
-// Note: IK and XX encoder usage is opposite (XX uses in stages encode0 where IK uses encode1)
-
-export function encode0 (message: MessageBuffer): Uint8ArrayList {
-  return new Uint8ArrayList(message.ne, message.ciphertext)
-}
-
-export function encode1 (message: MessageBuffer): Uint8ArrayList {
-  return new Uint8ArrayList(message.ne, message.ns, message.ciphertext)
-}
-
-export function encode2 (message: MessageBuffer): Uint8ArrayList {
-  return new Uint8ArrayList(message.ns, message.ciphertext)
-}
-
-export function decode0 (input: bytes): MessageBuffer {
-  if (input.length < 32) {
-    throw new Error('Cannot decode stage 0 MessageBuffer: length less than 32 bytes.')
-  }
-
-  return {
-    ne: input.subarray(0, 32),
-    ciphertext: input.subarray(32, input.length),
-    ns: uint8ArrayAlloc(0)
-  }
-}
-
-export function decode1 (input: bytes): MessageBuffer {
-  if (input.length < 80) {
-    throw new Error('Cannot decode stage 1 MessageBuffer: length less than 80 bytes.')
-  }
-
-  return {
-    ne: input.subarray(0, 32),
-    ns: input.subarray(32, 80),
-    ciphertext: input.subarray(80, input.length)
-  }
-}
-
-export function decode2 (input: bytes): MessageBuffer {
-  if (input.length < 48) {
-    throw new Error('Cannot decode stage 2 MessageBuffer: length less than 48 bytes.')
-  }
-
-  return {
-    ne: uint8ArrayAlloc(0),
-    ns: input.subarray(0, 48),
-    ciphertext: input.subarray(48, input.length)
-  }
-}