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

packages/libp2p/src/connection-manager/dial-queue.ts

@@ -15,8 +15,8 @@ import {
   MAX_PEER_ADDRS_TO_DIAL,
   LAST_DIAL_FAILURE_KEY
 } from './constants.js'
-import { combineSignals, resolveMultiaddrs } from './utils.js'
-import type { AddressSorter, AbortOptions, PendingDial, ComponentLogger, Logger, Connection, ConnectionGater, Metric, Metrics, PeerId, Address, PeerStore } from '@libp2p/interface'
+import { resolveMultiaddrs } from './utils.js'
+import type { AddressSorter, AbortOptions, ComponentLogger, Logger, Connection, ConnectionGater, Metrics, PeerId, Address, PeerStore, PeerRouting } from '@libp2p/interface'
 import type { TransportManager } from '@libp2p/interface-internal'
 
 export interface PendingDialTarget {
@@ -59,15 +59,11 @@ interface DialQueueComponents {
   transportManager: TransportManager
   connectionGater: ConnectionGater
   logger: ComponentLogger
+  peerRouting: PeerRouting
 }
 
 export class DialQueue {
-  public pendingDials: PendingDialInternal[]
-  public queue: PQueue
-  private readonly peerId: PeerId
-  private readonly peerStore: PeerStore
-  private readonly connectionGater: ConnectionGater
-  private readonly transportManager: TransportManager
+  public queue: Queue<Connection, DialQueueJobOptions>
   private readonly addressSorter: AddressSorter
   private readonly maxPeerAddrsToDial: number
   private readonly dialTimeout: number
@@ -76,6 +72,7 @@ export class DialQueue {
   private shutDownController: AbortController
   private readonly connections: PeerMap<Connection[]>
   private readonly log: Logger
+  private readonly components: DialQueueComponents
 
   constructor (components: DialQueueComponents, init: DialerInit = {}) {
     this.addressSorter = init.addressSorter ?? defaultOptions.addressSorter
@@ -84,10 +81,7 @@ export class DialQueue {
     this.connections = init.connections ?? new PeerMap()
     this.log = components.logger.forComponent('libp2p:connection-manager:dial-queue')
 
-    this.peerId = components.peerId
-    this.peerStore = components.peerStore
-    this.connectionGater = components.connectionGater
-    this.transportManager = components.transportManager
+    this.components = components
     this.shutDownController = new AbortController()
 
     setMaxListeners(Infinity, this.shutDownController.signal)
@@ -242,6 +236,75 @@ export class DialQueue {
         // remove our pending dial entry
         this.pendingDials = this.pendingDials.filter(p => p.id !== pendingDial.id)
 
+      try {
+        // load addresses from address book, resolve and dnsaddrs, filter
+        // undiallables, add peer IDs, etc
+        addrsToDial = await this.calculateMultiaddrs(peerId, options?.multiaddrs, {
+          ...options,
+          signal
+        })
+
+        addrsToDial.map(({ multiaddr }) => multiaddr.toString()).forEach(addr => {
+          options?.multiaddrs.add(addr)
+        })
+      } catch (err) {
+        signal.clear()
+        throw err
+      }
+
+      try {
+        let dialed = 0
+        const errors: Error[] = []
+
+        for (const address of addrsToDial) {
+          if (dialed === this.maxPeerAddrsToDial) {
+            this.log('dialed maxPeerAddrsToDial (%d) addresses for %p, not trying any others', dialed, peerId)
+
+            throw new CodeError('Peer had more than maxPeerAddrsToDial', codes.ERR_TOO_MANY_ADDRESSES)
+          }
+
+          dialed++
+
+          try {
+            const conn = await this.components.transportManager.dial(address.multiaddr, {
+              ...options,
+              signal
+            })
+
+            this.log('dial to %a succeeded', address.multiaddr)
+
+            return conn
+          } catch (err: any) {
+            this.log.error('dial failed to %a', address.multiaddr, err)
+
+            if (peerId != null) {
+              // record the failed dial
+              try {
+                await this.components.peerStore.patch(peerId, {
+                  metadata: {
+                    [LAST_DIAL_FAILURE_KEY]: uint8ArrayFromString(Date.now().toString())
+                  }
+                })
+              } catch (err: any) {
+                this.log.error('could not update last dial failure key for %p', peerId, err)
+              }
+            }
+
+            // the user/dial timeout/shutdown controller signal aborted
+            if (signal.aborted) {
+              throw new CodeError(err.message, ERR_TIMEOUT)
+            }
+
+            errors.push(err)
+          }
+        }
+
+        if (errors.length === 1) {
+          throw errors[0]
+        }
+
+        throw new AggregateCodeError(errors, 'All multiaddr dials failed', codes.ERR_TRANSPORT_DIAL_FAILED)
+      } finally {
         // clean up abort signals/controllers
         signal.clear()
       })
@@ -315,19 +378,20 @@ export class DialQueue {
   private async calculateMultiaddrs (peerId?: PeerId, addrs: Address[] = [], options: DialOptions = {}): Promise<Address[]> {
     // if a peer id or multiaddr(s) with a peer id, make sure it isn't our peer id and that we are allowed to dial it
     if (peerId != null) {
-      if (this.peerId.equals(peerId)) {
+      if (this.components.peerId.equals(peerId)) {
         throw new CodeError('Tried to dial self', codes.ERR_DIALED_SELF)
       }
 
-      if ((await this.connectionGater.denyDialPeer?.(peerId)) === true) {
+      if ((await this.components.connectionGater.denyDialPeer?.(peerId)) === true) {
         throw new CodeError('The dial request is blocked by gater.allowDialPeer', codes.ERR_PEER_DIAL_INTERCEPTED)
       }
 
-      // if just a peer id was passed, load available multiaddrs for this peer from the address book
+      // if just a peer id was passed, load available multiaddrs for this peer
+      // from the peer store
       if (addrs.length === 0) {
         this.log('loading multiaddrs for %p', peerId)
         try {
-          const peer = await this.peerStore.get(peerId)
+          const peer = await this.components.peerStore.get(peerId)
           addrs.push(...peer.addresses)
           this.log('loaded multiaddrs for %p', peerId, addrs.map(({ multiaddr }) => multiaddr.toString()))
         } catch (err: any) {
@@ -336,9 +400,23 @@ export class DialQueue {
           }
         }
       }
+
+      // if the peer store had no addresses for the peer, try to find them via
+      // peer routing
+      if (addrs.length === 0) {
+        const peer = await this.components.peerRouting.findPeer(peerId, options)
+
+        peer.multiaddrs.forEach(multiaddr => {
+          addrs.push({
+            multiaddr,
+            isCertified: false
+          })
+        })
+      }
     }
 
-    // resolve addresses - this can result in a one-to-many translation when dnsaddrs are resolved
+    // resolve addresses - this can result in a one-to-many translation when
+    // dnsaddrs are resolved
     let resolvedAddresses = (await Promise.all(
       addrs.map(async addr => {
         const result = await resolveMultiaddrs(addr.multiaddr, {
@@ -383,7 +461,7 @@ export class DialQueue {
 
     const filteredAddrs = resolvedAddresses.filter(addr => {
       // filter out any multiaddrs that we do not have transports for
-      if (this.transportManager.transportForMultiaddr(addr.multiaddr) == null) {
+      if (this.components.transportManager.transportForMultiaddr(addr.multiaddr) == null) {
         return false
       }
 
@@ -433,7 +511,7 @@ export class DialQueue {
     const gatedAdrs: Address[] = []
 
     for (const addr of dedupedMultiaddrs) {
-      if (this.connectionGater.denyDialMultiaddr != null && await this.connectionGater.denyDialMultiaddr(addr.multiaddr)) {
+      if (this.components.connectionGater.denyDialMultiaddr != null && await this.components.connectionGater.denyDialMultiaddr(addr.multiaddr)) {
         continue
       }