chore: Add prefix in encryptStream (#236)

* Add prefix in encryptStream * chore: fix lint error Co-authored-by: Cayman <[email protected]>
ChainSafe · Oct 31, 2022 · db66eeb · db66eeb
1 parent 4803ffd
commit db66eeb
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/src/crypto/streaming.ts b/src/crypto/streaming.ts
@@ -2,6 +2,7 @@ import type { Transform } from 'it-stream-types'
 import type { Uint8ArrayList } from 'uint8arraylist'
 import type { IHandshake } from '../@types/handshake-interface.js'
 import { NOISE_MSG_MAX_LENGTH_BYTES, NOISE_MSG_MAX_LENGTH_BYTES_WITHOUT_TAG } from '../constants.js'
+import { uint16BEEncode } from '../encoder.js'
 
 // Returns generator that encrypts payload from the user
 export function encryptStream (handshake: IHandshake): Transform<Uint8Array> {
@@ -14,6 +15,8 @@ export function encryptStream (handshake: IHandshake): Transform<Uint8Array> {
         }
 
         const data = handshake.encrypt(chunk.subarray(i, end), handshake.session)
+
+        yield uint16BEEncode(data.byteLength)
         yield data
       }
     }

diff --git a/src/encoder.ts b/src/encoder.ts
@@ -2,7 +2,7 @@ import { concat as uint8ArrayConcat } from 'uint8arrays/concat'
 import type { Uint8ArrayList } from 'uint8arraylist'
 import type { bytes } from './@types/basic.js'
 import type { MessageBuffer } from './@types/handshake.js'
-import type { LengthDecoderFunction, LengthEncoderFunction } from 'it-length-prefixed'
+import type { LengthDecoderFunction } from 'it-length-prefixed'
 
 const allocUnsafe = (len: number): Uint8Array => {
   if (globalThis.Buffer) {
@@ -12,7 +12,7 @@ const allocUnsafe = (len: number): Uint8Array => {
   return new Uint8Array(len)
 }
 
-export const uint16BEEncode: LengthEncoderFunction = (value: number) => {
+export const uint16BEEncode = (value: number): Uint8Array => {
   const target = allocUnsafe(2)
   new DataView(target.buffer, target.byteOffset, target.byteLength).setUint16(0, value, false)
   return target

diff --git a/src/noise.ts b/src/noise.ts
@@ -3,7 +3,7 @@ import type { SecuredConnection } from '@libp2p/interface-connection-encrypter'
 import { pbStream, ProtobufStream } from 'it-pb-stream'
 import { duplexPair } from 'it-pair/duplex'
 import { pipe } from 'it-pipe'
-import { encode, decode } from 'it-length-prefixed'
+import { decode } from 'it-length-prefixed'
 import type { Duplex } from 'it-stream-types'
 import type { bytes } from './@types/basic.js'
 import type { IHandshake } from './@types/handshake-interface.js'
@@ -173,8 +173,7 @@ export class Noise implements INoiseConnection {
 
     await pipe(
       secure, // write to wrapper
-      encryptStream(handshake), // data is encrypted
-      encode({ lengthEncoder: uint16BEEncode }), // prefix with message length
+      encryptStream(handshake), // encrypt data + prefix with message length
       network, // send to the remote peer
       decode({ lengthDecoder: uint16BEDecode }), // read message length prefix
       decryptStream(handshake), // decrypt the incoming data