From b6032c53aac7886f5654017c3105286ba8d716c8 Mon Sep 17 00:00:00 2001
From: Lion - dapplion <35266934+dapplion@users.noreply.github.com>
Date: Tue, 6 Dec 2022 03:48:58 +0700
Subject: [PATCH] Use state dialed to attestation target epoch (#4849)

---
 .../src/chain/validation/aggregateAndProof.ts | 27 +++++++++++++------
 .../src/chain/validation/attestation.ts       | 23 ++++++++++------
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/packages/beacon-node/src/chain/validation/aggregateAndProof.ts b/packages/beacon-node/src/chain/validation/aggregateAndProof.ts
index 451a2e3e5dc..71c7fbedd7d 100644
--- a/packages/beacon-node/src/chain/validation/aggregateAndProof.ts
+++ b/packages/beacon-node/src/chain/validation/aggregateAndProof.ts
@@ -76,14 +76,25 @@ export async function validateGossipAggregateAndProof(
   // -- i.e. get_ancestor(store, aggregate.data.beacon_block_root, compute_start_slot_at_epoch(store.finalized_checkpoint.epoch)) == store.finalized_checkpoint.root
   // > Altready check in `chain.forkChoice.hasBlock(attestation.data.beaconBlockRoot)`
 
-  const attHeadState = await chain.regen
-    .getState(attHeadBlock.stateRoot, RegenCaller.validateGossipAggregateAndProof)
-    .catch((e: Error) => {
-      throw new AttestationError(GossipAction.REJECT, {
-        code: AttestationErrorCode.MISSING_ATTESTATION_HEAD_STATE,
-        error: e as Error,
-      });
-    });
+  // TODO: Must be a state in the same chain as attHeadBlock, but dialed to target.epoch
+  const attHeadState =
+    computeEpochAtSlot(attHeadBlock.slot) < attEpoch
+      ? await chain.regen
+          .getCheckpointState(attTarget, RegenCaller.validateGossipAggregateAndProof)
+          .catch((e: Error) => {
+            throw new AttestationError(GossipAction.REJECT, {
+              code: AttestationErrorCode.MISSING_ATTESTATION_HEAD_STATE,
+              error: e as Error,
+            });
+          })
+      : await chain.regen
+          .getState(attHeadBlock.stateRoot, RegenCaller.validateGossipAggregateAndProof)
+          .catch((e: Error) => {
+            throw new AttestationError(GossipAction.REJECT, {
+              code: AttestationErrorCode.MISSING_ATTESTATION_HEAD_STATE,
+              error: e as Error,
+            });
+          });
 
   const committeeIndices: number[] = getCommitteeIndices(attHeadState, attSlot, attIndex);
 
diff --git a/packages/beacon-node/src/chain/validation/attestation.ts b/packages/beacon-node/src/chain/validation/attestation.ts
index 3e6b057346f..416417c76f9 100644
--- a/packages/beacon-node/src/chain/validation/attestation.ts
+++ b/packages/beacon-node/src/chain/validation/attestation.ts
@@ -76,14 +76,21 @@ export async function validateGossipAttestation(
   //  --i.e. get_ancestor(store, attestation.data.beacon_block_root, compute_start_slot_at_epoch(attestation.data.target.epoch)) == attestation.data.target.root
   // > Altready check in `verifyHeadBlockAndTargetRoot()`
 
-  const attHeadState = await chain.regen
-    .getState(attHeadBlock.stateRoot, RegenCaller.validateGossipAttestation)
-    .catch((e: Error) => {
-      throw new AttestationError(GossipAction.REJECT, {
-        code: AttestationErrorCode.MISSING_ATTESTATION_HEAD_STATE,
-        error: e as Error,
-      });
-    });
+  // TODO: Must be a state in the same chain as attHeadBlock, but dialed to target.epoch
+  const attHeadState =
+    computeEpochAtSlot(attHeadBlock.slot) < attEpoch
+      ? await chain.regen.getCheckpointState(attTarget, RegenCaller.validateGossipAttestation).catch((e: Error) => {
+          throw new AttestationError(GossipAction.REJECT, {
+            code: AttestationErrorCode.MISSING_ATTESTATION_HEAD_STATE,
+            error: e as Error,
+          });
+        })
+      : await chain.regen.getState(attHeadBlock.stateRoot, RegenCaller.validateGossipAttestation).catch((e: Error) => {
+          throw new AttestationError(GossipAction.REJECT, {
+            code: AttestationErrorCode.MISSING_ATTESTATION_HEAD_STATE,
+            error: e as Error,
+          });
+        });
 
   // [REJECT] The committee index is within the expected range
   // -- i.e. data.index < get_committee_count_per_slot(state, data.target.epoch)