@libp2p/multistream-select 4.0.6-05b52d69c → 4.0.6-13a870cbe

package/src/select.ts

@@ -1,14 +1,11 @@
 import { CodeError } from '@libp2p/interface/errors'
-import { handshake } from 'it-handshake'
-import merge from 'it-merge'
-import { pushable } from 'it-pushable'
-import { reader } from 'it-reader'
-import { Uint8ArrayList } from 'uint8arraylist'
+import { lpStream } from 'it-length-prefixed-stream'
 import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
+import { MAX_PROTOCOL_LENGTH } from './constants.js'
 import * as multistream from './multistream.js'
 import { PROTOCOL_ID } from './index.js'
-import type { ByteArrayInit, ByteListInit, MultistreamSelectInit, ProtocolStream } from './index.js'
-import type { Duplex, Source } from 'it-stream-types'
+import type { MultistreamSelectInit, ProtocolStream } from './index.js'
+import type { Duplex } from 'it-stream-types'
 
 /**
  * Negotiate a protocol to use from a list of protocols.
@@ -53,12 +50,11 @@ import type { Duplex, Source } from 'it-stream-types'
  * // }
  * ```
  */
-export async function select (stream: Duplex<AsyncGenerator<Uint8Array>, Source<Uint8Array>>, protocols: string | string[], options: ByteArrayInit): Promise<ProtocolStream<Uint8Array>>
-export async function select (stream: Duplex<AsyncGenerator<Uint8ArrayList | Uint8Array>, Source<Uint8ArrayList | Uint8Array>>, protocols: string | string[], options?: ByteListInit): Promise<ProtocolStream<Uint8ArrayList, Uint8ArrayList | Uint8Array>>
-export async function select (stream: any, protocols: string | string[], options?: MultistreamSelectInit): Promise<ProtocolStream<any>> {
+export async function select <Stream extends Duplex<any, any, any>> (stream: Stream, protocols: string | string[], options: MultistreamSelectInit): Promise<ProtocolStream<Stream>> {
   protocols = Array.isArray(protocols) ? [...protocols] : [protocols]
-  const { reader, writer, rest, stream: shakeStream } = handshake(stream)
-
+  const lp = lpStream(stream, {
+    maxDataLength: MAX_PROTOCOL_LENGTH
+  })
   const protocol = protocols.shift()
 
   if (protocol == null) {
@@ -66,39 +62,39 @@ export async function select (stream: any, protocols: string | string[], options
   }
 
   options?.log.trace('select: write ["%s", "%s"]', PROTOCOL_ID, protocol)
-  const p1 = uint8ArrayFromString(PROTOCOL_ID)
-  const p2 = uint8ArrayFromString(protocol)
-  multistream.writeAll(writer, [p1, p2], options)
+  const p1 = uint8ArrayFromString(`${PROTOCOL_ID}\n`)
+  const p2 = uint8ArrayFromString(`${protocol}\n`)
+  await multistream.writeAll(lp, [p1, p2], options)
 
-  let response = await multistream.readString(reader, options)
+  options?.log.trace('select: reading multistream-select header')
+  let response = await multistream.readString(lp, options)
   options?.log.trace('select: read "%s"', response)
 
   // Read the protocol response if we got the protocolId in return
   if (response === PROTOCOL_ID) {
-    response = await multistream.readString(reader, options)
+    options?.log.trace('select: reading protocol response')
+    response = await multistream.readString(lp, options)
     options?.log.trace('select: read "%s"', response)
   }
 
   // We're done
   if (response === protocol) {
-    rest()
-    return { stream: shakeStream, protocol }
+    return { stream: lp.unwrap(), protocol }
   }
 
   // We haven't gotten a valid ack, try the other protocols
   for (const protocol of protocols) {
     options?.log.trace('select: write "%s"', protocol)
-    multistream.write(writer, uint8ArrayFromString(protocol), options)
-    const response = await multistream.readString(reader, options)
+    await multistream.write(lp, uint8ArrayFromString(`${protocol}\n`), options)
+    options?.log.trace('select: reading protocol response')
+    const response = await multistream.readString(lp, options)
     options?.log.trace('select: read "%s" for "%s"', response, protocol)
 
     if (response === protocol) {
-      rest() // End our writer so others can start writing to stream
-      return { stream: shakeStream, protocol }
+      return { stream: lp.unwrap(), protocol }
     }
   }
 
-  rest()
   throw new CodeError('protocol selection failed', 'ERR_UNSUPPORTED_PROTOCOL')
 }
 
@@ -110,49 +106,51 @@ export async function select (stream: any, protocols: string | string[], options
  *
  * Use when it is known that the receiver supports the desired protocol.
  */
-export function lazySelect (stream: Duplex<Source<Uint8Array>, Source<Uint8Array>>, protocol: string): ProtocolStream<Uint8Array>
-export function lazySelect (stream: Duplex<Source<Uint8ArrayList | Uint8Array>, Source<Uint8ArrayList | Uint8Array>>, protocol: string): ProtocolStream<Uint8ArrayList, Uint8ArrayList | Uint8Array>
-export function lazySelect (stream: Duplex<any>, protocol: string): ProtocolStream<any> {
-  // This is a signal to write the multistream headers if the consumer tries to
-  // read from the source
-  const negotiateTrigger = pushable()
-  let negotiated = false
+export function lazySelect <Stream extends Duplex<any, any, any>> (stream: Stream, protocol: string, options: MultistreamSelectInit): ProtocolStream<Stream> {
+  const originalSink = stream.sink.bind(stream)
+  const originalSource = stream.source
+
+  const lp = lpStream({
+    sink: originalSink,
+    source: originalSource
+  }, {
+    maxDataLength: MAX_PROTOCOL_LENGTH
+  })
+
+  stream.sink = async source => {
+    options?.log.trace('lazy: write ["%s", "%s"]', PROTOCOL_ID, protocol)
+
+    await lp.writeV([
+      uint8ArrayFromString(`${PROTOCOL_ID}\n`),
+      uint8ArrayFromString(`${protocol}\n`)
+    ])
+
+    options?.log.trace('lazy: writing rest of "%s" stream', protocol)
+    await lp.unwrap().sink(source)
+  }
+
+  stream.source = (async function * () {
+    options?.log.trace('lazy: reading multistream select header')
+
+    let response = await multistream.readString(lp, options)
+    options?.log.trace('lazy: read multistream select header "%s"', response)
+
+    if (response === PROTOCOL_ID) {
+      response = await multistream.readString(lp, options)
+    }
+
+    options?.log.trace('lazy: read protocol "%s", expecting "%s"', response, protocol)
+
+    if (response !== protocol) {
+      throw new CodeError('protocol selection failed', 'ERR_UNSUPPORTED_PROTOCOL')
+    }
+
+    options?.log.trace('lazy: reading rest of "%s" stream', protocol)
+    yield * lp.unwrap().source
+  })()
+
   return {
-    stream: {
-      sink: async source => {
-        await stream.sink((async function * () {
-          let first = true
-          for await (const chunk of merge(source, negotiateTrigger)) {
-            if (first) {
-              first = false
-              negotiated = true
-              negotiateTrigger.end()
-              const p1 = uint8ArrayFromString(PROTOCOL_ID)
-              const p2 = uint8ArrayFromString(protocol)
-              const list = new Uint8ArrayList(multistream.encode(p1), multistream.encode(p2))
-              if (chunk.length > 0) list.append(chunk)
-              yield * list
-            } else {
-              yield chunk
-            }
-          }
-        })())
-      },
-      source: (async function * () {
-        if (!negotiated) negotiateTrigger.push(new Uint8Array())
-        const byteReader = reader(stream.source)
-        let response = await multistream.readString(byteReader)
-        if (response === PROTOCOL_ID) {
-          response = await multistream.readString(byteReader)
-        }
-        if (response !== protocol) {
-          throw new CodeError('protocol selection failed', 'ERR_UNSUPPORTED_PROTOCOL')
-        }
-        for await (const chunk of byteReader) {
-          yield * chunk
-        }
-      })()
-    },
+    stream,
     protocol
   }
 }