@libp2p/mplex 11.0.47-8484de8a2 → 11.0.47-87bc8d4fb

package/src/mplex.ts

@@ -1,16 +1,27 @@
-import { MuxerClosedError } from '@libp2p/interface'
-import { RateLimiter, AbstractStreamMuxer } from '@libp2p/utils'
+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 { toString as uint8ArrayToString } from 'uint8arrays'
-import { Decoder, MAX_MSG_QUEUE_SIZE, MAX_MSG_SIZE } from './decode.js'
+import { Decoder } from './decode.js'
+import { encode } from './encode.js'
+import { StreamInputBufferError } from './errors.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 { CreateStreamOptions, MultiaddrConnection, MessageStreamDirection } from '@libp2p/interface'
+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 { 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 = {
@@ -19,34 +30,93 @@ function printMessage (msg: Message): any {
   }
 
   if (msg.type === MessageTypes.NEW_STREAM) {
-    output.data = uint8ArrayToString(msg.data.subarray())
+    output.data = uint8ArrayToString(msg.data instanceof Uint8Array ? msg.data : msg.data.subarray())
   }
 
   if (msg.type === MessageTypes.MESSAGE_INITIATOR || msg.type === MessageTypes.MESSAGE_RECEIVER) {
-    output.data = uint8ArrayToString(msg.data.subarray(), 'base16')
+    output.data = uint8ArrayToString(msg.data instanceof Uint8Array ? msg.data : msg.data.subarray(), 'base16')
   }
 
   return output
 }
 
-export class MplexStreamMuxer extends AbstractStreamMuxer<MplexStream> {
+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
   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 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'
-    })
+  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.maxMessageSize = init.maxMessageSize ?? MAX_MSG_SIZE
-    this.maxUnprocessedMessageQueueSize = init.maxUnprocessedMessageQueueSize ?? MAX_MSG_QUEUE_SIZE
-    this.decoder = new Decoder(this.maxMessageSize, this.maxUnprocessedMessageQueueSize)
+    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()
+        }
+      }
+    })
+    this.source = pipe(
+      this._source,
+      source => encode(source)
+    )
+
+    /**
+     * Close controller
+     */
+    this.closeController = new AbortController()
 
     this.rateLimiter = new RateLimiter({
       points: init.disconnectThreshold ?? DISCONNECT_THRESHOLD,
@@ -54,75 +124,207 @@ export class MplexStreamMuxer extends AbstractStreamMuxer<MplexStream> {
     })
   }
 
-  onData (data: Uint8Array | Uint8ArrayList): void {
-    for (const msg of this.decoder.write(data)) {
-      this.handleMessage(msg)
+  /**
+   * 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)
     }
+    return streams
   }
 
   /**
    * Initiate a new stream with the given name. If no name is
    * provided, the id of the stream will be used.
    */
-  onCreateStream (options: CreateStreamOptions): MplexStream {
-    if (this.status !== 'open') {
+  newStream (name?: string): Stream {
+    if (this.closeController.signal.aborted) {
       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
+        }))
+      )
 
-    return this._newStream(id, 'outbound', options)
+      this._source.end()
+
+      // try to gracefully close the muxer
+      await this._source.onEmpty({
+        signal
+      })
+
+      this.closeController.abort()
+    } catch (err: any) {
+      this.abort(err)
+    }
   }
 
-  _newStream (id: number, direction: MessageStreamDirection, options?: CreateStreamOptions): MplexStream {
-    this.log('new %s stream %s', direction, id)
+  abort (err: Error): void {
+    if (this.closeController.signal.aborted) {
+      return
+    }
+
+    this.streams.forEach(s => { s.abort(err) })
+    this.closeController.abort(err)
+  }
 
-    const stream = createStream({
-      ...options,
-      id,
-      direction,
-      maxMsgSize: this.maxMessageSize,
-      log: this.log,
-      muxer: this
-    })
+  /**
+   * 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 })
+  }
+
+  _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')
+    }
 
+    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({ id, name, send, type, onEnd, maxMsgSize: this._init.maxMsgSize, log: this.log })
+    registry.set(id, stream)
     return stream
   }
 
-  handleMessage (message: Message): void {
+  /**
+   * 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
+
     if (this.log.enabled) {
       this.log.trace('incoming message', printMessage(message))
     }
 
     // Create a new stream?
     if (message.type === MessageTypes.NEW_STREAM) {
-      // 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'))
+      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
+        }
+
         return
       }
 
-      const stream = this._newStream(message.id, 'inbound', this.streamOptions)
-      this.onRemoteStream(stream)
+      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)
+      }
 
       return
     }
 
-    const id = `${(message.type & 1) === 1 ? 'i' : 'r'}${message.id}`
-    const stream = this.streams.find(s => s.id === id)
+    const list = (type & 1) === 1 ? this._streams.initiators : this._streams.receivers
+    const stream = list.get(id)
 
     if (stream == null) {
-      this.log('missing stream %s for message type %s', id, MessageTypeNames[message.type])
+      this.log('missing stream %s for message type %s', id, MessageTypeNames[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 {
-        this.rateLimiter.consume('missing-stream', 1)
+        await 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
@@ -133,28 +335,41 @@ export class MplexStreamMuxer extends AbstractStreamMuxer<MplexStream> {
       return
     }
 
+    const maxBufferSize = this._init.maxStreamBufferSize ?? MAX_STREAM_BUFFER_SIZE
+
     try {
-      switch (message.type) {
+      switch (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.onData(message.data)
+          stream.sourcePush(message.data)
           break
         case MessageTypes.CLOSE_INITIATOR:
         case MessageTypes.CLOSE_RECEIVER:
-          // The remote has stopped writing
-          stream.onRemoteCloseWrite()
+          // The remote has stopped writing, so we can stop reading
+          stream.remoteCloseWrite()
           break
         case MessageTypes.RESET_INITIATOR:
         case MessageTypes.RESET_RECEIVER:
           // The remote has errored, stop reading and writing to the stream immediately
-          stream.onRemoteReset()
+          stream.reset()
           break
         default:
-          this.log('unknown message type')
+          this.log('unknown message type %s', type)
       }
     } catch (err: any) {
-      this.log.error('error while processing message - %e', err)
+      this.log.error('error while processing message', err)
       stream.abort(err)
     }
   }

package/src/stream.ts

@@ -1,127 +1,95 @@
-import { AbstractStream } from '@libp2p/utils'
+import { AbstractStream } from '@libp2p/utils/abstract-stream'
 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 { MplexStreamMuxer } from './mplex.ts'
-import type { Logger, MessageStreamDirection } from '@libp2p/interface'
-import type { AbstractStreamInit, SendResult } from '@libp2p/utils'
-import type { AbortOptions } from 'it-pushable'
+import type { Message } from './message-types.js'
+import type { Logger } from '@libp2p/interface'
+import type { AbstractStreamInit } from '@libp2p/utils/abstract-stream'
 
 export interface Options {
   id: number
+  send(msg: Message): Promise<void>
   log: Logger
-  direction: MessageStreamDirection
+  name?: string
+  onEnd?(err?: Error): void
+  type?: 'initiator' | 'receiver'
   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 {
-  public readonly streamId: number
+  private readonly name: string
+  private readonly streamId: number
+  private readonly send: (msg: Message) => Promise<void>
   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))
+  }
 
-    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 sendNewStream (): Promise<void> {
+    await this.send({ id: this.streamId, type: InitiatorMessageTypes.NEW_STREAM, data: new Uint8ArrayList(uint8ArrayFromString(this.name)) })
   }
 
-  sendData (data: Uint8ArrayList): SendResult {
-    const list = new Uint8ArrayList()
-    const sentBytes = data.byteLength
+  async sendData (data: Uint8ArrayList): Promise<void> {
+    data = data.sublist()
 
     while (data.byteLength > 0) {
       const toSend = Math.min(data.byteLength, this.maxDataSize)
-      const slice = data.sublist(0, toSend)
-      data = data.sublist(toSend)
-
-      list.append(
-        encode({
-          id: this.streamId,
-          type: this.types.MESSAGE,
-          data: slice
-        })
-      )
-    }
-
-    return {
-      sentBytes,
-      canSendMore: this.muxer.send(list)
-    }
-  }
-
-  sendReset (): boolean {
-    return this.muxer.send(
-      encode({
+      await this.send({
         id: this.streamId,
-        type: this.types.RESET
+        type: this.types.MESSAGE,
+        data: data.sublist(0, toSend)
       })
-    )
-  }
 
-  async sendCloseWrite (options?: AbortOptions): Promise<void> {
-    this.muxer.send(
-      encode({
-        id: this.streamId,
-        type: this.types.CLOSE
-      })
-    )
-    options?.signal?.throwIfAborted()
+      data.consume(toSend)
+    }
   }
 
-  async sendCloseRead (options?: AbortOptions): Promise<void> {
-    options?.signal?.throwIfAborted()
-    // mplex does not support close read, only close write
+  async sendReset (): Promise<void> {
+    await this.send({ id: this.streamId, type: this.types.RESET })
   }
 
-  sendPause (): void {
-    // mplex does not support backpressure
+  async sendCloseWrite (): Promise<void> {
+    await this.send({ id: this.streamId, type: this.types.CLOSE })
   }
 
-  sendResume (): void {
-    // mplex does not support backpressure
+  async sendCloseRead (): Promise<void> {
+    // mplex does not support close read, only close write
   }
 }
 
 export function createStream (options: Options): MplexStream {
-  const { id, muxer, direction, maxMsgSize = MAX_MSG_SIZE } = options
+  const { id, name, send, onEnd, type = 'initiator', maxMsgSize = MAX_MSG_SIZE } = options
+  const direction = type === 'initiator' ? 'outbound' : 'inbound'
 
   return new MplexStream({
-    ...options,
-    id: direction === 'outbound' ? (`i${id}`) : `r${id}`,
+    id: type === 'initiator' ? (`i${id}`) : `r${id}`,
+    streamId: id,
+    name: `${name ?? id}`,
     direction,
     maxDataSize: maxMsgSize,
-    muxer,
-    log: options.log.newScope(`${direction}:${id}`),
-    protocol: ''
+    onEnd,
+    send,
+    log: options.log.newScope(`${direction}:${id}`)
   })
 }