@libp2p/mplex 11.0.46-cf9aab5c8 → 11.0.47-0f07e3df5

package/src/mplex.ts

@@ -1,27 +1,16 @@
-import { TooManyOutboundProtocolStreamsError, MuxerClosedError } from '@libp2p/interface'
-import { closeSource } from '@libp2p/utils/close-source'
-import { RateLimiter } from '@libp2p/utils/rate-limiter'
-import { pipe } from 'it-pipe'
-import { pushable } from 'it-pushable'
+import { MuxerClosedError } from '@libp2p/interface'
+import { RateLimiter, AbstractStreamMuxer } from '@libp2p/utils'
 import { toString as uint8ArrayToString } from 'uint8arrays'
-import { Decoder } from './decode.js'
-import { encode } from './encode.js'
-import { StreamInputBufferError } from './errors.js'
+import { Decoder, MAX_MSG_QUEUE_SIZE, MAX_MSG_SIZE } from './decode.js'
 import { MessageTypes, MessageTypeNames } from './message-types.js'
 import { createStream } from './stream.js'
 import type { MplexInit } from './index.js'
 import type { Message } from './message-types.js'
 import type { MplexStream } from './stream.js'
-import type { AbortOptions, ComponentLogger, Logger, Stream, StreamMuxer, StreamMuxerInit } from '@libp2p/interface'
-import type { Pushable } from 'it-pushable'
-import type { Sink, Source } from 'it-stream-types'
+import type { CreateStreamOptions, MultiaddrConnection, MessageStreamDirection } from '@libp2p/interface'
 import type { Uint8ArrayList } from 'uint8arraylist'
 
-const MAX_STREAMS_INBOUND_STREAMS_PER_CONNECTION = 1024
-const MAX_STREAMS_OUTBOUND_STREAMS_PER_CONNECTION = 1024
-const MAX_STREAM_BUFFER_SIZE = 1024 * 1024 * 4 // 4MB
 const DISCONNECT_THRESHOLD = 5
-const CLOSE_TIMEOUT = 500
 
 function printMessage (msg: Message): any {
   const output: any = {
@@ -30,93 +19,34 @@ function printMessage (msg: Message): any {
   }
 
   if (msg.type === MessageTypes.NEW_STREAM) {
-    output.data = uint8ArrayToString(msg.data instanceof Uint8Array ? msg.data : msg.data.subarray())
+    output.data = uint8ArrayToString(msg.data.subarray())
   }
 
   if (msg.type === MessageTypes.MESSAGE_INITIATOR || msg.type === MessageTypes.MESSAGE_RECEIVER) {
-    output.data = uint8ArrayToString(msg.data instanceof Uint8Array ? msg.data : msg.data.subarray(), 'base16')
+    output.data = uint8ArrayToString(msg.data.subarray(), 'base16')
   }
 
   return output
 }
 
-export interface MplexComponents {
-  logger: ComponentLogger
-}
-
-interface MplexStreamMuxerInit extends MplexInit, StreamMuxerInit {
-  /**
-   * The default timeout to use in ms when shutting down the muxer.
-   */
-  closeTimeout?: number
-}
-
-export class MplexStreamMuxer implements StreamMuxer {
-  public protocol = '/mplex/6.7.0'
-
-  public sink: Sink<Source<Uint8ArrayList | Uint8Array>, Promise<void>>
-  public source: AsyncGenerator<Uint8ArrayList | Uint8Array>
-
-  private readonly log: Logger
+export class MplexStreamMuxer extends AbstractStreamMuxer<MplexStream> {
   private _streamId: number
-  private readonly _streams: { initiators: Map<number, MplexStream>, receivers: Map<number, MplexStream> }
-  private readonly _init: MplexStreamMuxerInit
-  private readonly _source: Pushable<Message>
-  private readonly closeController: AbortController
   private readonly rateLimiter: RateLimiter
-  private readonly closeTimeout: number
-  private readonly logger: ComponentLogger
-
-  constructor (components: MplexComponents, init?: MplexStreamMuxerInit) {
-    init = init ?? {}
-
-    this.log = init.log?.newScope('mplex') ?? components.logger.forComponent('libp2p:mplex')
-    this.logger = components.logger
-    this._streamId = 0
-    this._streams = {
-      /**
-       * Stream to ids map
-       */
-      initiators: new Map<number, MplexStream>(),
-      /**
-       * Stream to ids map
-       */
-      receivers: new Map<number, MplexStream>()
-    }
-    this._init = init
-    this.closeTimeout = init.closeTimeout ?? CLOSE_TIMEOUT
-
-    /**
-     * An iterable sink
-     */
-    this.sink = this._createSink()
-
-    /**
-     * An iterable source
-     */
-    this._source = pushable<Message>({
-      objectMode: true,
-      onEnd: (): void => {
-        // the source has ended, we can't write any more messages to gracefully
-        // close streams so all we can do is destroy them
-        for (const stream of this._streams.initiators.values()) {
-          stream.destroy()
-        }
-
-        for (const stream of this._streams.receivers.values()) {
-          stream.destroy()
-        }
-      }
+  private readonly maxMessageSize: number
+  private readonly maxUnprocessedMessageQueueSize: number
+  private readonly decoder: Decoder
+
+  constructor (maConn: MultiaddrConnection, init: MplexInit) {
+    super(maConn, {
+      ...init,
+      protocol: '/mplex/6.7.0',
+      name: 'mplex'
     })
-    this.source = pipe(
-      this._source,
-      source => encode(source)
-    )
 
-    /**
-     * Close controller
-     */
-    this.closeController = new AbortController()
+    this._streamId = 0
+    this.maxMessageSize = init.maxMessageSize ?? MAX_MSG_SIZE
+    this.maxUnprocessedMessageQueueSize = init.maxUnprocessedMessageQueueSize ?? MAX_MSG_QUEUE_SIZE
+    this.decoder = new Decoder(this.maxMessageSize, this.maxUnprocessedMessageQueueSize)
 
     this.rateLimiter = new RateLimiter({
       points: init.disconnectThreshold ?? DISCONNECT_THRESHOLD,
@@ -124,207 +54,75 @@ export class MplexStreamMuxer implements StreamMuxer {
     })
   }
 
-  /**
-   * Returns a Map of streams and their ids
-   */
-  get streams (): Stream[] {
-    // Inbound and Outbound streams may have the same ids, so we need to make those unique
-    const streams: Stream[] = []
-    for (const stream of this._streams.initiators.values()) {
-      streams.push(stream)
-    }
-
-    for (const stream of this._streams.receivers.values()) {
-      streams.push(stream)
+  onData (data: Uint8Array | Uint8ArrayList): void {
+    for (const msg of this.decoder.write(data)) {
+      this.handleMessage(msg)
     }
-    return streams
   }
 
   /**
    * Initiate a new stream with the given name. If no name is
    * provided, the id of the stream will be used.
    */
-  newStream (name?: string): Stream {
-    if (this.closeController.signal.aborted) {
+  onCreateStream (options: CreateStreamOptions): MplexStream {
+    if (this.status !== 'open') {
       throw new MuxerClosedError('Muxer already closed')
     }
-    const id = this._streamId++
-    name = name == null ? id.toString() : name.toString()
-    const registry = this._streams.initiators
-    return this._newStream({ id, name, type: 'initiator', registry })
-  }
-
-  /**
-   * Close or abort all tracked streams and stop the muxer
-   */
-  async close (options?: AbortOptions): Promise<void> {
-    if (this.closeController.signal.aborted) {
-      return
-    }
-
-    const signal = options?.signal ?? AbortSignal.timeout(this.closeTimeout)
-
-    try {
-      // try to gracefully close all streams
-      await Promise.all(
-        this.streams.map(async s => s.close({
-          signal
-        }))
-      )
 
-      this._source.end()
-
-      // try to gracefully close the muxer
-      await this._source.onEmpty({
-        signal
-      })
-
-      this.closeController.abort()
-    } catch (err: any) {
-      this.abort(err)
-    }
-  }
-
-  abort (err: Error): void {
-    if (this.closeController.signal.aborted) {
-      return
-    }
-
-    this.streams.forEach(s => { s.abort(err) })
-    this.closeController.abort(err)
-  }
+    const id = this._streamId++
 
-  /**
-   * Called whenever an inbound stream is created
-   */
-  _newReceiverStream (options: { id: number, name: string }): MplexStream {
-    const { id, name } = options
-    const registry = this._streams.receivers
-    return this._newStream({ id, name, type: 'receiver', registry })
+    return this._newStream(id, 'outbound', options)
   }
 
-  _newStream (options: { id: number, name: string, type: 'initiator' | 'receiver', registry: Map<number, MplexStream> }): MplexStream {
-    const { id, name, type, registry } = options
-
-    this.log('new %s stream %s', type, id)
-
-    if (type === 'initiator' && this._streams.initiators.size === (this._init.maxOutboundStreams ?? MAX_STREAMS_OUTBOUND_STREAMS_PER_CONNECTION)) {
-      throw new TooManyOutboundProtocolStreamsError('Too many outbound streams open')
-    }
+  _newStream (id: number, direction: MessageStreamDirection, options?: CreateStreamOptions): MplexStream {
+    this.log('new %s stream %s', direction, id)
 
-    if (registry.has(id)) {
-      throw new Error(`${type} stream ${id} already exists!`)
-    }
-
-    const send = async (msg: Message): Promise<void> => {
-      if (this.log.enabled) {
-        this.log.trace('%s stream %s send', type, id, printMessage(msg))
-      }
-
-      this._source.push(msg)
-    }
-
-    const onEnd = (): void => {
-      this.log('%s stream with id %s and protocol %s ended', type, id, stream.protocol)
-      registry.delete(id)
-
-      if (this._init.onStreamEnd != null) {
-        this._init.onStreamEnd(stream)
-      }
-    }
+    const stream = createStream({
+      ...options,
+      id,
+      direction,
+      maxMsgSize: this.maxMessageSize,
+      log: this.log,
+      muxer: this
+    })
 
-    const stream = createStream({ id, name, send, type, onEnd, maxMsgSize: this._init.maxMsgSize, log: this.log })
-    registry.set(id, stream)
     return stream
   }
 
-  /**
-   * Creates a sink with an abortable source. Incoming messages will
-   * also have their size restricted. All messages will be varint decoded.
-   */
-  _createSink (): Sink<Source<Uint8ArrayList | Uint8Array>, Promise<void>> {
-    const sink: Sink<Source<Uint8ArrayList | Uint8Array>, Promise<void>> = async source => {
-      const abortListener = (): void => {
-        closeSource(source, this.log)
-      }
-
-      this.closeController.signal.addEventListener('abort', abortListener)
-
-      try {
-        const decoder = new Decoder(this._init.maxMsgSize, this._init.maxUnprocessedMessageQueueSize)
-
-        for await (const chunk of source) {
-          for (const msg of decoder.write(chunk)) {
-            await this._handleIncoming(msg)
-          }
-        }
-
-        this._source.end()
-      } catch (err: any) {
-        this.log('error in sink', err)
-        this._source.end(err) // End the source with an error
-      } finally {
-        this.closeController.signal.removeEventListener('abort', abortListener)
-      }
-    }
-
-    return sink
-  }
-
-  async _handleIncoming (message: Message): Promise<void> {
-    const { id, type } = message
-
+  handleMessage (message: Message): void {
     if (this.log.enabled) {
       this.log.trace('incoming message', printMessage(message))
     }
 
     // Create a new stream?
     if (message.type === MessageTypes.NEW_STREAM) {
-      if (this._streams.receivers.size === (this._init.maxInboundStreams ?? MAX_STREAMS_INBOUND_STREAMS_PER_CONNECTION)) {
-        this.log('too many inbound streams open')
-
-        // not going to allow this stream, send the reset message manually
-        // instead of setting it up just to tear it down
-        this._source.push({
-          id,
-          type: MessageTypes.RESET_RECEIVER
-        })
-
-        // if we've hit our stream limit, and the remote keeps trying to open
-        // more new streams, if they are doing this very quickly maybe they
-        // are attacking us and we should close the connection
-        try {
-          await this.rateLimiter.consume('new-stream', 1)
-        } catch {
-          this.log('rate limit hit when opening too many new streams over the inbound stream limit - closing remote connection')
-          // since there's no backpressure in mplex, the only thing we can really do to protect ourselves is close the connection
-          this.abort(new Error('Too many open streams'))
-          return
-        }
-
+      // close the connection if the remote opens too many streams too quickly
+      try {
+        this.rateLimiter.consume('new-stream', 1)
+      } catch {
+        this.log('rate limit hit when opening too many new streams over the inbound stream limit - closing remote connection')
+        // since there's no backpressure in mplex, the only thing we can really do to protect ourselves is close the connection
+        this.abort(new Error('Too many open streams'))
         return
       }
 
-      const stream = this._newReceiverStream({ id, name: uint8ArrayToString(message.data instanceof Uint8Array ? message.data : message.data.subarray()) })
-
-      if (this._init.onIncomingStream != null) {
-        this._init.onIncomingStream(stream)
-      }
+      const stream = this._newStream(message.id, 'inbound', this.streamOptions)
+      this.onRemoteStream(stream)
 
       return
     }
 
-    const list = (type & 1) === 1 ? this._streams.initiators : this._streams.receivers
-    const stream = list.get(id)
+    const id = `${(message.type & 1) === 1 ? 'i' : 'r'}${message.id}`
+    const stream = this.streams.find(s => s.id === id)
 
     if (stream == null) {
-      this.log('missing stream %s for message type %s', id, MessageTypeNames[type])
+      this.log('missing stream %s for message type %s', id, MessageTypeNames[message.type])
 
       // if the remote keeps sending us messages for streams that have been
       // closed or were never opened they may be attacking us so if they do
       // this very quickly all we can do is close the connection
       try {
-        await this.rateLimiter.consume('missing-stream', 1)
+        this.rateLimiter.consume('missing-stream', 1)
       } catch {
         this.log('rate limit hit when receiving messages for streams that do not exist - closing remote connection')
         // since there's no backpressure in mplex, the only thing we can really do to protect ourselves is close the connection
@@ -335,41 +133,28 @@ export class MplexStreamMuxer implements StreamMuxer {
       return
     }
 
-    const maxBufferSize = this._init.maxStreamBufferSize ?? MAX_STREAM_BUFFER_SIZE
-
     try {
-      switch (type) {
+      switch (message.type) {
         case MessageTypes.MESSAGE_INITIATOR:
         case MessageTypes.MESSAGE_RECEIVER:
-          if (stream.sourceReadableLength() > maxBufferSize) {
-            // Stream buffer has got too large, reset the stream
-            this._source.push({
-              id: message.id,
-              type: type === MessageTypes.MESSAGE_INITIATOR ? MessageTypes.RESET_RECEIVER : MessageTypes.RESET_INITIATOR
-            })
-
-            // Inform the stream consumer they are not fast enough
-            throw new StreamInputBufferError('Input buffer full - increase Mplex maxBufferSize to accommodate slow consumers')
-          }
-
           // We got data from the remote, push it into our local stream
-          stream.sourcePush(message.data)
+          stream.onData(message.data)
           break
         case MessageTypes.CLOSE_INITIATOR:
         case MessageTypes.CLOSE_RECEIVER:
-          // The remote has stopped writing, so we can stop reading
-          stream.remoteCloseWrite()
+          // The remote has stopped writing
+          stream.onRemoteCloseWrite()
           break
         case MessageTypes.RESET_INITIATOR:
         case MessageTypes.RESET_RECEIVER:
           // The remote has errored, stop reading and writing to the stream immediately
-          stream.reset()
+          stream.onRemoteReset()
           break
         default:
-          this.log('unknown message type %s', type)
+          this.log('unknown message type')
       }
     } catch (err: any) {
-      this.log.error('error while processing message', err)
+      this.log.error('error while processing message - %e', err)
       stream.abort(err)
     }
   }

package/src/stream.ts

@@ -1,95 +1,127 @@
-import { AbstractStream } from '@libp2p/utils/abstract-stream'
+import { AbstractStream } from '@libp2p/utils'
 import { Uint8ArrayList } from 'uint8arraylist'
 import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
 import { MAX_MSG_SIZE } from './decode.js'
+import { encode } from './encode.ts'
 import { InitiatorMessageTypes, ReceiverMessageTypes } from './message-types.js'
-import type { Message } from './message-types.js'
-import type { Logger } from '@libp2p/interface'
-import type { AbstractStreamInit } from '@libp2p/utils/abstract-stream'
+import type { MplexStreamMuxer } from './mplex.ts'
+import type { Logger, MessageStreamDirection } from '@libp2p/interface'
+import type { AbstractStreamInit, SendResult } from '@libp2p/utils'
+import type { AbortOptions } from 'it-pushable'
 
 export interface Options {
   id: number
-  send(msg: Message): Promise<void>
   log: Logger
-  name?: string
-  onEnd?(err?: Error): void
-  type?: 'initiator' | 'receiver'
+  direction: MessageStreamDirection
   maxMsgSize?: number
+  muxer: MplexStreamMuxer
 }
 
 interface MplexStreamInit extends AbstractStreamInit {
-  streamId: number
-  name: string
-  send(msg: Message): Promise<void>
-
-  /**
-   * The maximum allowable data size, any data larger than this will be
-   * chunked and sent in multiple data messages
-   */
   maxDataSize: number
+  muxer: MplexStreamMuxer
+  direction: MessageStreamDirection
 }
 
 export class MplexStream extends AbstractStream {
-  private readonly name: string
-  private readonly streamId: number
-  private readonly send: (msg: Message) => Promise<void>
+  public readonly streamId: number
   private readonly types: Record<string, number>
   private readonly maxDataSize: number
+  private readonly muxer: MplexStreamMuxer
 
   constructor (init: MplexStreamInit) {
     super(init)
 
     this.types = init.direction === 'outbound' ? InitiatorMessageTypes : ReceiverMessageTypes
-    this.send = init.send
-    this.name = init.name
-    this.streamId = init.streamId
     this.maxDataSize = init.maxDataSize
-  }
+    this.muxer = init.muxer
+    this.streamId = parseInt(this.id.substring(1))
 
-  async sendNewStream (): Promise<void> {
-    await this.send({ id: this.streamId, type: InitiatorMessageTypes.NEW_STREAM, data: new Uint8ArrayList(uint8ArrayFromString(this.name)) })
+    if (init.direction === 'outbound') {
+      // open the stream on the receiver end. do this in a microtask so the
+      // stream gets added to the streams array by the muxer superclass before
+      // we send the NEW_STREAM message, otherwise we create a race condition
+      // whereby we can receive the stream messages before the stream is added
+      // to the streams list
+      queueMicrotask(() => {
+        this.muxer.send(
+          encode({
+            id: this.streamId,
+            type: InitiatorMessageTypes.NEW_STREAM,
+            data: new Uint8ArrayList(uint8ArrayFromString(this.id))
+          })
+        )
+      })
+    }
   }
 
-  async sendData (data: Uint8ArrayList): Promise<void> {
-    data = data.sublist()
+  sendData (data: Uint8ArrayList): SendResult {
+    const list = new Uint8ArrayList()
+    const sentBytes = data.byteLength
 
     while (data.byteLength > 0) {
       const toSend = Math.min(data.byteLength, this.maxDataSize)
-      await this.send({
-        id: this.streamId,
-        type: this.types.MESSAGE,
-        data: data.sublist(0, toSend)
-      })
+      const slice = data.sublist(0, toSend)
+      data = data.sublist(toSend)
+
+      list.append(
+        encode({
+          id: this.streamId,
+          type: this.types.MESSAGE,
+          data: slice
+        })
+      )
+    }
 
-      data.consume(toSend)
+    return {
+      sentBytes,
+      canSendMore: this.muxer.send(list)
     }
   }
 
-  async sendReset (): Promise<void> {
-    await this.send({ id: this.streamId, type: this.types.RESET })
+  sendReset (): boolean {
+    return this.muxer.send(
+      encode({
+        id: this.streamId,
+        type: this.types.RESET
+      })
+    )
   }
 
-  async sendCloseWrite (): Promise<void> {
-    await this.send({ id: this.streamId, type: this.types.CLOSE })
+  async sendCloseWrite (options?: AbortOptions): Promise<void> {
+    this.muxer.send(
+      encode({
+        id: this.streamId,
+        type: this.types.CLOSE
+      })
+    )
+    options?.signal?.throwIfAborted()
   }
 
-  async sendCloseRead (): Promise<void> {
+  async sendCloseRead (options?: AbortOptions): Promise<void> {
+    options?.signal?.throwIfAborted()
     // mplex does not support close read, only close write
   }
+
+  sendPause (): void {
+    // mplex does not support backpressure
+  }
+
+  sendResume (): void {
+    // mplex does not support backpressure
+  }
 }
 
 export function createStream (options: Options): MplexStream {
-  const { id, name, send, onEnd, type = 'initiator', maxMsgSize = MAX_MSG_SIZE } = options
-  const direction = type === 'initiator' ? 'outbound' : 'inbound'
+  const { id, muxer, direction, maxMsgSize = MAX_MSG_SIZE } = options
 
   return new MplexStream({
-    id: type === 'initiator' ? (`i${id}`) : `r${id}`,
-    streamId: id,
-    name: `${name ?? id}`,
+    ...options,
+    id: direction === 'outbound' ? (`i${id}`) : `r${id}`,
     direction,
     maxDataSize: maxMsgSize,
-    onEnd,
-    send,
-    log: options.log.newScope(`${direction}:${id}`)
+    muxer,
+    log: options.log.newScope(`${direction}:${id}`),
+    protocol: ''
   })
 }