feat: add stream middleware (#3173)

Adds middleware handlers for protocol streams.  They are invoked
for incoming and outgoing streams and allow access to the stream
and connection before the handler (incoming) or caller (outgoing)
receive them.

This way middleware can wrap streams in transforms, or deny access,
or something else.

```ts
libp2p.use('/my/protocol/1.0.0', (stream, connection, next) => {
  const originalSource = stream.source

  // increment all byte values in the stream by one
  stream.source = (async function * () {
    for await (const buf of originalSource) {
      buf = buf.map(val => val + 1)

      yield buf
    }
  })()

  // pass the stream on to the next middleware
  next(stream, connection)
})
```

Author: achingbrain

.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()
+  })
 })