feat: add stream middleware

.github/dictionary.txt

@@ -14,6 +14,7 @@ additionals
 SECG
 Certicom
 RSAES
+unuse
 dialback
 chacha
 peerStore

packages/interface-internal/src/registrar.ts

@@ -1,4 +1,4 @@
-import type { StreamHandler, StreamHandlerOptions, StreamHandlerRecord, Topology, AbortOptions } from '@libp2p/interface'
+import type { StreamHandler, StreamHandlerOptions, StreamHandlerRecord, Topology, StreamMiddleware, AbortOptions } from '@libp2p/interface'
 
 export type {
   /**
@@ -63,6 +63,30 @@ export interface Registrar {
    */
   getHandler(protocol: string): StreamHandlerRecord
 
+  /**
+   * Retrieve any registered middleware for a given protocol.
+   *
+   * @param protocol - The protocol to fetch middleware for
+   * @returns A list of `StreamMiddleware` implementations
+   */
+  use(protocol: string, middleware: StreamMiddleware[]): void
+
+  /**
+   * Retrieve any registered middleware for a given protocol.
+   *
+   * @param protocol - The protocol to fetch middleware for
+   * @returns A list of `StreamMiddleware` implementations
+   */
+  unuse(protocol: string): void
+
+  /**
+   * Retrieve any registered middleware for a given protocol.
+   *
+   * @param protocol - The protocol to fetch middleware for
+   * @returns A list of `StreamMiddleware` implementations
+   */
+  getMiddleware(protocol: string): StreamMiddleware[]
+
   /**
    * Register a topology handler for a protocol - the topology will be
    * invoked when peers are discovered on the network that support the

packages/interface/src/index.ts

@@ -23,7 +23,7 @@ import type { PeerInfo } from './peer-info.js'
 import type { PeerRouting } from './peer-routing.js'
 import type { Address, Peer, PeerStore } from './peer-store.js'
 import type { Startable } from './startable.js'
-import type { StreamHandler, StreamHandlerOptions } from './stream-handler.js'
+import type { StreamHandler, StreamHandlerOptions, StreamMiddleware } from './stream-handler.js'
 import type { Stream } from './stream.js'
 import type { Topology } from './topology.js'
 import type { Listener, OutboundConnectionUpgradeEvents } from './transport.js'
@@ -781,6 +781,33 @@ export interface Libp2p<T extends ServiceMap = ServiceMap> extends Startable, Ty
    */
   unregister(id: string): void
 
+  /**
+   * Registers one or more middleware implementations that will be invoked for
+   * incoming and outgoing protocol streams that match the passed protocol.
+   *
+   * @example
+   *
+   * ```TypeScript
+   * libp2p.use('/my/protocol/1.0.0', (stream, connection, next) => {
+   *   // do something with stream and/or connection
+   *   next(stream, connection)
+   * })
+   * ```
+   */
+  use (protocol: string, middleware: StreamMiddleware | StreamMiddleware[]): void
+
+  /**
+   * Deregisters all middleware for the passed protocol.
+   *
+   * @example
+   *
+   * ```TypeScript
+   * libp2p.unuse('/my/protocol/1.0.0')
+   * // any previously registered middleware will no longer be invoked
+   * ```
+   */
+  unuse (protocol: string): void
+
   /**
    * Returns the public key for the passed PeerId. If the PeerId is of the 'RSA'
    * type this may mean searching the routing if the peer's key is not present

packages/interface/src/stream-handler.ts

@@ -7,6 +7,13 @@ export interface StreamHandler {
   (stream: Stream, connection: Connection): void | Promise<void>
 }
 
+/**
+ * Stream middleware allows accessing stream data outside of the stream handler
+ */
+export interface StreamMiddleware {
+  (stream: Stream, connection: Connection, next: (stream: Stream, connection: Connection) => void): void | Promise<void>
+}
+
 export interface StreamHandlerOptions extends AbortOptions {
   /**
    * How many incoming streams can be open for this protocol at the same time on each connection
@@ -33,6 +40,11 @@ export interface StreamHandlerOptions extends AbortOptions {
    * protocol(s), the existing handler will be discarded.
    */
   force?: true
+
+  /**
+   * Middleware allows accessing stream data outside of the stream handler
+   */
+  middleware?: StreamMiddleware[]
 }
 
 export interface StreamHandlerRecord {

packages/libp2p/src/connection.ts

@@ -126,7 +126,7 @@ export class Connection extends TypedEventEmitter<MessageStreamEvents> implement
     }
 
     this.log.trace('starting new stream for protocols %s', protocols)
-    const muxedStream = await this.muxer.createStream({
+    let muxedStream = await this.muxer.createStream({
       ...options,
 
       // most underlying transports only support negotiating a single protocol
@@ -177,6 +177,24 @@ export class Connection extends TypedEventEmitter<MessageStreamEvents> implement
 
       this.components.metrics?.trackProtocolStream(muxedStream)
 
+      const middleware = this.components.registrar.getMiddleware(muxedStream.protocol)
+
+      middleware.push((stream, connection, next) => {
+        next(stream, connection)
+      })
+
+      let i = 0
+      let connection: ConnectionInterface = this
+
+      while (i < middleware.length) {
+        // eslint-disable-next-line no-loop-func
+        middleware[i](muxedStream, connection, (s, c) => {
+          muxedStream = s
+          connection = c
+          i++
+        })
+      }
+
       return muxedStream
     } catch (err: any) {
       if (muxedStream.status === 'open') {
@@ -190,7 +208,7 @@ export class Connection extends TypedEventEmitter<MessageStreamEvents> implement
   }
 
   private async onIncomingStream (evt: CustomEvent<Stream>): Promise<void> {
-    const muxedStream = evt.detail
+    let muxedStream = evt.detail
 
     const signal = AbortSignal.timeout(this.inboundStreamProtocolNegotiationTimeout)
     setMaxListeners(Infinity, signal)
@@ -235,7 +253,22 @@ export class Connection extends TypedEventEmitter<MessageStreamEvents> implement
         throw new LimitedConnectionError('Cannot open protocol stream on limited connection')
       }
 
-      await handler(muxedStream, this)
+      const middleware = this.components.registrar.getMiddleware(muxedStream.protocol)
+
+      middleware.push(async (stream, connection, next) => {
+        await handler(stream, connection)
+        next(stream, connection)
+      })
+
+      let connection: ConnectionInterface = this
+
+      for (const m of middleware) {
+        // eslint-disable-next-line no-loop-func
+        await m(muxedStream, connection, (s, c) => {
+          muxedStream = s
+          connection = c
+        })
+      }
     } catch (err: any) {
       muxedStream.abort(err)
     }

packages/libp2p/src/libp2p.ts

@@ -24,7 +24,7 @@ import { userAgent } from './user-agent.js'
 import * as pkg from './version.js'
 import type { Components } from './components.js'
 import type { Libp2p as Libp2pInterface, Libp2pInit } from './index.js'
-import type { PeerRouting, ContentRouting, Libp2pEvents, PendingDial, ServiceMap, AbortOptions, ComponentLogger, Logger, Connection, NewStreamOptions, Stream, Metrics, PeerId, PeerInfo, PeerStore, Topology, Libp2pStatus, IsDialableOptions, DialOptions, PublicKey, Ed25519PeerId, Secp256k1PeerId, RSAPublicKey, RSAPeerId, URLPeerId, Ed25519PublicKey, Secp256k1PublicKey, StreamHandler, StreamHandlerOptions } from '@libp2p/interface'
+import type { PeerRouting, ContentRouting, Libp2pEvents, PendingDial, ServiceMap, AbortOptions, ComponentLogger, Logger, Connection, NewStreamOptions, Stream, Metrics, PeerId, PeerInfo, PeerStore, Topology, Libp2pStatus, IsDialableOptions, DialOptions, PublicKey, Ed25519PeerId, Secp256k1PeerId, RSAPublicKey, RSAPeerId, URLPeerId, Ed25519PublicKey, Secp256k1PublicKey, StreamHandler, StreamHandlerOptions, StreamMiddleware } from '@libp2p/interface'
 import type { Multiaddr } from '@multiformats/multiaddr'
 
 export class Libp2p<T extends ServiceMap = ServiceMap> extends TypedEventEmitter<Libp2pEvents> implements Libp2pInterface<T> {
@@ -401,6 +401,14 @@ export class Libp2p<T extends ServiceMap = ServiceMap> extends TypedEventEmitter
     this.components.registrar.unregister(id)
   }
 
+  use (protocol: string, middleware: StreamMiddleware | StreamMiddleware[]): void {
+    this.components.registrar.use(protocol, Array.isArray(middleware) ? middleware : [middleware])
+  }
+
+  unuse (protocol: string): void {
+    this.components.registrar.unuse(protocol)
+  }
+
   async isDialable (multiaddr: Multiaddr, options: IsDialableOptions = {}): Promise<boolean> {
     return this.components.connectionManager.isDialable(multiaddr, options)
   }

packages/libp2p/src/registrar.ts

@@ -1,7 +1,7 @@
 import { InvalidParametersError } from '@libp2p/interface'
 import { mergeOptions, trackedMap } from '@libp2p/utils'
 import { DuplicateProtocolHandlerError, UnhandledProtocolError } from './errors.js'
-import type { IdentifyResult, Libp2pEvents, Logger, PeerUpdate, PeerId, PeerStore, Topology, StreamHandler, StreamHandlerRecord, StreamHandlerOptions, AbortOptions, Metrics } from '@libp2p/interface'
+import type { IdentifyResult, Libp2pEvents, Logger, PeerUpdate, PeerId, PeerStore, Topology, StreamHandler, StreamHandlerRecord, StreamHandlerOptions, AbortOptions, Metrics, StreamMiddleware } from '@libp2p/interface'
 import type { Registrar as RegistrarInterface } from '@libp2p/interface-internal'
 import type { ComponentLogger } from '@libp2p/logger'
 import type { TypedEventTarget } from 'main-event'
@@ -25,10 +25,12 @@ export class Registrar implements RegistrarInterface {
   private readonly topologies: Map<string, Map<string, Topology>>
   private readonly handlers: Map<string, StreamHandlerRecord>
   private readonly components: RegistrarComponents
+  private readonly middleware: Map<string, StreamMiddleware[]>
 
   constructor (components: RegistrarComponents) {
     this.components = components
     this.log = components.logger.forComponent('libp2p:registrar')
+    this.middleware = new Map()
     this.topologies = new Map()
     components.metrics?.registerMetricGroup('libp2p_registrar_topologies', {
       calculate: () => {
@@ -164,6 +166,18 @@ export class Registrar implements RegistrarInterface {
     }
   }
 
+  use (protocol: string, middleware: StreamMiddleware[]): void {
+    this.middleware.set(protocol, middleware)
+  }
+
+  unuse (protocol: string): void {
+    this.middleware.delete(protocol)
+  }
+
+  getMiddleware (protocol: string): StreamMiddleware[] {
+    return this.middleware.get(protocol) ?? []
+  }
+
   /**
    * Remove a disconnected peer from the record
    */

packages/libp2p/test/connection/index.spec.ts

@@ -11,7 +11,7 @@ import { fromString as uint8ArrayFromString } from 'uint8arrays/from-string'
 import { createConnection } from '../../src/connection.js'
 import { UnhandledProtocolError } from '../../src/errors.ts'
 import type { ConnectionComponents, ConnectionInit } from '../../src/connection.js'
-import type { MultiaddrConnection, PeerStore, StreamMuxer } from '@libp2p/interface'
+import type { MultiaddrConnection, PeerStore, Stream, StreamMuxer } from '@libp2p/interface'
 import type { Registrar } from '@libp2p/interface-internal'
 import type { StubbedInstance } from 'sinon-ts'
 
@@ -38,6 +38,7 @@ describe('connection', () => {
       },
       options: {}
     })
+    registrar.getMiddleware.withArgs(ECHO_PROTOCOL).returns([])
 
     components = {
       peerStore,
@@ -223,6 +224,7 @@ describe('connection', () => {
       }
     })
     registrar.getProtocols.returns([protocol])
+    registrar.getMiddleware.withArgs(protocol).returns([])
 
     const connection = createConnection(components, init)
     expect(connection.streams).to.have.lengthOf(0)
@@ -259,6 +261,7 @@ describe('connection', () => {
       }
     })
     registrar.getProtocols.returns([protocol])
+    registrar.getMiddleware.withArgs(protocol).returns([])
 
     const connection = createConnection(components, init)
     expect(connection.streams).to.have.lengthOf(0)
@@ -274,6 +277,7 @@ describe('connection', () => {
     const protocol = '/test/protocol'
 
     registrar.getHandler.withArgs(protocol).throws(new UnhandledProtocolError())
+    registrar.getMiddleware.withArgs(protocol).returns([])
 
     const connection = createConnection(components, init)
     expect(connection.streams).to.have.lengthOf(0)
@@ -289,4 +293,102 @@ describe('connection', () => {
     await expect(connection.newStream(protocol, opts)).to.eventually.be.rejected
       .with.property('name', 'TooManyOutboundProtocolStreamsError')
   })
+
+  it('should support outgoing stream middleware', async () => {
+    const streamProtocol = '/test/protocol'
+
+    const middleware1 = Sinon.stub().callsFake((stream, connection, next) => {
+      next(stream, connection)
+    })
+    const middleware2 = Sinon.stub().callsFake((stream, connection, next) => {
+      next(stream, connection)
+    })
+
+    const middleware = [
+      middleware1,
+      middleware2
+    ]
+
+    registrar.getMiddleware.withArgs(streamProtocol).returns(middleware)
+    registrar.getHandler.withArgs(streamProtocol).returns({
+      handler: () => {},
+      options: {}
+    })
+
+    const connection = createConnection(components, init)
+
+    await connection.newStream(streamProtocol)
+
+    expect(middleware1.called).to.be.true()
+    expect(middleware2.called).to.be.true()
+  })
+
+  it('should support incoming stream middleware', async () => {
+    const streamProtocol = '/test/protocol'
+
+    const middleware1 = Sinon.stub().callsFake((stream, connection, next) => {
+      next(stream, connection)
+    })
+    const middleware2 = Sinon.stub().callsFake((stream, connection, next) => {
+      next(stream, connection)
+    })
+
+    const middleware = [
+      middleware1,
+      middleware2
+    ]
+
+    registrar.getMiddleware.withArgs(streamProtocol).returns(middleware)
+    registrar.getHandler.withArgs(streamProtocol).returns({
+      handler: () => {},
+      options: {}
+    })
+
+    const muxer = stubInterface<StreamMuxer>({
+      streams: []
+    })
+
+    createConnection(components, {
+      ...init,
+      muxer
+    })
+
+    expect(muxer.addEventListener.getCall(0).args[0]).to.equal('stream')
+    const onIncomingStream = muxer.addEventListener.getCall(0).args[1]
+
+    if (onIncomingStream == null) {
+      throw new Error('No incoming stream handler registered')
+    }
+
+    const incomingStream = stubInterface<Stream>({
+      protocol: streamProtocol
+    })
+
+    if (typeof onIncomingStream !== 'function') {
+      throw new Error('Stream handler was not function')
+    }
+
+    onIncomingStream(new CustomEvent('stream', {
+      detail: incomingStream
+    }))
+    /*
+    const incomingStream = stubInterface<Stream>({
+      id: 'stream-id',
+      log: logger('test-stream'),
+      direction: 'outbound',
+      sink: async (source) => drain(source),
+      source: map((async function * () {
+        yield '/multistream/1.0.0\n'
+        yield `${streamProtocol}\n`
+      })(), str => encode.single(uint8ArrayFromString(str)))
+    })
+*/
+    //    onIncomingStream?.(incomingStream)
+
+    // incoming stream is opened asynchronously
+    await delay(100)
+
+    expect(middleware1.called).to.be.true()
+    expect(middleware2.called).to.be.true()
+  })
 })