@aztec/p2p 0.0.0-test.1 → 0.0.1-commit.03f7ef2

package/src/mem_pools/attestation_pool/kv_attestation_pool.ts

@@ -1,18 +1,27 @@
-import { Fr } from '@aztec/foundation/fields';
+import { SlotNumber } from '@aztec/foundation/branded-types';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { toArray } from '@aztec/foundation/iterable';
 import { createLogger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore, AztecAsyncMap, AztecAsyncMultiMap } from '@aztec/kv-store';
-import { BlockAttestation } from '@aztec/stdlib/p2p';
+import { BlockAttestation, BlockProposal } from '@aztec/stdlib/p2p';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
-import { PoolInstrumentation, PoolName } from '../instrumentation.js';
+import { ProposalSlotCapExceededError } from '../../errors/attestation-pool.error.js';
+import { PoolInstrumentation, PoolName, type PoolStatsCallback } from '../instrumentation.js';
 import type { AttestationPool } from './attestation_pool.js';
 
+export const MAX_PROPOSALS_PER_SLOT = 5;
+export const ATTESTATION_CAP_BUFFER = 10;
+
 export class KvAttestationPool implements AttestationPool {
   private metrics: PoolInstrumentation<BlockAttestation>;
 
   private attestations: AztecAsyncMap<string, Buffer>;
-  private proposalsForSlot: AztecAsyncMultiMap<string, string>;
+  private proposals: AztecAsyncMap<
+    /*  proposal.payload.archive  */ string,
+    /* buffer representation of proposal */ Buffer
+  >;
+  private proposalsForSlot: AztecAsyncMultiMap<number, string>;
   private attestationsForProposal: AztecAsyncMultiMap<string, string>;
 
   constructor(
@@ -21,10 +30,24 @@ export class KvAttestationPool implements AttestationPool {
     private log = createLogger('aztec:attestation_pool'),
   ) {
     this.attestations = store.openMap('attestations');
+    this.proposals = store.openMap('proposals');
     this.proposalsForSlot = store.openMultiMap('proposals_for_slot');
     this.attestationsForProposal = store.openMultiMap('attestations_for_proposal');
 
-    this.metrics = new PoolInstrumentation(telemetry, PoolName.ATTESTATION_POOL);
+    this.metrics = new PoolInstrumentation(telemetry, PoolName.ATTESTATION_POOL, this.poolStats);
+  }
+
+  private poolStats: PoolStatsCallback = async () => {
+    return {
+      itemCount: await this.attestations.sizeAsync(),
+    };
+  };
+
+  public async isEmpty(): Promise<boolean> {
+    for await (const _ of this.attestations.entriesAsync()) {
+      return false;
+    }
+    return true;
   }
 
   private getProposalKey(slot: number | bigint | Fr | string, proposalId: Fr | string | Buffer): string {
@@ -33,8 +56,8 @@ export class KvAttestationPool implements AttestationPool {
       typeof proposalId === 'string'
         ? proposalId
         : Buffer.isBuffer(proposalId)
-        ? Fr.fromBuffer(proposalId).toString()
-        : proposalId.toString();
+          ? Fr.fromBuffer(proposalId).toString()
+          : proposalId.toString();
 
     return `${slotStr}-${proposalIdStr}`;
   }
@@ -46,26 +69,52 @@ export class KvAttestationPool implements AttestationPool {
   public async addAttestations(attestations: BlockAttestation[]): Promise<void> {
     await this.store.transactionAsync(async () => {
       for (const attestation of attestations) {
-        const slotNumber = attestation.payload.header.globalVariables.slotNumber;
+        const slotNumber = attestation.payload.header.slotNumber;
         const proposalId = attestation.archive;
-        const address = (await attestation.getSender()).toString();
+        const sender = attestation.getSender();
+
+        // Skip attestations with invalid signatures
+        if (!sender) {
+          this.log.warn(`Skipping attestation with invalid signature for slot ${slotNumber}`, {
+            signature: attestation.signature.toString(),
+            slotNumber,
+            proposalId,
+          });
+          continue;
+        }
+
+        const address = sender.toString();
 
         await this.attestations.set(this.getAttestationKey(slotNumber, proposalId, address), attestation.toBuffer());
 
-        await this.proposalsForSlot.set(slotNumber.toString(), proposalId.toString());
+        await this.proposalsForSlot.set(slotNumber, proposalId.toString());
         await this.attestationsForProposal.set(
           this.getProposalKey(slotNumber, proposalId),
           this.getAttestationKey(slotNumber, proposalId, address),
         );
 
-        this.log.verbose(`Added attestation for slot ${slotNumber} from ${address}`);
+        this.log.verbose(`Added attestation for slot ${slotNumber} from ${address}`, {
+          signature: attestation.signature.toString(),
+          slotNumber,
+          address,
+          proposalId,
+        });
       }
     });
+  }
+
+  public async getAttestationsForSlot(slot: SlotNumber): Promise<BlockAttestation[]> {
+    const proposalIds = await toArray(this.proposalsForSlot.getValuesAsync(slot));
+    const attestations: BlockAttestation[] = [];
 
-    this.metrics.recordAddedObjects(attestations.length);
+    for (const proposalId of proposalIds) {
+      attestations.push(...(await this.getAttestationsForSlotAndProposal(slot, proposalId)));
+    }
+
+    return attestations;
   }
 
-  public async getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]> {
+  public async getAttestationsForSlotAndProposal(slot: SlotNumber, proposalId: string): Promise<BlockAttestation[]> {
     const attestationIds = await toArray(
       this.attestationsForProposal.getValuesAsync(this.getProposalKey(slot, proposalId)),
     );
@@ -87,21 +136,20 @@ export class KvAttestationPool implements AttestationPool {
     return attestations;
   }
 
-  public async deleteAttestationsOlderThan(oldestSlot: bigint): Promise<void> {
-    const olderThan = await toArray(this.proposalsForSlot.keysAsync({ end: new Fr(oldestSlot).toString() }));
+  public async deleteAttestationsOlderThan(oldestSlot: SlotNumber): Promise<void> {
+    const olderThan = await toArray(this.proposalsForSlot.keysAsync({ end: oldestSlot }));
     for (const oldSlot of olderThan) {
-      await this.deleteAttestationsForSlot(BigInt(oldSlot));
+      await this.deleteAttestationsForSlot(SlotNumber(oldSlot));
     }
   }
 
-  public async deleteAttestationsForSlot(slot: bigint): Promise<void> {
-    const slotFr = new Fr(slot);
+  public async deleteAttestationsForSlot(slot: SlotNumber): Promise<void> {
     let numberOfAttestations = 0;
     await this.store.transactionAsync(async () => {
-      const proposalIds = await toArray(this.proposalsForSlot.getValuesAsync(slotFr.toString()));
+      const proposalIds = await toArray(this.proposalsForSlot.getValuesAsync(slot));
       for (const proposalId of proposalIds) {
         const attestations = await toArray(
-          this.attestationsForProposal.getValuesAsync(this.getProposalKey(slotFr, proposalId)),
+          this.attestationsForProposal.getValuesAsync(this.getProposalKey(slot, proposalId)),
         );
 
         numberOfAttestations += attestations.length;
@@ -109,18 +157,20 @@ export class KvAttestationPool implements AttestationPool {
           await this.attestations.delete(attestation);
         }
 
-        await this.attestationsForProposal.delete(this.getProposalKey(slotFr, proposalId));
+        await this.proposals.delete(proposalId);
+        await this.attestationsForProposal.delete(this.getProposalKey(slot, proposalId));
       }
-    });
 
-    this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot}`);
-    this.metrics.recordRemovedObjects(numberOfAttestations);
+      // Delete from proposalsForSlot
+      await this.proposalsForSlot.delete(slot);
+
+      this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot}`);
+    });
   }
 
-  public async deleteAttestationsForSlotAndProposal(slot: bigint, proposalId: string): Promise<void> {
+  public async deleteAttestationsForSlotAndProposal(slot: SlotNumber, proposalId: string): Promise<void> {
     let numberOfAttestations = 0;
     await this.store.transactionAsync(async () => {
-      const slotString = new Fr(slot).toString();
       const attestations = await toArray(
         this.attestationsForProposal.getValuesAsync(this.getProposalKey(slot, proposalId)),
       );
@@ -130,22 +180,34 @@ export class KvAttestationPool implements AttestationPool {
         await this.attestations.delete(attestation);
       }
 
-      await this.proposalsForSlot.deleteValue(slotString, proposalId);
-      await this.attestationsForProposal.delete(this.getProposalKey(slotString, proposalId));
-    });
+      await this.proposals.delete(proposalId);
+      await this.proposalsForSlot.deleteValue(slot, proposalId);
+      await this.attestationsForProposal.delete(this.getProposalKey(slot, proposalId));
 
-    this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot} and proposal ${proposalId}`);
-    this.metrics.recordRemovedObjects(numberOfAttestations);
+      this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot} and proposal ${proposalId}`);
+    });
   }
 
   public async deleteAttestations(attestations: BlockAttestation[]): Promise<void> {
     await this.store.transactionAsync(async () => {
       for (const attestation of attestations) {
-        const slotNumber = attestation.payload.header.globalVariables.slotNumber;
+        const slotNumber = attestation.payload.header.slotNumber;
         const proposalId = attestation.archive;
-        const address = (await attestation.getSender()).toString();
+        const sender = attestation.getSender();
+
+        // Skip attestations with invalid signatures
+        if (!sender) {
+          this.log.warn(`Skipping deletion of attestation with invalid signature for slot ${slotNumber}`);
+          continue;
+        }
+
+        const address = sender.toString();
+        const key = this.getAttestationKey(slotNumber, proposalId, address);
+
+        if (await this.attestations.hasAsync(key)) {
+          await this.attestations.delete(key);
+        }
 
-        await this.attestations.delete(this.getAttestationKey(slotNumber, proposalId, address));
         await this.attestationsForProposal.deleteValue(
           this.getProposalKey(slotNumber, proposalId),
           this.getAttestationKey(slotNumber, proposalId, address),
@@ -154,6 +216,83 @@ export class KvAttestationPool implements AttestationPool {
         this.log.debug(`Deleted attestation for slot ${slotNumber} from ${address}`);
       }
     });
-    this.metrics.recordRemovedObjects(attestations.length);
+  }
+
+  public async hasAttestation(attestation: BlockAttestation): Promise<boolean> {
+    const slotNumber = attestation.payload.header.slotNumber;
+    const proposalId = attestation.archive;
+    const sender = attestation.getSender();
+
+    // Attestations with invalid signatures are never in the pool
+    if (!sender) {
+      return false;
+    }
+
+    const address = sender.toString();
+    const key = this.getAttestationKey(slotNumber, proposalId, address);
+
+    return await this.attestations.hasAsync(key);
+  }
+
+  public async getBlockProposal(id: string): Promise<BlockProposal | undefined> {
+    const buffer = await this.proposals.getAsync(id);
+    try {
+      if (buffer && buffer.length > 0) {
+        return BlockProposal.fromBuffer(buffer);
+      }
+    } catch {
+      return Promise.resolve(undefined);
+    }
+
+    return Promise.resolve(undefined);
+  }
+
+  public async hasBlockProposal(idOrProposal: string | BlockProposal): Promise<boolean> {
+    const id = typeof idOrProposal === 'string' ? idOrProposal : idOrProposal.payload.archive.toString();
+    return await this.proposals.hasAsync(id);
+  }
+
+  public async addBlockProposal(blockProposal: BlockProposal): Promise<void> {
+    await this.store.transactionAsync(async () => {
+      const slotKey = blockProposal.slotNumber;
+      const proposalId = blockProposal.archive.toString();
+
+      if (!(await this.canAddProposal(blockProposal))) {
+        throw new ProposalSlotCapExceededError(
+          `Maximum proposals per slot reached: slot=${slotKey} cap=${MAX_PROPOSALS_PER_SLOT} proposal=${proposalId}`,
+        );
+      }
+
+      await this.proposalsForSlot.set(slotKey, proposalId);
+      // Always update the stored proposal buffer so re-adds overwrite with latest data
+      await this.proposals.set(proposalId, blockProposal.toBuffer());
+    });
+  }
+
+  public async hasReachedProposalCap(slot: SlotNumber): Promise<boolean> {
+    const uniqueProposalCount = await this.proposalsForSlot.getValueCountAsync(slot);
+    return uniqueProposalCount >= MAX_PROPOSALS_PER_SLOT;
+  }
+
+  public async hasReachedAttestationCap(slot: SlotNumber, proposalId: string, committeeSize: number): Promise<boolean> {
+    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
+    return (await this.attestationsForProposal.getValueCountAsync(this.getProposalKey(slot, proposalId))) >= limit;
+  }
+
+  public async canAddProposal(block: BlockProposal): Promise<boolean> {
+    return (
+      (await this.proposals.hasAsync(block.archive.toString())) || !(await this.hasReachedProposalCap(block.slotNumber))
+    );
+  }
+
+  public async canAddAttestation(attestation: BlockAttestation, committeeSize: number): Promise<boolean> {
+    return (
+      (await this.hasAttestation(attestation)) ||
+      !(await this.hasReachedAttestationCap(
+        attestation.payload.header.slotNumber,
+        attestation.archive.toString(),
+        committeeSize,
+      ))
+    );
   }
 }

package/src/mem_pools/attestation_pool/memory_attestation_pool.ts

@@ -1,21 +1,50 @@
+import type { SlotNumber } from '@aztec/foundation/branded-types';
 import { createLogger } from '@aztec/foundation/log';
-import type { BlockAttestation } from '@aztec/stdlib/p2p';
+import type { BlockAttestation, BlockProposal } from '@aztec/stdlib/p2p';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
-import { PoolInstrumentation, PoolName } from '../instrumentation.js';
+import { PoolInstrumentation, PoolName, type PoolStatsCallback } from '../instrumentation.js';
 import type { AttestationPool } from './attestation_pool.js';
+import { ATTESTATION_CAP_BUFFER, MAX_PROPOSALS_PER_SLOT } from './kv_attestation_pool.js';
 
 export class InMemoryAttestationPool implements AttestationPool {
   private metrics: PoolInstrumentation<BlockAttestation>;
 
-  private attestations: Map</*slot=*/ bigint, Map</*proposalId*/ string, Map</*address=*/ string, BlockAttestation>>>;
+  // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
+  private attestations: Map<
+    /*slot=*/ SlotNumber,
+    Map</*proposalId*/ string, Map</*address=*/ string, BlockAttestation>>
+  >;
+  private proposals: Map<string, BlockProposal>;
 
-  constructor(telemetry: TelemetryClient = getTelemetryClient(), private log = createLogger('p2p:attestation_pool')) {
+  constructor(
+    telemetry: TelemetryClient = getTelemetryClient(),
+    private log = createLogger('p2p:attestation_pool'),
+  ) {
     this.attestations = new Map();
-    this.metrics = new PoolInstrumentation(telemetry, PoolName.ATTESTATION_POOL);
+    this.proposals = new Map();
+    this.metrics = new PoolInstrumentation(telemetry, PoolName.ATTESTATION_POOL, this.poolStats);
   }
 
-  public getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]> {
+  private poolStats: PoolStatsCallback = () => {
+    return Promise.resolve({
+      itemCount: this.attestations.size,
+    });
+  };
+
+  public isEmpty(): Promise<boolean> {
+    return Promise.resolve(this.attestations.size === 0);
+  }
+
+  public getAttestationsForSlot(slot: SlotNumber): Promise<BlockAttestation[]> {
+    return Promise.resolve(
+      Array.from(this.attestations.get(slot)?.values() ?? []).flatMap(proposalAttestationMap =>
+        Array.from(proposalAttestationMap.values()),
+      ),
+    );
+  }
+
+  public getAttestationsForSlotAndProposal(slot: SlotNumber, proposalId: string): Promise<BlockAttestation[]> {
     const slotAttestationMap = this.attestations.get(slot);
     if (slotAttestationMap) {
       const proposalAttestationMap = slotAttestationMap.get(proposalId);
@@ -26,27 +55,40 @@ export class InMemoryAttestationPool implements AttestationPool {
     return Promise.resolve([]);
   }
 
-  public async addAttestations(attestations: BlockAttestation[]): Promise<void> {
+  public addAttestations(attestations: BlockAttestation[]): Promise<void> {
     for (const attestation of attestations) {
       // Perf: order and group by slot before insertion
-      const slotNumber = attestation.payload.header.globalVariables.slotNumber;
+      const slotNumber = attestation.payload.header.slotNumber;
 
       const proposalId = attestation.archive.toString();
-      const address = await attestation.getSender();
+      const sender = attestation.getSender();
 
-      const slotAttestationMap = getSlotOrDefault(this.attestations, slotNumber.toBigInt());
+      // Skip attestations with invalid signatures
+      if (!sender) {
+        this.log.warn(`Skipping attestation with invalid signature for slot ${slotNumber}`, {
+          signature: attestation.signature.toString(),
+          slotNumber,
+          proposalId,
+        });
+        continue;
+      }
+
+      const slotAttestationMap = getSlotOrDefault(this.attestations, slotNumber);
       const proposalAttestationMap = getProposalOrDefault(slotAttestationMap, proposalId);
-      proposalAttestationMap.set(address.toString(), attestation);
+      proposalAttestationMap.set(sender.toString(), attestation);
 
-      this.log.verbose(`Added attestation for slot ${slotNumber} from ${address}`);
+      this.log.verbose(`Added attestation for slot ${slotNumber} from ${sender}`, {
+        signature: attestation.signature.toString(),
+        slotNumber,
+        address: sender,
+        proposalId,
+      });
     }
 
-    // TODO: set these to pending or something ????
-    this.metrics.recordAddedObjects(attestations.length);
     return Promise.resolve();
   }
 
-  #getNumberOfAttestationsInSlot(slot: bigint): number {
+  #getNumberOfAttestationsInSlot(slot: SlotNumber): number {
     let total = 0;
     const slotAttestationMap = getSlotOrDefault(this.attestations, slot);
 
@@ -58,7 +100,7 @@ export class InMemoryAttestationPool implements AttestationPool {
     return total;
   }
 
-  public async deleteAttestationsOlderThan(oldestSlot: bigint): Promise<void> {
+  public async deleteAttestationsOlderThan(oldestSlot: SlotNumber): Promise<void> {
     const olderThan = [];
 
     // Entries are iterated in insertion order, so we can break as soon as we find a slot that is older than the oldestSlot.
@@ -79,18 +121,25 @@ export class InMemoryAttestationPool implements AttestationPool {
     return Promise.resolve();
   }
 
-  public deleteAttestationsForSlot(slot: bigint): Promise<void> {
+  public deleteAttestationsForSlot(slot: SlotNumber): Promise<void> {
     // We count the number of attestations we are removing
     const numberOfAttestations = this.#getNumberOfAttestationsInSlot(slot);
+    const proposalIdsToDelete = this.attestations.get(slot)?.keys();
+    let proposalIdsToDeleteCount = 0;
+    proposalIdsToDelete?.forEach(proposalId => {
+      this.proposals.delete(proposalId);
+      proposalIdsToDeleteCount++;
+    });
 
     this.attestations.delete(slot);
-    this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot}`);
+    this.log.verbose(
+      `Removed ${numberOfAttestations} attestations and ${proposalIdsToDeleteCount} proposals for slot ${slot}`,
+    );
 
-    this.metrics.recordRemovedObjects(numberOfAttestations);
     return Promise.resolve();
   }
 
-  public deleteAttestationsForSlotAndProposal(slot: bigint, proposalId: string): Promise<void> {
+  public deleteAttestationsForSlotAndProposal(slot: SlotNumber, proposalId: string): Promise<void> {
     const slotAttestationMap = getSlotOrDefault(this.attestations, slot);
     if (slotAttestationMap) {
       if (slotAttestationMap.has(proposalId)) {
@@ -99,29 +148,105 @@ export class InMemoryAttestationPool implements AttestationPool {
         slotAttestationMap.delete(proposalId);
 
         this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot} and proposal ${proposalId}`);
-        this.metrics.recordRemovedObjects(numberOfAttestations);
       }
     }
+
+    this.proposals.delete(proposalId);
     return Promise.resolve();
   }
 
-  public async deleteAttestations(attestations: BlockAttestation[]): Promise<void> {
+  public deleteAttestations(attestations: BlockAttestation[]): Promise<void> {
     for (const attestation of attestations) {
-      const slotNumber = attestation.payload.header.globalVariables.slotNumber;
-      const slotAttestationMap = this.attestations.get(slotNumber.toBigInt());
+      const slotNumber = attestation.payload.header.slotNumber;
+      const slotAttestationMap = this.attestations.get(slotNumber);
       if (slotAttestationMap) {
         const proposalId = attestation.archive.toString();
         const proposalAttestationMap = getProposalOrDefault(slotAttestationMap, proposalId);
         if (proposalAttestationMap) {
-          const address = await attestation.getSender();
-          proposalAttestationMap.delete(address.toString());
-          this.log.debug(`Deleted attestation for slot ${slotNumber} from ${address}`);
+          const sender = attestation.getSender();
+
+          // Skip attestations with invalid signatures
+          if (!sender) {
+            this.log.warn(`Skipping deletion of attestation with invalid signature for slot ${slotNumber}`);
+            continue;
+          }
+
+          proposalAttestationMap.delete(sender.toString());
+          this.log.debug(`Deleted attestation for slot ${slotNumber} from ${sender}`);
         }
       }
     }
-    this.metrics.recordRemovedObjects(attestations.length);
     return Promise.resolve();
   }
+
+  public hasAttestation(attestation: BlockAttestation): Promise<boolean> {
+    const slotNumber = attestation.payload.header.slotNumber;
+    const proposalId = attestation.archive.toString();
+    const sender = attestation.getSender();
+
+    // Attestations with invalid signatures are never in the pool
+    if (!sender) {
+      return Promise.resolve(false);
+    }
+
+    const slotAttestationMap = this.attestations.get(slotNumber);
+    if (!slotAttestationMap) {
+      return Promise.resolve(false);
+    }
+
+    const proposalAttestationMap = slotAttestationMap.get(proposalId);
+    if (!proposalAttestationMap) {
+      return Promise.resolve(false);
+    }
+
+    return Promise.resolve(proposalAttestationMap.has(sender.toString()));
+  }
+
+  public addBlockProposal(blockProposal: BlockProposal): Promise<void> {
+    // We initialize slot-proposal mapping if it does not exist
+    // This is important to ensure we can delete this proposal if there were not attestations for it
+    const slotProposalMapping = getSlotOrDefault(this.attestations, blockProposal.slotNumber);
+    slotProposalMapping.set(blockProposal.payload.archive.toString(), new Map<string, BlockAttestation>());
+
+    this.proposals.set(blockProposal.payload.archive.toString(), blockProposal);
+    return Promise.resolve();
+  }
+
+  public getBlockProposal(id: string): Promise<BlockProposal | undefined> {
+    return Promise.resolve(this.proposals.get(id));
+  }
+
+  public hasBlockProposal(idOrProposal: string | BlockProposal): Promise<boolean> {
+    const id = typeof idOrProposal === 'string' ? idOrProposal : idOrProposal.payload.archive.toString();
+    return Promise.resolve(this.proposals.has(id));
+  }
+
+  public hasReachedProposalCap(slot: SlotNumber): Promise<boolean> {
+    const slotAttestationMap = this.attestations.get(slot);
+    const proposalCount = slotAttestationMap?.size ?? 0;
+    return Promise.resolve(proposalCount >= MAX_PROPOSALS_PER_SLOT);
+  }
+
+  public hasReachedAttestationCap(slot: SlotNumber, proposalId: string, committeeSize: number): Promise<boolean> {
+    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
+    const count = this.attestations.get(slot)?.get(proposalId)?.size ?? 0;
+    return Promise.resolve(limit <= 0 || count >= limit);
+  }
+
+  public async canAddProposal(block: BlockProposal): Promise<boolean> {
+    return this.proposals.has(block.archive.toString()) || !(await this.hasReachedProposalCap(block.slotNumber));
+  }
+
+  public async canAddAttestation(attestation: BlockAttestation, committeeSize: number): Promise<boolean> {
+    const sender = attestation.getSender();
+    const slot = attestation.payload.header.slotNumber;
+    const pid = attestation.archive.toString();
+    return (
+      !!sender &&
+      ((this.attestations.get(slot)?.get(pid)?.has(sender.toString()) ?? false) ||
+        !(await this.hasReachedAttestationCap(slot, pid, committeeSize)))
+    );
+  }
 }
 
 /**
@@ -133,8 +258,9 @@ export class InMemoryAttestationPool implements AttestationPool {
  * @returns The slot mapping
  */
 function getSlotOrDefault(
-  map: Map<bigint, Map<string, Map<string, BlockAttestation>>>,
-  slot: bigint,
+  // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
+  map: Map<SlotNumber, Map<string, Map<string, BlockAttestation>>>,
+  slot: SlotNumber,
 ): Map<string, Map<string, BlockAttestation>> {
   if (!map.has(slot)) {
     map.set(slot, new Map<string, Map<string, BlockAttestation>>());

package/src/mem_pools/attestation_pool/mocks.ts

@@ -1,13 +1,12 @@
-import type { Secp256k1Signer } from '@aztec/foundation/crypto';
-import { Fr } from '@aztec/foundation/fields';
+import type { Secp256k1Signer } from '@aztec/foundation/crypto/secp256k1-signer';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import {
   BlockAttestation,
   ConsensusPayload,
   SignatureDomainSeparator,
   getHashedSignaturePayloadEthSignedMessage,
 } from '@aztec/stdlib/p2p';
-import { makeHeader } from '@aztec/stdlib/testing';
-import { TxHash } from '@aztec/stdlib/tx';
+import { makeL2BlockHeader } from '@aztec/stdlib/testing';
 
 import { type LocalAccount, generatePrivateKey, privateKeyToAccount } from 'viem/accounts';
 
@@ -27,18 +26,20 @@ export const generateAccount = (): LocalAccount => {
  * @param slot The slot number the attestation is for
  * @returns A Block Attestation
  */
-export const mockAttestation = async (
+export const mockAttestation = (
   signer: Secp256k1Signer,
   slot: number = 0,
   archive: Fr = Fr.random(),
-  txs: TxHash[] = [0, 1, 2, 3, 4, 5].map(() => TxHash.random()),
-): Promise<BlockAttestation> => {
+): BlockAttestation => {
   // Use arbitrary numbers for all other than slot
-  const header = makeHeader(1, 2, slot);
-  const payload = new ConsensusPayload(header, archive, txs);
+  const header = makeL2BlockHeader(1, 2, slot);
+  const payload = new ConsensusPayload(header.toCheckpointHeader(), archive);
 
-  const hash = await getHashedSignaturePayloadEthSignedMessage(payload, SignatureDomainSeparator.blockAttestation);
-  const signature = signer.sign(hash);
+  const attestationHash = getHashedSignaturePayloadEthSignedMessage(payload, SignatureDomainSeparator.blockAttestation);
+  const attestationSignature = signer.sign(attestationHash);
 
-  return new BlockAttestation(payload, signature);
+  const proposalHash = getHashedSignaturePayloadEthSignedMessage(payload, SignatureDomainSeparator.blockProposal);
+  const proposerSignature = signer.sign(proposalHash);
+
+  return new BlockAttestation(payload, attestationSignature, proposerSignature);
 };