@libp2p/multistream-select 4.0.10 → 5.0.0-10ea19700

package/src/multistream.ts

@@ -1,98 +1,45 @@
-import { CodeError } from '@libp2p/interface/errors'
-import { logger } from '@libp2p/logger'
-import { abortableSource } from 'abortable-iterator'
-import first from 'it-first'
-import * as lp from 'it-length-prefixed'
-import { pipe } from 'it-pipe'
-import { Uint8ArrayList } from 'uint8arraylist'
+import { CodeError } from '@libp2p/interface'
+import { type Uint8ArrayList } from 'uint8arraylist'
 import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
 import { toString as uint8ArrayToString } from 'uint8arrays/to-string'
-import { MAX_PROTOCOL_LENGTH } from './constants.js'
-import type { MultistreamSelectInit } from '.'
-import type { AbortOptions } from '@libp2p/interface'
-import type { Pushable } from 'it-pushable'
-import type { Reader } from 'it-reader'
-import type { Source } from 'it-stream-types'
-
-const log = logger('libp2p:mss')
+import type { AbortOptions, LoggerOptions } from '@libp2p/interface'
+import type { LengthPrefixedStream } from 'it-length-prefixed-stream'
+import type { Duplex, Source } from 'it-stream-types'
 
 const NewLine = uint8ArrayFromString('\n')
 
-export function encode (buffer: Uint8Array | Uint8ArrayList): Uint8ArrayList {
-  const list = new Uint8ArrayList(buffer, NewLine)
-
-  return lp.encode.single(list)
-}
-
 /**
  * `write` encodes and writes a single buffer
  */
-export function write (writer: Pushable<any>, buffer: Uint8Array | Uint8ArrayList, options: MultistreamSelectInit = {}): void {
-  const encoded = encode(buffer)
-
-  if (options.writeBytes === true) {
-    writer.push(encoded.subarray())
-  } else {
-    writer.push(encoded)
-  }
+export async function write (writer: LengthPrefixedStream<Duplex<AsyncGenerator<Uint8Array | Uint8ArrayList>, Source<Uint8Array>>>, buffer: Uint8Array | Uint8ArrayList, options?: AbortOptions): Promise<void> {
+  await writer.write(buffer, options)
 }
 
 /**
  * `writeAll` behaves like `write`, except it encodes an array of items as a single write
  */
-export function writeAll (writer: Pushable<any>, buffers: Uint8Array[], options: MultistreamSelectInit = {}): void {
-  const list = new Uint8ArrayList()
-
-  for (const buf of buffers) {
-    list.append(encode(buf))
-  }
-
-  if (options.writeBytes === true) {
-    writer.push(list.subarray())
-  } else {
-    writer.push(list)
-  }
+export async function writeAll (writer: LengthPrefixedStream<Duplex<AsyncGenerator<Uint8Array | Uint8ArrayList>, Source<Uint8Array>>>, buffers: Uint8Array[], options?: AbortOptions): Promise<void> {
+  await writer.writeV(buffers, options)
 }
 
-export async function read (reader: Reader, options?: AbortOptions): Promise<Uint8ArrayList> {
-  let byteLength = 1 // Read single byte chunks until the length is known
-  const varByteSource = { // No return impl - we want the reader to remain readable
-    [Symbol.asyncIterator]: () => varByteSource,
-    next: async () => reader.next(byteLength)
-  }
-
-  let input: Source<Uint8ArrayList> = varByteSource
-
-  // If we have been passed an abort signal, wrap the input source in an abortable
-  // iterator that will throw if the operation is aborted
-  if (options?.signal != null) {
-    input = abortableSource(varByteSource, options.signal)
-  }
-
-  // Once the length has been parsed, read chunk for that length
-  const onLength = (l: number): void => {
-    byteLength = l
-  }
-
-  const buf = await pipe(
-    input,
-    (source) => lp.decode(source, { onLength, maxDataLength: MAX_PROTOCOL_LENGTH }),
-    async (source) => first(source)
-  )
-
-  if (buf == null || buf.length === 0) {
-    throw new CodeError('no buffer returned', 'ERR_INVALID_MULTISTREAM_SELECT_MESSAGE')
-  }
+/**
+ * Read a length-prefixed buffer from the passed stream, stripping the final newline character
+ */
+export async function read (reader: LengthPrefixedStream<Duplex<AsyncGenerator<Uint8Array | Uint8ArrayList>, Source<Uint8Array>>>, options: AbortOptions & LoggerOptions): Promise<Uint8ArrayList> {
+  const buf = await reader.read(options)
 
-  if (buf.get(buf.byteLength - 1) !== NewLine[0]) {
-    log.error('Invalid mss message - missing newline - %s', buf.subarray())
+  if (buf.byteLength === 0 || buf.get(buf.byteLength - 1) !== NewLine[0]) {
+    options.log.error('Invalid mss message - missing newline', buf)
     throw new CodeError('missing newline', 'ERR_INVALID_MULTISTREAM_SELECT_MESSAGE')
   }
 
   return buf.sublist(0, -1) // Remove newline
 }
 
-export async function readString (reader: Reader, options?: AbortOptions): Promise<string> {
+/**
+ * Read a length-prefixed string from the passed stream, stripping the final newline character
+ */
+export async function readString (reader: LengthPrefixedStream<Duplex<AsyncGenerator<Uint8Array | Uint8ArrayList>, Source<Uint8Array>>>, options: AbortOptions & LoggerOptions): Promise<string> {
   const buf = await read(reader, options)
 
   return uint8ArrayToString(buf.subarray())

package/src/select.ts

@@ -1,17 +1,22 @@
-import { CodeError } from '@libp2p/interface/errors'
-import { logger } from '@libp2p/logger'
-import { handshake } from 'it-handshake'
-import merge from 'it-merge'
-import { pushable } from 'it-pushable'
-import { reader } from 'it-reader'
+import { CodeError } from '@libp2p/interface'
+import { lpStream } from 'it-length-prefixed-stream'
+import pDefer from 'p-defer'
+import * as varint from 'uint8-varint'
 import { Uint8ArrayList } from 'uint8arraylist'
 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 { AbortOptions } from '@libp2p/interface'
+import type { Duplex } from 'it-stream-types'
 
-const log = logger('libp2p:mss:select')
+export interface SelectStream extends Duplex<any, any, any> {
+  readStatus?: string
+  closeWrite?(options?: AbortOptions): Promise<void>
+  closeRead?(options?: AbortOptions): Promise<void>
+  close?(options?: AbortOptions): Promise<void>
+}
 
 /**
  * Negotiate a protocol to use from a list of protocols.
@@ -56,106 +61,281 @@ const log = logger('libp2p:mss:select')
  * // }
  * ```
  */
-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 SelectStream> (stream: Stream, protocols: string | string[], options: MultistreamSelectInit): Promise<ProtocolStream<Stream>> {
   protocols = Array.isArray(protocols) ? [...protocols] : [protocols]
-  const { reader, writer, rest, stream: shakeStream } = handshake(stream)
 
+  if (protocols.length === 1) {
+    return optimisticSelect(stream, protocols[0], options)
+  }
+
+  const lp = lpStream(stream, {
+    ...options,
+    maxDataLength: MAX_PROTOCOL_LENGTH
+  })
   const protocol = protocols.shift()
 
   if (protocol == null) {
     throw new Error('At least one protocol must be specified')
   }
 
-  log.trace('select: write ["%s", "%s"]', PROTOCOL_ID, protocol)
-  const p1 = uint8ArrayFromString(PROTOCOL_ID)
-  const p2 = uint8ArrayFromString(protocol)
-  multistream.writeAll(writer, [p1, p2], options)
+  options.log.trace('select: write ["%s", "%s"]', PROTOCOL_ID, protocol)
+  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)
-  log.trace('select: read "%s"', response)
+  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)
-    log.trace('select: read "%s"', response)
+    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) {
-    log.trace('select: write "%s"', protocol)
-    multistream.write(writer, uint8ArrayFromString(protocol), options)
-    const response = await multistream.readString(reader, options)
-    log.trace('select: read "%s" for "%s"', response, protocol)
+    options.log.trace('select: write "%s"', protocol)
+    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')
 }
 
 /**
- * Lazily negotiates a protocol.
+ * Optimistically negotiates a protocol.
  *
  * It *does not* block writes waiting for the other end to respond. Instead, it
  * simply assumes the negotiation went successfully and starts writing data.
  *
  * 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()
+function optimisticSelect <Stream extends SelectStream> (stream: Stream, protocol: string, options: MultistreamSelectInit): ProtocolStream<Stream> {
+  const originalSink = stream.sink.bind(stream)
+  const originalSource = stream.source
+
   let negotiated = false
-  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')
+  let negotiating = false
+  const doneNegotiating = pDefer()
+
+  let sentProtocol = false
+  let sendingProtocol = false
+  const doneSendingProtocol = pDefer()
+
+  let readProtocol = false
+  let readingProtocol = false
+  const doneReadingProtocol = pDefer()
+
+  const lp = lpStream({
+    sink: originalSink,
+    source: originalSource
+  }, {
+    ...options,
+    maxDataLength: MAX_PROTOCOL_LENGTH
+  })
+
+  stream.sink = async source => {
+    const { sink } = lp.unwrap()
+
+    await sink(async function * () {
+      let sentData = false
+
+      for await (const buf of source) {
+        // started reading before the source yielded, wait for protocol send
+        if (sendingProtocol) {
+          await doneSendingProtocol.promise
         }
-        for await (const chunk of byteReader) {
-          yield * chunk
+
+        // writing before reading, send the protocol and the first chunk of data
+        if (!sentProtocol) {
+          sendingProtocol = true
+
+          options.log.trace('optimistic: write ["%s", "%s", data(%d)] in sink', PROTOCOL_ID, protocol, buf.byteLength)
+
+          const protocolString = `${protocol}\n`
+
+          // send protocols in first chunk of data written to transport
+          yield new Uint8ArrayList(
+            Uint8Array.from([19]), // length of PROTOCOL_ID plus newline
+            uint8ArrayFromString(`${PROTOCOL_ID}\n`),
+            varint.encode(protocolString.length),
+            uint8ArrayFromString(protocolString),
+            buf
+          ).subarray()
+
+          options.log.trace('optimistic: wrote ["%s", "%s", data(%d)] in sink', PROTOCOL_ID, protocol, buf.byteLength)
+
+          sentProtocol = true
+          sendingProtocol = false
+          doneSendingProtocol.resolve()
+        } else {
+          yield buf
         }
-      })()
-    },
+
+        sentData = true
+      }
+
+      // special case - the source passed to the sink has ended but we didn't
+      // negotiated the protocol yet so do it now
+      if (!sentData) {
+        await negotiate()
+      }
+    }())
+  }
+
+  async function negotiate (): Promise<void> {
+    if (negotiating) {
+      options.log.trace('optimistic: already negotiating %s stream', protocol)
+      await doneNegotiating.promise
+      return
+    }
+
+    negotiating = true
+
+    try {
+      // we haven't sent the protocol yet, send it now
+      if (!sentProtocol) {
+        options.log.trace('optimistic: doing send protocol for %s stream', protocol)
+        await doSendProtocol()
+      }
+
+      // if we haven't read the protocol response yet, do it now
+      if (!readProtocol) {
+        options.log.trace('optimistic: doing read protocol for %s stream', protocol)
+        await doReadProtocol()
+      }
+    } finally {
+      negotiating = false
+      negotiated = true
+      doneNegotiating.resolve()
+    }
+  }
+
+  async function doSendProtocol (): Promise<void> {
+    if (sendingProtocol) {
+      await doneSendingProtocol.promise
+      return
+    }
+
+    sendingProtocol = true
+
+    try {
+      options.log.trace('optimistic: write ["%s", "%s", data] in source', PROTOCOL_ID, protocol)
+      await lp.writeV([
+        uint8ArrayFromString(`${PROTOCOL_ID}\n`),
+        uint8ArrayFromString(`${protocol}\n`)
+      ])
+      options.log.trace('optimistic: wrote ["%s", "%s", data] in source', PROTOCOL_ID, protocol)
+    } finally {
+      sentProtocol = true
+      sendingProtocol = false
+      doneSendingProtocol.resolve()
+    }
+  }
+
+  async function doReadProtocol (): Promise<void> {
+    if (readingProtocol) {
+      await doneReadingProtocol.promise
+      return
+    }
+
+    readingProtocol = true
+
+    try {
+      options.log.trace('optimistic: reading multistream select header')
+      let response = await multistream.readString(lp, options)
+      options.log.trace('optimistic: read multistream select header "%s"', response)
+
+      if (response === PROTOCOL_ID) {
+        response = await multistream.readString(lp, options)
+      }
+
+      options.log.trace('optimistic: read protocol "%s", expecting "%s"', response, protocol)
+
+      if (response !== protocol) {
+        throw new CodeError('protocol selection failed', 'ERR_UNSUPPORTED_PROTOCOL')
+      }
+    } finally {
+      readProtocol = true
+      readingProtocol = false
+      doneReadingProtocol.resolve()
+    }
+  }
+
+  stream.source = (async function * () {
+    // make sure we've done protocol negotiation before we read stream data
+    await negotiate()
+
+    options.log.trace('optimistic: reading data from "%s" stream', protocol)
+    yield * lp.unwrap().source
+  })()
+
+  if (stream.closeRead != null) {
+    const originalCloseRead = stream.closeRead.bind(stream)
+
+    stream.closeRead = async (opts) => {
+      // we need to read & write to negotiate the protocol so ensure we've done
+      // this before closing the readable end of the stream
+      if (!negotiated) {
+        await negotiate().catch(err => {
+          options.log.error('could not negotiate protocol before close read', err)
+        })
+      }
+
+      // protocol has been negotiated, ok to close the readable end
+      await originalCloseRead(opts)
+    }
+  }
+
+  if (stream.closeWrite != null) {
+    const originalCloseWrite = stream.closeWrite.bind(stream)
+
+    stream.closeWrite = async (opts) => {
+      // we need to read & write to negotiate the protocol so ensure we've done
+      // this before closing the writable end of the stream
+      if (!negotiated) {
+        await negotiate().catch(err => {
+          options.log.error('could not negotiate protocol before close write', err)
+        })
+      }
+
+      // protocol has been negotiated, ok to close the writable end
+      await originalCloseWrite(opts)
+    }
+  }
+
+  if (stream.close != null) {
+    const originalClose = stream.close.bind(stream)
+
+    stream.close = async (opts) => {
+      // the stream is being closed, don't try to negotiate a protocol if we
+      // haven't already
+      if (!negotiated) {
+        negotiated = true
+        negotiating = false
+        doneNegotiating.resolve()
+      }
+
+      // protocol has been negotiated, ok to close the writable end
+      await originalClose(opts)
+    }
+  }
+
+  return {
+    stream,
     protocol
   }
 }