fix: add time out to incoming KAD streams (#3167)

Close the stream if the RPC invocation has not completed after a configurable window.

Author: achingbrain

packages/kad-dht/package.json

@@ -61,7 +61,6 @@
     "it-all": "^3.0.8",
     "it-drain": "^3.0.9",
     "it-length": "^3.0.8",
-    "it-length-prefixed": "^10.0.1",
     "it-map": "^3.1.3",
     "it-merge": "^3.0.11",
     "it-parallel": "^3.0.11",
@@ -95,6 +94,7 @@
     "execa": "^9.5.3",
     "it-filter": "^3.1.3",
     "it-last": "^3.0.8",
+    "it-length-prefixed": "^10.0.1",
     "it-pair": "^2.0.6",
     "it-stream-types": "^2.0.2",
     "lodash.random": "^3.2.0",

packages/kad-dht/src/index.ts

@@ -693,6 +693,14 @@ export interface KadDHTInit {
    * @default 10_000
    */
   onPeerConnectTimeout?: number
+
+  /**
+   * When acting as a DHT server, all incoming RPC requests must complete within
+   * this timeout in ms otherwise the stream will be closed.
+   *
+   * @default 10_000
+   */
+  incomingMessageTimeout?: number
 }
 
 export interface KadDHTComponents {

packages/kad-dht/src/rpc/index.ts

@@ -1,5 +1,5 @@
-import * as lp from 'it-length-prefixed'
-import { pipe } from 'it-pipe'
+import { TimeoutError } from '@libp2p/interface'
+import { pbStream } from 'it-protobuf-stream'
 import { Message, MessageType } from '../message/dht.js'
 import { AddProviderHandler } from './handlers/add-provider.js'
 import { FindNodeHandler } from './handlers/find-node.js'
@@ -30,6 +30,7 @@ export interface RPCInit {
   metricsPrefix: string
   datastorePrefix: string
   peerInfoMapper: PeerInfoMapper
+  incomingMessageTimeout?: number
 }
 
 export interface RPCComponents extends GetValueHandlerComponents, PutValueHandlerComponents, FindNodeHandlerComponents, GetProvidersHandlerComponents {
@@ -38,21 +39,22 @@ export interface RPCComponents extends GetValueHandlerComponents, PutValueHandle
 
 export class RPC {
   private readonly handlers: Record<string, DHTMessageHandler>
-  private readonly routingTable: RoutingTable
   private readonly log: Logger
   private readonly metrics: {
     operations?: CounterGroup
     errors?: CounterGroup
   }
 
+  private readonly incomingMessageTimeout: number
+
   constructor (components: RPCComponents, init: RPCInit) {
     this.metrics = {
       operations: components.metrics?.registerCounterGroup(`${init.metricsPrefix}_inbound_rpc_requests_total`),
       errors: components.metrics?.registerCounterGroup(`${init.metricsPrefix}_inbound_rpc_errors_total`)
     }
 
     this.log = components.logger.forComponent(`${init.logPrefix}:rpc`)
-    this.routingTable = init.routingTable
+    this.incomingMessageTimeout = init.incomingMessageTimeout ?? 10_000
     this.handlers = {
       [MessageType.GET_VALUE.toString()]: new GetValueHandler(components, init),
       [MessageType.PUT_VALUE.toString()]: new PutValueHandler(components, init),
@@ -92,34 +94,46 @@ export class RPC {
    * Handle incoming streams on the dht protocol
    */
   onIncomingStream (data: IncomingStreamData): void {
-    let message = 'unknown'
+    const message = 'unknown'
 
     Promise.resolve().then(async () => {
       const { stream, connection } = data
-      const peerId = connection.remotePeer
-
-      const self = this
-
-      await pipe(
-        stream,
-        (source) => lp.decode(source),
-        async function * (source) {
-          for await (const msg of source) {
-            // handle the message
-            const desMessage = Message.decode(msg)
-            message = desMessage.type
-            self.log('incoming %s from %p', desMessage.type, peerId)
-            const res = await self.handleMessage(peerId, desMessage)
-
-            // Not all handlers will return a response
-            if (res != null) {
-              yield Message.encode(res)
-            }
+
+      const abortListener = (): void => {
+        stream.abort(new TimeoutError())
+      }
+
+      let signal = AbortSignal.timeout(this.incomingMessageTimeout)
+      signal.addEventListener('abort', abortListener)
+
+      const messages = pbStream(stream).pb(Message)
+
+      try {
+        while (true) {
+          const message = await messages.read({
+            signal
+          })
+
+          // handle the message
+          this.log('incoming %s from %p', message.type, connection.remotePeer)
+          const res = await this.handleMessage(connection.remotePeer, message)
+
+          // Not all handlers will return a response
+          if (res != null) {
+            await messages.write(res, {
+              signal
+            })
           }
-        },
-        (source) => lp.encode(source),
-        stream
-      )
+
+          // we have received a message so reset the timeout controller to
+          // allow the remote to send another
+          signal.removeEventListener('abort', abortListener)
+          signal = AbortSignal.timeout(this.incomingMessageTimeout)
+          signal.addEventListener('abort', abortListener)
+        }
+      } catch (err: any) {
+        stream.abort(err)
+      }
     })
       .catch(err => {
         this.log.error('error handling %s RPC message from %p - %e', message, data.connection.remotePeer, err)

packages/kad-dht/test/libp2p-routing.spec.ts

@@ -70,7 +70,7 @@ function createPeer (peerId: PeerId, peer: Partial<Peer> = {}): Peer {
     id: peerId,
     addresses: [{
       isCertified: false,
-      multiaddr: multiaddr(`/ip4/58.42.62.62/tcp/${Math.random() * (maxPort - minPort) + minPort}`)
+      multiaddr: multiaddr(`/ip4/58.42.62.62/tcp/${Math.round(Math.random() * (maxPort - minPort) + minPort)}`)
     }],
     tags: new Map(),
     metadata: new Map(),