@aztec/p2p 0.0.0-test.0 → 0.0.1-commit.023c3e5

package/src/services/reqresp/connection-sampler/batch_connection_sampler.ts

@@ -18,11 +18,18 @@ import type { ConnectionSampler } from './connection_sampler.js';
  * If a peer fails, it is replaced while maintaining the same bucket.
  */
 export class BatchConnectionSampler {
-  private readonly logger = createLogger('p2p:reqresp:batch-connection-sampler');
   private readonly batch: PeerId[] = [];
   private readonly requestsPerPeer: number;
+  /** Tracks peer-index combinations that returned empty/invalid responses */
+  private readonly failedPeerIndices: Map<string, Set<number>> = new Map();
 
-  constructor(private readonly connectionSampler: ConnectionSampler, batchSize: number, maxPeers: number) {
+  constructor(
+    private readonly connectionSampler: ConnectionSampler,
+    batchSize: number,
+    maxPeers: number,
+    exclude?: PeerId[],
+    private readonly logger = createLogger('p2p:reqresp:batch-connection-sampler'),
+  ) {
     if (maxPeers <= 0) {
       throw new Error('Max peers cannot be 0');
     }
@@ -34,14 +41,17 @@ export class BatchConnectionSampler {
     this.requestsPerPeer = Math.max(1, Math.floor(batchSize / maxPeers));
 
     // Sample initial peers
-    this.batch = this.connectionSampler.samplePeersBatch(maxPeers);
+    const excluding = exclude && new Map(exclude.map(peerId => [peerId.toString(), true] as const));
+    this.batch = this.connectionSampler.samplePeersBatch(maxPeers, excluding);
   }
 
   /**
-   * Gets the peer responsible for handling a specific request index
+   * Gets the peer responsible for handling a specific request index.
+   * If the primary peer has previously failed for this index, tries other peers.
+   * If all batch peers have failed, attempts to sample a new peer.
    *
    * @param index - The request index
-   * @returns The peer assigned to handle this request
+   * @returns The peer assigned to handle this request, or undefined if no peer available
    */
   getPeerForRequest(index: number): PeerId | undefined {
     if (this.batch.length === 0) {
@@ -49,8 +59,65 @@ export class BatchConnectionSampler {
     }
 
     // Calculate which peer bucket this index belongs to
-    const peerIndex = Math.floor(index / this.requestsPerPeer) % this.batch.length;
-    return this.batch[peerIndex];
+    const primaryPeerIndex = Math.floor(index / this.requestsPerPeer) % this.batch.length;
+
+    // Try peers starting from primary, wrapping around
+    for (let offset = 0; offset < this.batch.length; offset++) {
+      const peerIndex = (primaryPeerIndex + offset) % this.batch.length;
+      const peer = this.batch[peerIndex];
+      const peerKey = peer.toString();
+
+      const failedIndices = this.failedPeerIndices.get(peerKey);
+      if (!failedIndices || !failedIndices.has(index)) {
+        return peer;
+      }
+    }
+
+    // All batch peers have failed for this index - try to sample a new peer
+    const newPeer = this.sampleNewPeer();
+    if (newPeer) {
+      return newPeer;
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Attempts to sample a new peer that isn't already in the batch.
+   * If successful, adds the peer to the batch.
+   *
+   * @returns The new peer if one was sampled, undefined otherwise
+   */
+  private sampleNewPeer(): PeerId | undefined {
+    // Exclude all current batch peers
+    const excluding = new Map(this.batch.map(p => [p.toString(), true] as const));
+    const newPeer = this.connectionSampler.getPeer(excluding);
+
+    if (newPeer) {
+      this.batch.push(newPeer);
+      this.logger.trace('Sampled new peer for exhausted index', { newPeer: newPeer.toString() });
+      return newPeer;
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Marks that a peer returned an empty/invalid response for a specific request index.
+   * The peer will not be assigned this index again.
+   *
+   * @param peerId - The peer that failed
+   * @param index - The request index that failed
+   */
+  markPeerFailedForIndex(peerId: PeerId, index: number): void {
+    const peerKey = peerId.toString();
+    let failedIndices = this.failedPeerIndices.get(peerKey);
+    if (!failedIndices) {
+      failedIndices = new Set();
+      this.failedPeerIndices.set(peerKey, failedIndices);
+    }
+    failedIndices.add(index);
+    this.logger.trace('Marked peer failed for index', { peerId: peerKey, index });
   }
 
   /**
@@ -66,15 +133,15 @@ export class BatchConnectionSampler {
     }
 
     const excluding = new Map([[peerId.toString(), true]]);
-    const newPeer = this.connectionSampler.getPeer(excluding);
+    const newPeer = this.connectionSampler.getPeer(excluding); // Q: Shouldn't we accumulate all excluded peers? Otherwise the sampler could return us a previously excluded peer?
 
     if (newPeer) {
       this.batch[index] = newPeer;
-      this.logger.trace(`Replaced peer ${peerId} with ${newPeer}`, { peerId, newPeer });
+      this.logger.trace('Replaced peer', { peerId, newPeer });
     } else {
       // If we couldn't get a replacement, remove the peer and compact the array
       this.batch.splice(index, 1);
-      this.logger.trace(`Removed peer ${peerId}`, { peerId });
+      this.logger.trace('Removed peer', { peerId });
     }
   }
 

package/src/services/reqresp/connection-sampler/connection_sampler.ts

@@ -1,15 +1,10 @@
+import { AbortError } from '@aztec/foundation/error';
 import { createLogger } from '@aztec/foundation/log';
-import { SerialQueue } from '@aztec/foundation/queue';
 
 import type { Libp2p, PeerId, Stream } from '@libp2p/interface';
 
 const MAX_SAMPLE_ATTEMPTS = 4;
 
-interface StreamAndPeerId {
-  stream: Stream;
-  peerId: PeerId;
-}
-
 export class RandomSampler {
   random(max: number) {
     return Math.floor(Math.random() * max);
@@ -17,31 +12,33 @@ export class RandomSampler {
 }
 
 /**
- * A class that samples peers from the libp2p node and returns a peer that we don't already have a connection open to.
+ * A class that samples peers from the libp2p node and returns a peer that we don't already have a reqresp connection open to.
  * If we already have a connection open, we try to sample a different peer.
  * We do this MAX_SAMPLE_ATTEMPTS times, if we still don't find a peer we just go for it.
  *
  * @dev Close must always be called on connections, else memory leak
  */
 export class ConnectionSampler {
-  private readonly logger = createLogger('p2p:reqresp:connection-sampler');
   private cleanupInterval: NodeJS.Timeout;
-  private abortController: AbortController = new AbortController();
 
-  private readonly activeConnectionsCount: Map<PeerId, number> = new Map();
-  private readonly streams: Map<string, StreamAndPeerId> = new Map();
+  // Map from stringified peer id to number of active connections
+  protected readonly activeConnectionsCount: Map<string, number> = new Map();
+
+  // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
+  protected readonly streams: Set<Stream> = new Set();
 
-  // Serial queue to ensure that we only dial one peer at a time
-  private dialQueue: SerialQueue = new SerialQueue();
+  private abortOnStop: AbortController = new AbortController();
 
   constructor(
     private readonly libp2p: Libp2p,
-    private readonly cleanupIntervalMs: number = 60000, // Default to 1 minute
-    private readonly sampler: RandomSampler = new RandomSampler(), // Allow randomness to be mocked for testing
+    private readonly sampler: RandomSampler,
+    private readonly logger = createLogger('p2p:reqresp:connection-sampler'),
+    private readonly opts: { cleanupIntervalMs?: number; p2pOptimisticNegotiation?: boolean } = {},
   ) {
-    this.cleanupInterval = setInterval(() => void this.cleanupStaleConnections(), this.cleanupIntervalMs);
-
-    this.dialQueue.start();
+    this.cleanupInterval = setInterval(
+      () => void this.cleanupStaleConnections(),
+      this.opts.cleanupIntervalMs ?? 60_000,
+    );
   }
 
   /**
@@ -49,13 +46,11 @@ export class ConnectionSampler {
    */
   async stop() {
     this.logger.info('Stopping connection sampler');
+    this.abortOnStop.abort(new AbortError('Connection sampler stopped'));
     clearInterval(this.cleanupInterval);
 
-    this.abortController.abort();
-    await this.dialQueue.end();
-
     // Close all active streams
-    const closePromises = Array.from(this.streams.keys()).map(streamId => this.close(streamId));
+    const closePromises = Array.from(this.streams.values()).map(stream => this.close(stream));
     await Promise.all(closePromises);
     this.logger.info('Connection sampler stopped');
   }
@@ -69,70 +64,128 @@ export class ConnectionSampler {
   getPeer(excluding?: Map<string, boolean>): PeerId | undefined {
     // In libp2p getPeers performs a shallow copy, so this array can be sliced from safetly
     const peers = this.libp2p.getPeers();
+    const { peer } = this.getPeerFromList(peers, excluding);
+    return peer;
+  }
 
+  /**
+   * Samples a peer from a list of peers, excluding those that have active (reqresp) connections or are in the exclusion list
+   *
+   * @param peers - The list of peers to sample from
+   * @param excluding - The peers to exclude from the sampling
+   * @returns - A peer from the list, or undefined if no peers are available,
+   *          - a boolean indicating if the peer has active connections, and
+   *          - all sampled peers - to enable optional resampling
+   *
+   * @dev The provided list peers, should be mutated by this function. This allows batch sampling
+   *      to be performed without making extra copies of the list.
+   */
+  getPeerFromList(
+    peers: PeerId[],
+    excluding?: Map<string, boolean>,
+  ): {
+    peer: PeerId | undefined;
+    sampledPeers: PeerId[];
+  } {
     if (peers.length === 0) {
-      return undefined;
+      return { peer: undefined, sampledPeers: [] };
     }
 
-    let randomIndex = this.sampler.random(peers.length);
-    let attempts = 0;
-
-    // Keep sampling while:
-    // - we haven't exceeded max attempts AND
-    // - either the peer has active connections OR is in the exclusion list
-    while (
-      attempts < MAX_SAMPLE_ATTEMPTS &&
-      ((this.activeConnectionsCount.get(peers[randomIndex]) ?? 0) > 0 ||
-        (excluding?.get(peers[randomIndex]?.toString()) ?? false))
-    ) {
+    const sampledPeers: PeerId[] = [];
+    // Try to find a peer that has no active connections and is not in the exclusion list
+    for (let attempts = 0; attempts < MAX_SAMPLE_ATTEMPTS && peers.length > 0; attempts++) {
+      const randomIndex = this.sampler.random(peers.length);
+      const peer = peers[randomIndex];
+      const hasActiveConnections = (this.activeConnectionsCount.get(peer.toString()) ?? 0) > 0;
+      const isExcluded = excluding?.get(peer.toString()) ?? false;
+
+      // Remove this peer from consideration
       peers.splice(randomIndex, 1);
-      randomIndex = this.sampler.random(peers.length);
-      attempts++;
+
+      // If peer is suitable (no active connections and not excluded), return it
+      if (!hasActiveConnections && !isExcluded) {
+        this.logger.trace('Sampled peer', {
+          attempts,
+          peer,
+        });
+        return { peer, sampledPeers };
+      }
+
+      // Keep track of peers that have active reqresp channels, batch sampling will use these to resample
+      sampledPeers.push(peer);
     }
 
-    this.logger.trace(`Sampled peer in ${attempts} attempts`, {
-      attempts,
-      peer: peers[randomIndex]?.toString(),
+    // If we've exhausted our attempts or peers list is empty, return the last peer if available
+    const lastPeer = peers.length > 0 ? peers[this.sampler.random(peers.length)] : undefined;
+
+    this.logger.trace('Sampled peer', {
+      attempts: MAX_SAMPLE_ATTEMPTS,
+      peer: lastPeer?.toString(),
     });
-    return peers[randomIndex];
+    return { peer: lastPeer, sampledPeers };
+  }
+
+  /*
+   * Returns all peers sorted by connection count ascending,
+   * meaning that the peers with the least number of active connections are earlier in an array
+   *
+   * @param: excluding - peers to exclude
+   * @return: list of peer ids
+   * */
+  public getPeerListSortedByConnectionCountAsc(excluding?: Set<string>): PeerId[] {
+    return this.libp2p
+      .getPeers()
+      .filter(id => !excluding?.has(id.toString()))
+      .map(id => ({ id, count: this.activeConnectionsCount.get(id.toString()) ?? 0 }))
+      .sort((a, b) => a.count - b.count)
+      .map(p => p.id);
   }
 
   /**
    * Samples a batch of unique peers from the libp2p node, prioritizing peers without active connections
    *
    * @param numberToSample - The number of peers to sample
+   * @param excluding - The peers to exclude from the sampling
    * @returns Array of unique sampled peers, prioritizing those without active connections
    */
-  samplePeersBatch(numberToSample: number): PeerId[] {
+  samplePeersBatch(numberToSample: number, excluding?: Map<string, boolean>): PeerId[] {
     const peers = this.libp2p.getPeers();
-    const sampledPeers: PeerId[] = [];
-    const peersWithConnections: PeerId[] = []; // Hold onto peers with active connections incase we need to sample more
-
-    for (const peer of peers) {
-      const activeConnections = this.activeConnectionsCount.get(peer) ?? 0;
-      if (activeConnections === 0) {
-        if (sampledPeers.push(peer) === numberToSample) {
-          return sampledPeers;
-        }
-      } else {
-        peersWithConnections.push(peer);
+    this.logger.debug('Sampling peers batch', { numberToSample, peers });
+
+    // Only sample as many peers as we have available
+    numberToSample = Math.min(numberToSample, peers.length);
+
+    const batch: PeerId[] = [];
+    // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
+    const withActiveConnections: Set<PeerId> = new Set();
+    for (let i = 0; i < numberToSample; i++) {
+      const { peer, sampledPeers } = this.getPeerFromList(peers, excluding);
+      if (peer) {
+        batch.push(peer);
+      }
+      if (sampledPeers.length > 0) {
+        sampledPeers.forEach(peer => withActiveConnections.add(peer));
       }
     }
+    const lengthWithoutConnections = batch.length;
 
     // If we still need more peers, sample from those with connections
-    while (sampledPeers.length < numberToSample && peersWithConnections.length > 0) {
-      const randomIndex = this.sampler.random(peersWithConnections.length);
-      const [peer] = peersWithConnections.splice(randomIndex, 1);
-      sampledPeers.push(peer);
+    while (batch.length < numberToSample && withActiveConnections.size > 0) {
+      const randomIndex = this.sampler.random(withActiveConnections.size);
+
+      const peer = Array.from(withActiveConnections)[randomIndex];
+      withActiveConnections.delete(peer);
+      batch.push(peer);
     }
 
-    this.logger.trace(`Batch sampled ${sampledPeers.length} unique peers`, {
-      peers: sampledPeers,
-      withoutConnections: sampledPeers.length - peersWithConnections.length,
-      withConnections: peersWithConnections.length,
+    this.logger.trace('Batch sampled peers', {
+      length: batch.length,
+      peers: batch,
+      withoutConnections: lengthWithoutConnections,
+      withConnections: numberToSample - lengthWithoutConnections,
     });
 
-    return sampledPeers;
+    return batch;
   }
 
   // Set of passthrough functions to keep track of active connections
@@ -142,22 +195,27 @@ export class ConnectionSampler {
    *
    * @param peerId - The peer id
    * @param protocol - The protocol
+   * @param timeout - Abort connection if it takes too long
    * @returns The stream
    */
-  async dialProtocol(peerId: PeerId, protocol: string): Promise<Stream> {
-    // Dialling at the same time can cause race conditions where two different streams
-    // end up with the same id, hence a serial queue
-    const stream = await this.dialQueue.put(() =>
-      this.libp2p.dialProtocol(peerId, protocol, { signal: this.abortController.signal }),
-    );
+  async dialProtocol(peerId: PeerId, protocol: string, timeout?: number): Promise<Stream> {
+    const stream = await this.libp2p.dialProtocol(peerId, protocol, {
+      signal: AbortSignal.any(
+        timeout ? [this.abortOnStop.signal, AbortSignal.timeout(timeout!)] : [this.abortOnStop.signal],
+      ),
+      negotiateFully: !this.opts.p2pOptimisticNegotiation,
+    });
+    stream.metadata.peerId = peerId;
+    this.streams.add(stream);
 
-    this.streams.set(stream.id, { stream, peerId });
-    const updatedActiveConnectionsCount = (this.activeConnectionsCount.get(peerId) ?? 0) + 1;
-    this.activeConnectionsCount.set(peerId, updatedActiveConnectionsCount);
+    const peerIdString = peerId.toString();
+    const updatedActiveConnectionsCount = (this.activeConnectionsCount.get(peerIdString) ?? 0) + 1;
+    this.activeConnectionsCount.set(peerIdString, updatedActiveConnectionsCount);
 
-    this.logger.trace(`Dialed protocol ${protocol} with peer ${peerId.toString()}`, {
+    this.logger.trace('Dialed protocol', {
       streamId: stream.id,
-      peerId: peerId.toString(),
+      protocol,
+      peerId: peerIdString,
       activeConnectionsCount: updatedActiveConnectionsCount,
     });
     return stream;
@@ -165,34 +223,39 @@ export class ConnectionSampler {
 
   /**
    * Closes a stream and updates the active connections count
-   *
-   * @param streamId - The stream id
    */
-  async close(streamId: string): Promise<void> {
-    try {
-      const streamAndPeerId = this.streams.get(streamId);
-      if (!streamAndPeerId) {
-        this.logger.warn(`Stream ${streamId} not found`);
-        return;
-      }
+  async close(stream: Stream): Promise<void> {
+    let peerId: PeerId | undefined = undefined;
 
-      const { stream, peerId } = streamAndPeerId;
+    try {
+      peerId = stream.metadata.peerId;
+      let updatedActiveConnectionsCount = undefined;
 
-      const updatedActiveConnectionsCount = (this.activeConnectionsCount.get(peerId) ?? 1) - 1;
-      this.activeConnectionsCount.set(peerId, updatedActiveConnectionsCount);
+      if (!peerId) {
+        this.logger.warn(`Stream ${stream.id} does not have a peerId set`);
+      } else {
+        updatedActiveConnectionsCount = (this.activeConnectionsCount.get(peerId.toString()) ?? 1) - 1;
+        this.activeConnectionsCount.set(peerId.toString(), updatedActiveConnectionsCount);
+      }
 
-      this.logger.trace(`Closing connection to peer ${peerId.toString()}`, {
-        streamId,
-        peerId: peerId.toString(),
+      this.logger.trace('Closing connection', {
+        streamId: stream.id,
+        peerId: peerId?.toString(),
         protocol: stream.protocol,
         activeConnectionsCount: updatedActiveConnectionsCount,
       });
 
-      await stream?.close();
+      if (!this.streams.has(stream)) {
+        this.logger.debug(`Stream ${stream.id} is not in the active streams set`);
+      }
+
+      await stream.close();
     } catch (error) {
-      this.logger.error(`Failed to close connection to peer with stream id ${streamId}`, error);
+      this.logger.error(`Failed to close connection to peer ${peerId ?? 'unknown'} with stream id ${stream.id}`, error);
+      // graceful close failed, abort the stream
+      stream.abort(new AbortError('Failed to close stream gracefully'));
     } finally {
-      this.streams.delete(streamId);
+      this.streams.delete(stream);
     }
   }
 
@@ -202,15 +265,23 @@ export class ConnectionSampler {
   private async cleanupStaleConnections() {
     // Look for streams without anything in the activeConnectionsCount
     // If we find anything, close the stream
-    for (const [streamId, { peerId }] of this.streams.entries()) {
+    for (const stream of this.streams.values()) {
       try {
         // Check if we have lost track of accounting
-        if (this.activeConnectionsCount.get(peerId) === 0) {
-          await this.close(streamId);
-          this.logger.debug(`Cleaned up stale connection ${streamId} to peer ${peerId.toString()}`);
+        const peerId: PeerId = stream.metadata.peerId;
+        if (!peerId) {
+          this.logger.warn(`Stream ${stream.id} does not have a peerId set`);
+        } else if (this.activeConnectionsCount.get(peerId.toString()) === 0) {
+          await this.close(stream);
+          this.logger.debug('Cleaned up stale connection', { streamId: stream.id, peerId: peerId.toString() });
         }
       } catch (error) {
-        this.logger.error(`Error cleaning up stale connection ${streamId}`, { error });
+        this.logger.error(
+          `Error cleaning up stale connection to peer ${stream.metadata.peerId?.toString() ?? 'unknown'} stream ${
+            stream.id
+          }`,
+          { error },
+        );
       }
     }
   }

package/src/services/reqresp/constants.ts

@@ -0,0 +1,14 @@
+/**
+ * Constants for P2P message deserialization bounds checking.
+ * These constants define maximum allowed sizes during deserialization
+ * to prevent DoS attacks via maliciously crafted messages.
+ */
+
+/** Max transactions per block for deserialization validation (~300x default of 32) */
+export { MAX_TXS_PER_BLOCK } from '@aztec/stdlib/deserialization';
+
+/** Max version string length (e.g., "1.0.0-alpha.123") */
+export const MAX_VERSION_STRING_LENGTH = 64;
+
+/** Max block hash string length (hex: 0x + 64 chars, with generous headroom) */
+export const MAX_BLOCK_HASH_STRING_LENGTH = 128;

package/src/services/reqresp/index.ts

@@ -2,3 +2,5 @@
  * Request Response protocol allows nodes to ask their peers for data
  * that they missed via the traditional gossip protocol.
  */
+export * from './protocols/index.js';
+export * from './interface.js';