@libp2p/mplex 1.0.0

package/src/encode.ts

@@ -0,0 +1,79 @@
+import varint from 'varint'
+import { Message, MessageTypes } from './message-types.js'
+import type { Source } from 'it-stream-types'
+
+const POOL_SIZE = 10 * 1024
+
+function allocUnsafe (size: number) {
+  if (globalThis.Buffer != null) {
+    return Buffer.allocUnsafe(size)
+  }
+
+  return new Uint8Array(size)
+}
+
+class Encoder {
+  private _pool: Uint8Array
+  private _poolOffset: number
+
+  constructor () {
+    this._pool = allocUnsafe(POOL_SIZE)
+    this._poolOffset = 0
+  }
+
+  /**
+   * Encodes the given message and returns it and its header
+   */
+  write (msg: Message): Uint8Array[] {
+    const pool = this._pool
+    let offset = this._poolOffset
+
+    varint.encode(msg.id << 3 | msg.type, pool, offset)
+    offset += varint.encode.bytes
+
+    if ((msg.type === MessageTypes.NEW_STREAM || msg.type === MessageTypes.MESSAGE_INITIATOR || msg.type === MessageTypes.MESSAGE_RECEIVER) && msg.data != null) {
+      varint.encode(msg.data.length, pool, offset)
+    } else {
+      varint.encode(0, pool, offset)
+    }
+
+    offset += varint.encode.bytes
+
+    const header = pool.slice(this._poolOffset, offset)
+
+    if (POOL_SIZE - offset < 100) {
+      this._pool = allocUnsafe(POOL_SIZE)
+      this._poolOffset = 0
+    } else {
+      this._poolOffset = offset
+    }
+
+    if ((msg.type === MessageTypes.NEW_STREAM || msg.type === MessageTypes.MESSAGE_INITIATOR || msg.type === MessageTypes.MESSAGE_RECEIVER) && msg.data != null) {
+      return [
+        header,
+        msg.data instanceof Uint8Array ? msg.data : msg.data.slice()
+      ]
+    }
+
+    return [
+      header
+    ]
+  }
+}
+
+const encoder = new Encoder()
+
+/**
+ * Encode and yield one or more messages
+ */
+export async function * encode (source: Source<Message | Message[]>) {
+  for await (const msg of source) {
+    if (Array.isArray(msg)) {
+      for (const m of msg) {
+        yield * encoder.write(m)
+      }
+    } else {
+      yield * encoder.write(msg)
+    }
+  }
+}

package/src/index.ts

@@ -0,0 +1,281 @@
+import { pipe } from 'it-pipe'
+import { Pushable, pushableV } from 'it-pushable'
+import { abortableSource } from 'abortable-iterator'
+import { encode } from './encode.js'
+import { decode } from './decode.js'
+import { restrictSize } from './restrict-size.js'
+import { MessageTypes, MessageTypeNames, Message } from './message-types.js'
+import { createStream } from './stream.js'
+import { toString as uint8ArrayToString } from 'uint8arrays'
+import { trackedMap } from '@libp2p/tracked-map'
+import { logger } from '@libp2p/logger'
+import type { AbortOptions } from '@libp2p/interfaces'
+import type { Sink } from 'it-stream-types'
+import type { Muxer } from '@libp2p/interfaces/stream-muxer'
+import type { Stream } from '@libp2p/interfaces/connection'
+import type { ComponentMetricsTracker } from '@libp2p/interfaces/metrics'
+import each from 'it-foreach'
+
+const log = logger('libp2p:mplex')
+
+function printMessage (msg: Message) {
+  const output: any = {
+    ...msg,
+    type: `${MessageTypeNames[msg.type]} (${msg.type})`
+  }
+
+  if (msg.type === MessageTypes.NEW_STREAM) {
+    output.data = uint8ArrayToString(msg.data instanceof Uint8Array ? msg.data : msg.data.slice())
+  }
+
+  if (msg.type === MessageTypes.MESSAGE_INITIATOR || msg.type === MessageTypes.MESSAGE_RECEIVER) {
+    output.data = uint8ArrayToString(msg.data instanceof Uint8Array ? msg.data : msg.data.slice(), 'base16')
+  }
+
+  return output
+}
+
+export interface MplexStream extends Stream {
+  source: Pushable<Uint8Array>
+}
+
+export interface MplexOptions extends AbortOptions {
+  onStream?: (...args: any[]) => void
+  onStreamEnd?: (...args: any[]) => void
+  maxMsgSize?: number
+  metrics?: ComponentMetricsTracker
+}
+
+export class Mplex implements Muxer {
+  static multicodec = '/mplex/6.7.0'
+
+  public sink: Sink<Uint8Array>
+  public source: AsyncIterable<Uint8Array>
+
+  private _streamId: number
+  private readonly _streams: { initiators: Map<number, MplexStream>, receivers: Map<number, MplexStream> }
+  private readonly _options: MplexOptions
+  private readonly _source: { push: (val: Message) => void, end: (err?: Error) => void }
+
+  constructor (options?: MplexOptions) {
+    options = options ?? {}
+
+    this._streamId = 0
+    this._streams = {
+      /**
+       * Stream to ids map
+       */
+      initiators: trackedMap<number, MplexStream>({ metrics: options.metrics, component: 'mplex', metric: 'initiatorStreams' }),
+      /**
+       * Stream to ids map
+       */
+      receivers: trackedMap<number, MplexStream>({ metrics: options.metrics, component: 'mplex', metric: 'receiverStreams' })
+    }
+    this._options = options
+
+    /**
+     * An iterable sink
+     */
+    this.sink = this._createSink()
+
+    /**
+     * An iterable source
+     */
+    const source = this._createSource()
+    this._source = source
+    this.source = source
+  }
+
+  /**
+   * Returns a Map of streams and their ids
+   */
+  get streams () {
+    // Inbound and Outbound streams may have the same ids, so we need to make those unique
+    const streams: Stream[] = []
+    this._streams.initiators.forEach(stream => {
+      streams.push(stream)
+    })
+    this._streams.receivers.forEach(stream => {
+      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.
+   */
+  newStream (name?: string): Stream {
+    const id = this._streamId++
+    name = name == null ? id.toString() : name.toString()
+    const registry = this._streams.initiators
+    return this._newStream({ id, name, type: 'initiator', registry })
+  }
+
+  /**
+   * Called whenever an inbound stream is created
+   */
+  _newReceiverStream (options: { id: number, name: string }) {
+    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> }) {
+    const { id, name, type, registry } = options
+
+    log('new %s stream %s %s', type, id, name)
+
+    if (registry.has(id)) {
+      throw new Error(`${type} stream ${id} already exists!`)
+    }
+
+    const send = (msg: Message) => {
+      if (log.enabled) {
+        log('%s stream %s send', type, id, printMessage(msg))
+      }
+
+      if (msg.type === MessageTypes.NEW_STREAM || msg.type === MessageTypes.MESSAGE_INITIATOR || msg.type === MessageTypes.MESSAGE_RECEIVER) {
+        msg.data = msg.data instanceof Uint8Array ? msg.data : msg.data.slice()
+      }
+
+      this._source.push(msg)
+    }
+
+    const onEnd = () => {
+      log('%s stream %s %s ended', type, id, name)
+      registry.delete(id)
+
+      if (this._options.onStreamEnd != null) {
+        this._options.onStreamEnd(stream)
+      }
+    }
+
+    const stream = createStream({ id, name, send, type, onEnd, maxMsgSize: this._options.maxMsgSize })
+    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 () {
+    const sink: Sink<Uint8Array> = async source => {
+      if (this._options.signal != null) {
+        source = abortableSource(source, this._options.signal)
+      }
+
+      try {
+        await pipe(
+          source,
+          source => each(source, (buf) => {
+            // console.info('incoming', uint8ArrayToString(buf, 'base64'))
+          }),
+          decode,
+          restrictSize(this._options.maxMsgSize),
+          async source => {
+            for await (const msg of source) {
+              this._handleIncoming(msg)
+            }
+          }
+        )
+
+        this._source.end()
+      } catch (err: any) {
+        log('error in sink', err)
+        this._source.end(err) // End the source with an error
+      }
+    }
+
+    return sink
+  }
+
+  /**
+   * Creates a source that restricts outgoing message sizes
+   * and varint encodes them
+   */
+  _createSource () {
+    const onEnd = (err?: Error) => {
+      const { initiators, receivers } = this._streams
+      // Abort all the things!
+      for (const s of initiators.values()) {
+        s.abort(err)
+      }
+      for (const s of receivers.values()) {
+        s.abort(err)
+      }
+    }
+    const source = pushableV<Message>({ onEnd })
+    /*
+    const p = pipe(
+      source,
+      source => each(source, (msgs) => {
+        if (log.enabled) {
+          msgs.forEach(msg => {
+            log('outgoing message', printMessage(msg))
+          })
+        }
+      }),
+      source => encode(source),
+      source => each(source, (buf) => {
+        console.info('outgoing', uint8ArrayToString(buf, 'base64'))
+      })
+    )
+
+    return Object.assign(p, {
+      push: source.push,
+      end: source.end,
+      return: source.return
+    })
+*/
+    return Object.assign(encode(source), {
+      push: source.push,
+      end: source.end,
+      return: source.return
+    })
+  }
+
+  _handleIncoming (message: Message) {
+    const { id, type } = message
+
+    if (log.enabled) {
+      log('incoming message', printMessage(message))
+    }
+
+    // Create a new stream?
+    if (message.type === MessageTypes.NEW_STREAM) {
+      const stream = this._newReceiverStream({ id, name: uint8ArrayToString(message.data instanceof Uint8Array ? message.data : message.data.slice()) })
+
+      if (this._options.onStream != null) {
+        this._options.onStream(stream)
+      }
+
+      return
+    }
+
+    const list = (type & 1) === 1 ? this._streams.initiators : this._streams.receivers
+    const stream = list.get(id)
+
+    if (stream == null) {
+      return log('missing stream %s', id)
+    }
+
+    switch (type) {
+      case MessageTypes.MESSAGE_INITIATOR:
+      case MessageTypes.MESSAGE_RECEIVER:
+        stream.source.push(message.data.slice())
+        break
+      case MessageTypes.CLOSE_INITIATOR:
+      case MessageTypes.CLOSE_RECEIVER:
+        stream.close()
+        break
+      case MessageTypes.RESET_INITIATOR:
+      case MessageTypes.RESET_RECEIVER:
+        stream.reset()
+        break
+      default:
+        log('unknown message type %s', type)
+    }
+  }
+}

package/src/message-types.ts

@@ -0,0 +1,79 @@
+import type { Uint8ArrayList } from 'uint8arraylist'
+
+type INITIATOR_NAME = 'NEW_STREAM' | 'MESSAGE' | 'CLOSE' | 'RESET'
+type RECEIVER_NAME = 'MESSAGE' | 'CLOSE' | 'RESET'
+type NAME = 'NEW_STREAM' | 'MESSAGE_INITIATOR' | 'CLOSE_INITIATOR' | 'RESET_INITIATOR' | 'MESSAGE_RECEIVER' | 'CLOSE_RECEIVER' | 'RESET_RECEIVER'
+type CODE = 0 | 1 | 2 | 3 | 4 | 5 | 6
+
+export enum MessageTypes {
+  NEW_STREAM = 0,
+  MESSAGE_RECEIVER = 1,
+  MESSAGE_INITIATOR = 2,
+  CLOSE_RECEIVER = 3,
+  CLOSE_INITIATOR = 4,
+  RESET_RECEIVER = 5,
+  RESET_INITIATOR = 6
+}
+
+export const MessageTypeNames: Record<CODE, NAME> = Object.freeze({
+  0: 'NEW_STREAM',
+  1: 'MESSAGE_RECEIVER',
+  2: 'MESSAGE_INITIATOR',
+  3: 'CLOSE_RECEIVER',
+  4: 'CLOSE_INITIATOR',
+  5: 'RESET_RECEIVER',
+  6: 'RESET_INITIATOR'
+})
+
+export const InitiatorMessageTypes: Record<INITIATOR_NAME, CODE> = Object.freeze({
+  NEW_STREAM: MessageTypes.NEW_STREAM,
+  MESSAGE: MessageTypes.MESSAGE_INITIATOR,
+  CLOSE: MessageTypes.CLOSE_INITIATOR,
+  RESET: MessageTypes.RESET_INITIATOR
+})
+
+export const ReceiverMessageTypes: Record<RECEIVER_NAME, CODE> = Object.freeze({
+  MESSAGE: MessageTypes.MESSAGE_RECEIVER,
+  CLOSE: MessageTypes.CLOSE_RECEIVER,
+  RESET: MessageTypes.RESET_RECEIVER
+})
+
+export interface NewStreamMessage {
+  id: number
+  type: MessageTypes.NEW_STREAM
+  data: Uint8Array | Uint8ArrayList
+}
+
+export interface MessageReceiverMessage {
+  id: number
+  type: MessageTypes.MESSAGE_RECEIVER
+  data: Uint8Array | Uint8ArrayList
+}
+
+export interface MessageInitiatorMessage {
+  id: number
+  type: MessageTypes.MESSAGE_INITIATOR
+  data: Uint8Array | Uint8ArrayList
+}
+
+export interface CloseReceiverMessage {
+  id: number
+  type: MessageTypes.CLOSE_RECEIVER
+}
+
+export interface CloseInitiatorMessage {
+  id: number
+  type: MessageTypes.CLOSE_INITIATOR
+}
+
+export interface ResetReceiverMessage {
+  id: number
+  type: MessageTypes.RESET_RECEIVER
+}
+
+export interface ResetInitiatorMessage {
+  id: number
+  type: MessageTypes.RESET_INITIATOR
+}
+
+export type Message = NewStreamMessage | MessageReceiverMessage | MessageInitiatorMessage | CloseReceiverMessage | CloseInitiatorMessage | ResetReceiverMessage | ResetInitiatorMessage

package/src/restrict-size.ts

@@ -0,0 +1,36 @@
+import { Message, MessageTypes } from './message-types.js'
+import type { Source, Transform } from 'it-stream-types'
+
+export const MAX_MSG_SIZE = 1 << 20 // 1MB
+
+/**
+ * Creates an iterable transform that restricts message sizes to
+ * the given maximum size.
+ */
+export function restrictSize (max?: number): Transform<Message | Message[], Message> {
+  const maxSize = max ?? MAX_MSG_SIZE
+
+  const checkSize = (msg: Message) => {
+    if (msg.type !== MessageTypes.NEW_STREAM && msg.type !== MessageTypes.MESSAGE_INITIATOR && msg.type !== MessageTypes.MESSAGE_RECEIVER) {
+      return
+    }
+
+    if (msg.data.byteLength > maxSize) {
+      throw Object.assign(new Error('message size too large!'), { code: 'ERR_MSG_TOO_BIG' })
+    }
+  }
+
+  return (source: Source<Message | Message[]>) => {
+    return (async function * restrictSize () {
+      for await (const msg of source) {
+        if (Array.isArray(msg)) {
+          msg.forEach(checkSize)
+          yield * msg
+        } else {
+          checkSize(msg)
+          yield msg
+        }
+      }
+    })()
+  }
+}

package/src/stream.ts

@@ -0,0 +1,182 @@
+import { abortableSource } from 'abortable-iterator'
+import { pushable } from 'it-pushable'
+import errCode from 'err-code'
+import { MAX_MSG_SIZE } from './restrict-size.js'
+import { anySignal } from 'any-signal'
+import { InitiatorMessageTypes, ReceiverMessageTypes } from './message-types.js'
+import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
+import { Uint8ArrayList } from 'uint8arraylist'
+import { logger } from '@libp2p/logger'
+import type { Message } from './message-types.js'
+import type { Timeline } from '@libp2p/interfaces/connection'
+import type { Source } from 'it-stream-types'
+import type { MplexStream } from './index.js'
+
+const log = logger('libp2p:mplex:stream')
+
+const ERR_MPLEX_STREAM_RESET = 'ERR_MPLEX_STREAM_RESET'
+const ERR_MPLEX_STREAM_ABORT = 'ERR_MPLEX_STREAM_ABORT'
+
+export interface Options {
+  id: number
+  send: (msg: Message) => void
+  name?: string
+  onEnd?: (err?: Error) => void
+  type?: 'initiator' | 'receiver'
+  maxMsgSize?: number
+}
+
+export function createStream (options: Options): MplexStream {
+  const { id, name, send, onEnd, type = 'initiator', maxMsgSize = MAX_MSG_SIZE } = options
+
+  const abortController = new AbortController()
+  const resetController = new AbortController()
+  const Types = type === 'initiator' ? InitiatorMessageTypes : ReceiverMessageTypes
+  const externalId = type === 'initiator' ? (`i${id}`) : `r${id}`
+  const streamName = `${name == null ? id : name}`
+
+  let sourceEnded = false
+  let sinkEnded = false
+  let endErr: Error | undefined
+
+  const timeline: Timeline = {
+    open: Date.now()
+  }
+
+  const onSourceEnd = (err?: Error) => {
+    if (sourceEnded) {
+      return
+    }
+
+    sourceEnded = true
+    log('%s stream %s source end', type, streamName, err)
+
+    if (err != null && endErr == null) {
+      endErr = err
+    }
+
+    if (sinkEnded) {
+      stream.timeline.close = Date.now()
+
+      if (onEnd != null) {
+        onEnd(endErr)
+      }
+    }
+  }
+
+  const onSinkEnd = (err?: Error) => {
+    if (sinkEnded) {
+      return
+    }
+
+    sinkEnded = true
+    log('%s stream %s sink end - err: %o', type, streamName, err)
+
+    if (err != null && endErr == null) {
+      endErr = err
+    }
+
+    if (sourceEnded) {
+      timeline.close = Date.now()
+
+      if (onEnd != null) {
+        onEnd(endErr)
+      }
+    }
+  }
+
+  const stream = {
+    // Close for reading
+    close: () => {
+      stream.source.end()
+    },
+    // Close for reading and writing (local error)
+    abort: (err?: Error) => {
+      log('%s stream %s abort', type, streamName, err)
+      // End the source with the passed error
+      stream.source.end(err)
+      abortController.abort()
+      onSinkEnd(err)
+    },
+    // Close immediately for reading and writing (remote error)
+    reset: () => {
+      const err = errCode(new Error('stream reset'), ERR_MPLEX_STREAM_RESET)
+      resetController.abort()
+      stream.source.end(err)
+      onSinkEnd(err)
+    },
+    sink: async (source: Source<Uint8Array>) => {
+      source = abortableSource(source, anySignal([
+        abortController.signal,
+        resetController.signal
+      ]))
+
+      try {
+        if (type === 'initiator') { // If initiator, open a new stream
+          send({ id, type: InitiatorMessageTypes.NEW_STREAM, data: uint8ArrayFromString(streamName) })
+        }
+
+        const uint8ArrayList = new Uint8ArrayList()
+
+        for await (const data of source) {
+          uint8ArrayList.append(data)
+
+          while (uint8ArrayList.length !== 0) {
+            if (uint8ArrayList.length <= maxMsgSize) {
+              send({ id, type: Types.MESSAGE, data: uint8ArrayList.subarray() })
+              uint8ArrayList.consume(uint8ArrayList.length)
+              break
+            }
+
+            const toSend = uint8ArrayList.length - maxMsgSize
+            send({ id, type: Types.MESSAGE, data: uint8ArrayList.subarray(0, toSend) })
+            uint8ArrayList.consume(toSend)
+          }
+        }
+      } catch (err: any) {
+        if (err.type === 'aborted' && err.message === 'The operation was aborted') {
+          if (resetController.signal.aborted) {
+            err.message = 'stream reset'
+            err.code = ERR_MPLEX_STREAM_RESET
+          }
+
+          if (abortController.signal.aborted) {
+            err.message = 'stream aborted'
+            err.code = ERR_MPLEX_STREAM_ABORT
+          }
+        }
+
+        // Send no more data if this stream was remotely reset
+        if (err.code === ERR_MPLEX_STREAM_RESET) {
+          log('%s stream %s reset', type, name)
+        } else {
+          log('%s stream %s error', type, name, err)
+          try {
+            send({ id, type: Types.RESET })
+          } catch (err) {
+            log('%s stream %s error sending reset', type, name, err)
+          }
+        }
+
+        stream.source.end(err)
+        onSinkEnd(err)
+        return
+      }
+
+      try {
+        send({ id, type: Types.CLOSE })
+      } catch (err) {
+        log('%s stream %s error sending close', type, name, err)
+      }
+
+      onSinkEnd()
+    },
+    source: pushable<Uint8Array>({
+      onEnd: onSourceEnd
+    }),
+    timeline,
+    id: externalId
+  }
+
+  return stream
+}