@aztec/p2p 0.0.1-commit.b655e406 → 0.0.1-commit.d1f2d6c

package/src/mem_pools/attestation_pool/kv_attestation_pool.ts

@@ -1,45 +1,64 @@
-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, BlockProposal } from '@aztec/stdlib/p2p';
+import {
+  BlockProposal,
+  CheckpointAttestation,
+  CheckpointProposal,
+  type CheckpointProposalCore,
+} from '@aztec/stdlib/p2p';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
+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 metrics: PoolInstrumentation<CheckpointAttestation>;
 
-  private attestations: AztecAsyncMap<string, Buffer>;
   private proposals: AztecAsyncMap<
     /*  proposal.payload.archive  */ string,
     /* buffer representation of proposal */ Buffer
   >;
-  private proposalsForSlot: AztecAsyncMultiMap<string, string>;
-  private attestationsForProposal: AztecAsyncMultiMap<string, string>;
+
+  // Checkpoint attestation storage
+  private checkpointAttestations: AztecAsyncMap<string, Buffer>;
+  private checkpointProposals: AztecAsyncMap<string, Buffer>;
+  private checkpointProposalsForSlot: AztecAsyncMultiMap<number, string>;
+  private checkpointAttestationsForProposal: AztecAsyncMultiMap<string, string>;
 
   constructor(
     private store: AztecAsyncKVStore,
     telemetry: TelemetryClient = getTelemetryClient(),
     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');
+
+    // Initialize checkpoint attestation storage
+    this.checkpointAttestations = store.openMap('checkpoint_attestations');
+    this.checkpointProposals = store.openMap('checkpoint_proposals');
+    this.checkpointProposalsForSlot = store.openMultiMap('checkpoint_proposals_for_slot');
+    this.checkpointAttestationsForProposal = store.openMultiMap('checkpoint_attestations_for_proposal');
 
     this.metrics = new PoolInstrumentation(telemetry, PoolName.ATTESTATION_POOL, this.poolStats);
   }
 
   private poolStats: PoolStatsCallback = async () => {
     return {
-      itemCount: await this.attestations.sizeAsync(),
+      itemCount: await this.checkpointAttestations.sizeAsync(),
     };
   };
 
   public async isEmpty(): Promise<boolean> {
-    for await (const _ of this.attestations.entriesAsync()) {
+    for await (const _ of this.checkpointAttestations.entriesAsync()) {
+      return false;
+    }
+    for await (const _ of this.proposals.entriesAsync()) {
       return false;
     }
     return true;
@@ -61,7 +80,81 @@ export class KvAttestationPool implements AttestationPool {
     return `${this.getProposalKey(slot, proposalId)}-${address}`;
   }
 
-  public async addAttestations(attestations: BlockAttestation[]): Promise<void> {
+  public async addBlockProposal(blockProposal: BlockProposal): Promise<void> {
+    await this.store.transactionAsync(async () => {
+      const proposalId = blockProposal.archive.toString();
+      // Strip signedTxs before storing to avoid persisting full tx data
+      await this.proposals.set(proposalId, blockProposal.withoutSignedTxs().toBuffer());
+    });
+  }
+
+  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.archive.toString();
+    return await this.proposals.hasAsync(id);
+  }
+
+  public async addCheckpointProposal(proposal: CheckpointProposal): Promise<void> {
+    if (!(await this.canAddCheckpointProposal(proposal))) {
+      throw new ProposalSlotCapExceededError(
+        `Maximum checkpoint proposals per slot reached: slot=${proposal.slotNumber} cap=${MAX_PROPOSALS_PER_SLOT} proposal=${proposal.archive.toString()}`,
+      );
+    }
+
+    // Extract and validate the block proposal if present
+    const blockProposal = proposal.getBlockProposal();
+    if (blockProposal && !(await this.canAddProposal(blockProposal))) {
+      throw new ProposalSlotCapExceededError(
+        `Maximum block proposals per slot reached when extracting from checkpoint: slot=${proposal.slotNumber} proposal=${blockProposal.archive.toString()}`,
+      );
+    }
+
+    await this.store.transactionAsync(async () => {
+      const slotKey = proposal.slotNumber;
+      const proposalId = proposal.archive.toString();
+
+      await this.checkpointProposalsForSlot.set(slotKey, proposalId);
+      // Store the checkpoint proposal as core (without lastBlock) to avoid duplication
+      await this.checkpointProposals.set(proposalId, proposal.toCore().toBuffer());
+
+      // Store the extracted block proposal separately
+      if (blockProposal) {
+        await this.proposals.set(blockProposal.archive.toString(), blockProposal.withoutSignedTxs().toBuffer());
+      }
+    });
+  }
+
+  public async getCheckpointProposal(id: string): Promise<CheckpointProposalCore | undefined> {
+    const buffer = await this.checkpointProposals.getAsync(id);
+    try {
+      if (buffer && buffer.length > 0) {
+        return CheckpointProposal.fromBuffer(buffer);
+      }
+    } catch {
+      return Promise.resolve(undefined);
+    }
+
+    return Promise.resolve(undefined);
+  }
+
+  public async hasCheckpointProposal(idOrProposal: string | CheckpointProposal): Promise<boolean> {
+    const id = typeof idOrProposal === 'string' ? idOrProposal : idOrProposal.archive.toString();
+    return await this.checkpointProposals.hasAsync(id);
+  }
+
+  public async addCheckpointAttestations(attestations: CheckpointAttestation[]): Promise<void> {
     await this.store.transactionAsync(async () => {
       for (const attestation of attestations) {
         const slotNumber = attestation.payload.header.slotNumber;
@@ -70,7 +163,7 @@ export class KvAttestationPool implements AttestationPool {
 
         // Skip attestations with invalid signatures
         if (!sender) {
-          this.log.warn(`Skipping attestation with invalid signature for slot ${slotNumber.toBigInt()}`, {
+          this.log.warn(`Skipping checkpoint attestation with invalid signature for slot ${slotNumber}`, {
             signature: attestation.signature.toString(),
             slotNumber,
             proposalId,
@@ -80,15 +173,18 @@ export class KvAttestationPool implements AttestationPool {
 
         const address = sender.toString();
 
-        await this.attestations.set(this.getAttestationKey(slotNumber, proposalId, address), attestation.toBuffer());
+        await this.checkpointAttestations.set(
+          this.getAttestationKey(slotNumber, proposalId, address),
+          attestation.toBuffer(),
+        );
 
-        await this.proposalsForSlot.set(slotNumber.toString(), proposalId.toString());
-        await this.attestationsForProposal.set(
+        await this.checkpointProposalsForSlot.set(slotNumber, proposalId.toString());
+        await this.checkpointAttestationsForProposal.set(
           this.getProposalKey(slotNumber, proposalId),
           this.getAttestationKey(slotNumber, proposalId, address),
         );
 
-        this.log.verbose(`Added attestation for slot ${slotNumber.toBigInt()} from ${address}`, {
+        this.log.verbose(`Added checkpoint attestation for slot ${slotNumber} from ${address}`, {
           signature: attestation.signature.toString(),
           slotNumber,
           address,
@@ -98,122 +194,72 @@ export class KvAttestationPool implements AttestationPool {
     });
   }
 
-  public async getAttestationsForSlot(slot: bigint): Promise<BlockAttestation[]> {
-    const slotFr = new Fr(slot);
-    const proposalIds = await toArray(this.proposalsForSlot.getValuesAsync(slotFr.toString()));
-    const attestations: BlockAttestation[] = [];
+  public async getCheckpointAttestationsForSlot(slot: SlotNumber): Promise<CheckpointAttestation[]> {
+    const proposalIds = await toArray(this.checkpointProposalsForSlot.getValuesAsync(slot));
+    const attestations: CheckpointAttestation[] = [];
 
     for (const proposalId of proposalIds) {
-      attestations.push(...(await this.getAttestationsForSlotAndProposal(slot, proposalId)));
+      attestations.push(...(await this.getCheckpointAttestationsForSlotAndProposal(slot, proposalId)));
     }
 
     return attestations;
   }
 
-  public async getAttestationsForSlotAndProposal(slot: bigint, proposalId: string): Promise<BlockAttestation[]> {
+  public async getCheckpointAttestationsForSlotAndProposal(
+    slot: SlotNumber,
+    proposalId: string,
+  ): Promise<CheckpointAttestation[]> {
     const attestationIds = await toArray(
-      this.attestationsForProposal.getValuesAsync(this.getProposalKey(slot, proposalId)),
+      this.checkpointAttestationsForProposal.getValuesAsync(this.getProposalKey(slot, proposalId)),
     );
-    const attestations: BlockAttestation[] = [];
+    const attestations: CheckpointAttestation[] = [];
 
-    // alternatively iterate this.attestaions starting from slot-proposal-EthAddress.zero
     for (const id of attestationIds) {
-      const buf = await this.attestations.getAsync(id);
+      const buf = await this.checkpointAttestations.getAsync(id);
 
       if (!buf) {
-        // this should not happen unless we lost writes
-        throw new Error('Attestation not found ' + id);
+        throw new Error('Checkpoint attestation not found ' + id);
       }
 
-      const attestation = BlockAttestation.fromBuffer(buf);
+      const attestation = CheckpointAttestation.fromBuffer(buf);
       attestations.push(attestation);
     }
 
     return attestations;
   }
 
-  public async deleteAttestationsOlderThan(oldestSlot: bigint): Promise<void> {
-    const olderThan = await toArray(this.proposalsForSlot.keysAsync({ end: new Fr(oldestSlot).toString() }));
+  public async deleteCheckpointAttestationsOlderThan(oldestSlot: SlotNumber): Promise<void> {
+    const olderThan = await toArray(this.checkpointProposalsForSlot.keysAsync({ end: oldestSlot }));
     for (const oldSlot of olderThan) {
-      await this.deleteAttestationsForSlot(BigInt(oldSlot));
+      await this.deleteCheckpointAttestationsForSlot(SlotNumber(oldSlot));
     }
   }
 
-  public async deleteAttestationsForSlot(slot: bigint): Promise<void> {
-    const slotFr = new Fr(slot);
+  private async deleteCheckpointAttestationsForSlot(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.checkpointProposalsForSlot.getValuesAsync(slot));
       for (const proposalId of proposalIds) {
         const attestations = await toArray(
-          this.attestationsForProposal.getValuesAsync(this.getProposalKey(slotFr, proposalId)),
+          this.checkpointAttestationsForProposal.getValuesAsync(this.getProposalKey(slot, proposalId)),
         );
 
         numberOfAttestations += attestations.length;
         for (const attestation of attestations) {
-          await this.attestations.delete(attestation);
+          await this.checkpointAttestations.delete(attestation);
         }
 
-        await this.proposals.delete(proposalId);
-        await this.attestationsForProposal.delete(this.getProposalKey(slotFr, proposalId));
+        await this.checkpointProposals.delete(proposalId);
+        await this.checkpointAttestationsForProposal.delete(this.getProposalKey(slot, proposalId));
       }
 
-      this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot}`);
-    });
-  }
-
-  public async deleteAttestationsForSlotAndProposal(slot: bigint, 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)),
-      );
-
-      numberOfAttestations += attestations.length;
-      for (const attestation of attestations) {
-        await this.attestations.delete(attestation);
-      }
-
-      await this.proposals.delete(proposalId);
-      await this.proposalsForSlot.deleteValue(slotString, proposalId);
-      await this.attestationsForProposal.delete(this.getProposalKey(slotString, proposalId));
-
-      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.slotNumber;
-        const proposalId = attestation.archive;
-        const sender = attestation.getSender();
-
-        // Skip attestations with invalid signatures
-        if (!sender) {
-          this.log.warn(`Skipping deletion of attestation with invalid signature for slot ${slotNumber.toBigInt()}`);
-          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.checkpointProposalsForSlot.delete(slot);
 
-        await this.attestationsForProposal.deleteValue(
-          this.getProposalKey(slotNumber, proposalId),
-          this.getAttestationKey(slotNumber, proposalId, address),
-        );
-
-        this.log.debug(`Deleted attestation for slot ${slotNumber} from ${address}`);
-      }
+      this.log.verbose(`Removed ${numberOfAttestations} checkpoint attestations for slot ${slot}`);
     });
   }
 
-  public async hasAttestation(attestation: BlockAttestation): Promise<boolean> {
+  public async hasCheckpointAttestation(attestation: CheckpointAttestation): Promise<boolean> {
     const slotNumber = attestation.payload.header.slotNumber;
     const proposalId = attestation.archive;
     const sender = attestation.getSender();
@@ -226,31 +272,49 @@ export class KvAttestationPool implements AttestationPool {
     const address = sender.toString();
     const key = this.getAttestationKey(slotNumber, proposalId, address);
 
-    return await this.attestations.hasAsync(key);
+    return await this.checkpointAttestations.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);
-    }
+  public canAddProposal(_block: BlockProposal): Promise<boolean> {
+    // TODO(palla/mbps): implement proposal cap logic
+    return Promise.resolve(true);
+  }
 
-    return Promise.resolve(undefined);
+  public async canAddCheckpointProposal(proposal: CheckpointProposal): Promise<boolean> {
+    // TODO(palla/mbps): Adjust checkpoint proposal limit to 1. Also connect to slashing condition in the caller.
+    return (
+      (await this.checkpointProposals.hasAsync(proposal.archive.toString())) ||
+      !(await this.hasReachedCheckpointProposalCap(proposal.slotNumber))
+    );
   }
 
-  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 canAddCheckpointAttestation(
+    attestation: CheckpointAttestation,
+    committeeSize: number,
+  ): Promise<boolean> {
+    return (
+      (await this.hasCheckpointAttestation(attestation)) ||
+      !(await this.hasReachedCheckpointAttestationCap(
+        attestation.payload.header.slotNumber,
+        attestation.archive.toString(),
+        committeeSize,
+      ))
+    );
   }
 
-  public async addBlockProposal(blockProposal: BlockProposal): Promise<void> {
-    await this.store.transactionAsync(async () => {
-      await this.proposalsForSlot.set(blockProposal.slotNumber.toString(), blockProposal.archive.toString());
-      await this.proposals.set(blockProposal.payload.archive.toString(), blockProposal.toBuffer());
-    });
+  public async hasReachedCheckpointProposalCap(slot: SlotNumber): Promise<boolean> {
+    const uniqueProposalCount = await this.checkpointProposalsForSlot.getValueCountAsync(slot);
+    return uniqueProposalCount >= MAX_PROPOSALS_PER_SLOT;
+  }
+
+  public async hasReachedCheckpointAttestationCap(
+    slot: SlotNumber,
+    proposalId: string,
+    committeeSize: number,
+  ): Promise<boolean> {
+    const limit = committeeSize + ATTESTATION_CAP_BUFFER;
+    return (
+      (await this.checkpointAttestationsForProposal.getValueCountAsync(this.getProposalKey(slot, proposalId))) >= limit
+    );
   }
 }