fix: ensure dht query is aborted on early exit

If query results are consumed from a `for await..of`-style loop,
and that loop is exited from before the results are complete, ensure
we abort any running sub-queries.

Author: achingbrain

packages/kad-dht/src/query/manager.ts

@@ -1,7 +1,8 @@
-import { AbortError, TypedEventEmitter, CustomEvent, setMaxListeners } from '@libp2p/interface'
+import { TypedEventEmitter, CustomEvent, setMaxListeners } from '@libp2p/interface'
 import { PeerSet } from '@libp2p/peer-collections'
 import { anySignal } from 'any-signal'
 import merge from 'it-merge'
+import { raceSignal } from 'race-signal'
 import { toString as uint8ArrayToString } from 'uint8arrays/to-string'
 import {
   ALPHA, K, DEFAULT_QUERY_TIMEOUT
@@ -127,7 +128,16 @@ export class QueryManager implements Startable {
       }
     }
 
-    const signal = anySignal([this.shutDownController.signal, options.signal])
+    // if the user breaks out of a for..await of loop iterating over query
+    // results we need to cancel any in-flight network requests
+    const queryEarlyExitController = new AbortController()
+    setMaxListeners(Infinity, queryEarlyExitController.signal)
+
+    const signal = anySignal([
+      this.shutDownController.signal,
+      queryEarlyExitController.signal,
+      options.signal
+    ])
 
     // this signal will get listened to for every invocation of queryFunc
     // so make sure we don't make a lot of noise in the logs
@@ -138,19 +148,13 @@ export class QueryManager implements Startable {
     // query a subset of peers up to `kBucketSize / 2` in length
     const startTime = Date.now()
     const cleanUp = new TypedEventEmitter<CleanUpEvents>()
+    let queryFinished = false
 
     try {
       if (options.isSelfQuery !== true && this.initialQuerySelfHasRun != null) {
         log('waiting for initial query-self query before continuing')
 
-        await Promise.race([
-          new Promise((resolve, reject) => {
-            signal.addEventListener('abort', () => {
-              reject(new AbortError('Query was aborted before self-query ran'))
-            })
-          }),
-          this.initialQuerySelfHasRun.promise
-        ])
+        await raceSignal(this.initialQuerySelfHasRun.promise, signal)
 
         this.initialQuerySelfHasRun = undefined
       }
@@ -192,19 +196,26 @@ export class QueryManager implements Startable {
 
       // Execute the query along each disjoint path and yield their results as they become available
       for await (const event of merge(...paths)) {
-        yield event
-
         if (event.name === 'QUERY_ERROR') {
-          log('error', event.error)
+          log.error('query error', event.error)
         }
+
+        yield event
       }
+
+      queryFinished = true
     } catch (err: any) {
       if (!this.running && err.code === 'ERR_QUERY_ABORTED') {
         // ignore query aborted errors that were thrown during query manager shutdown
       } else {
         throw err
       }
     } finally {
+      if (!queryFinished) {
+        log('query exited early')
+        queryEarlyExitController.abort()
+      }
+
       signal.clear()
 
       this.queries--

packages/kad-dht/test/query.spec.ts

@@ -19,7 +19,7 @@ import { QueryManager, type QueryManagerInit } from '../src/query/manager.js'
 import { convertBuffer } from '../src/utils.js'
 import { createPeerId, createPeerIds } from './utils/create-peer-id.js'
 import { sortClosestPeers } from './utils/sort-closest-peers.js'
-import type { QueryFunc } from '../src/query/types.js'
+import type { QueryContext, QueryFunc } from '../src/query/types.js'
 import type { RoutingTable } from '../src/routing-table/index.js'
 import type { PeerId } from '@libp2p/interface'
 
@@ -29,12 +29,9 @@ interface TopologyEntry {
   value?: Uint8Array
   closerPeers?: number[]
   event: QueryEvent
+  context?: QueryContext
 }
-type Topology = Record<string, {
-  delay?: number | undefined
-  error?: Error | undefined
-  event: QueryEvent
-}>
+type Topology = Record<string, TopologyEntry>
 
 describe('QueryManager', () => {
   let ourPeerId: PeerId
@@ -55,7 +52,7 @@ describe('QueryManager', () => {
   }
 
   function createTopology (opts: Record<number, { delay?: number, error?: Error, value?: Uint8Array, closerPeers?: number[] }>): Topology {
-    const topology: Record<string, { delay?: number, error?: Error, event: QueryEvent }> = {}
+    const topology: Topology = {}
 
     Object.keys(opts).forEach(key => {
       const id = parseInt(key)
@@ -94,9 +91,12 @@ describe('QueryManager', () => {
     return topology
   }
 
-  function createQueryFunction (topology: Record<string, { delay?: number, event: QueryEvent }>): QueryFunc {
-    const queryFunc: QueryFunc = async function * ({ peer }) {
+  function createQueryFunction (topology: Topology): QueryFunc {
+    const queryFunc: QueryFunc = async function * (context) {
+      const { peer } = context
+
       const res = topology[peer.toString()]
+      res.context = context
 
       if (res.delay != null) {
         await delay(res.delay)
@@ -870,4 +870,43 @@ describe('QueryManager', () => {
 
     await manager.stop()
   })
+
+  it('should abort the query if we break out of the loop early', async () => {
+    const manager = new QueryManager({
+      peerId: ourPeerId,
+      logger: defaultLogger()
+    }, {
+      ...defaultInit(),
+      disjointPaths: 2
+    })
+    await manager.start()
+
+    // 1 -> 0 [pathComplete]
+    // 4 -> 3 [delay] -> 2 [pathComplete]
+    const topology = createTopology({
+      // quick value path
+      0: { value: uint8ArrayFromString('true') },
+      1: { closerPeers: [0] },
+      // slow value path
+      2: { value: uint8ArrayFromString('true') },
+      3: { delay: 100, closerPeers: [2] },
+      4: { closerPeers: [3] }
+    })
+
+    routingTable.closestPeers.returns([peers[1], peers[4]])
+
+    for await (const event of manager.run(key, createQueryFunction(topology))) {
+      if (event.name === 'VALUE') {
+        expect(event.from.toString()).to.equal(peers[0].toString())
+
+        // break out of loop early
+        break
+      }
+    }
+
+    // should have aborted query on slow path
+    expect(topology[peers[3].toString()]).to.have.nested.property('context.signal.aborted', true)
+
+    await manager.stop()
+  })
 })