From 4f42f63e4e6cb8dbc0c7237dbfe0e4fee8b660c1 Mon Sep 17 00:00:00 2001
From: Danny Ryan <dannyjryan@gmail.com>
Date: Tue, 5 Nov 2019 11:35:58 -0700
Subject: [PATCH] only allow attestatiosn to be considered from current and
 previous epoch

---
 specs/core/0_fork-choice.md                   |  3 ++
 .../test/fork_choice/test_on_attestation.py   | 52 +++++++++++++++++--
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/specs/core/0_fork-choice.md b/specs/core/0_fork-choice.md
index e08ea09b6..09c2fd31c 100644
--- a/specs/core/0_fork-choice.md
+++ b/specs/core/0_fork-choice.md
@@ -224,6 +224,9 @@ def on_block(store: Store, block: BeaconBlock) -> None:
 def on_attestation(store: Store, attestation: Attestation) -> None:
     target = attestation.data.target
 
+    # Attestations must be from the current or previous epoch
+    current_epoch = compute_epoch_at_slot(get_current_slot(store))
+    assert target.epoch in [current_epoch, current_epoch - 1 if current_epoch > GENESIS_EPOCH else GENESIS_EPOCH]
     # Cannot calculate the current shuffling if have not seen the target
     assert target.root in store.blocks
 
diff --git a/test_libs/pyspec/eth2spec/test/fork_choice/test_on_attestation.py b/test_libs/pyspec/eth2spec/test/fork_choice/test_on_attestation.py
index 70375ef27..3a513cf2f 100644
--- a/test_libs/pyspec/eth2spec/test/fork_choice/test_on_attestation.py
+++ b/test_libs/pyspec/eth2spec/test/fork_choice/test_on_attestation.py
@@ -29,10 +29,9 @@ def run_on_attestation(spec, state, store, attestation, valid=True):
 
 @with_all_phases
 @spec_state_test
-def test_on_attestation(spec, state):
+def test_on_attestation_current_epoch(spec, state):
     store = spec.get_genesis_store(state)
-    time = 100
-    spec.on_tick(store, time)
+    spec.on_tick(store, store.time + spec.SECONDS_PER_SLOT * 2)
 
     block = build_empty_block_for_next_slot(spec, state)
     state_transition_and_sign_block(spec, state, block)
@@ -41,9 +40,53 @@ def test_on_attestation(spec, state):
     spec.on_block(store, block)
 
     attestation = get_valid_attestation(spec, state, slot=block.slot)
+    assert attestation.data.target.epoch == spec.GENESIS_EPOCH
+    assert spec.compute_epoch_at_slot(spec.get_current_slot(store)) == spec.GENESIS_EPOCH
+
     run_on_attestation(spec, state, store, attestation)
 
 
+@with_all_phases
+@spec_state_test
+def test_on_attestation_previous_epoch(spec, state):
+    store = spec.get_genesis_store(state)
+    spec.on_tick(store, store.time + spec.SECONDS_PER_SLOT * spec.SLOTS_PER_EPOCH)
+
+    block = build_empty_block_for_next_slot(spec, state)
+    state_transition_and_sign_block(spec, state, block)
+
+    # store block in store
+    spec.on_block(store, block)
+
+    attestation = get_valid_attestation(spec, state, slot=block.slot)
+    assert attestation.data.target.epoch == spec.GENESIS_EPOCH
+    assert spec.compute_epoch_at_slot(spec.get_current_slot(store)) == spec.GENESIS_EPOCH + 1
+
+    run_on_attestation(spec, state, store, attestation)
+
+
+@with_all_phases
+@spec_state_test
+def test_on_attestation_past_epoch(spec, state):
+    store = spec.get_genesis_store(state)
+
+    # move time forward 2 epochs
+    time = store.time + 2 * spec.SECONDS_PER_SLOT * spec.SLOTS_PER_EPOCH
+    spec.on_tick(store, time)
+
+    # create and store block from 3 epochs ago
+    block = build_empty_block_for_next_slot(spec, state)
+    state_transition_and_sign_block(spec, state, block)
+    spec.on_block(store, block)
+
+    # create attestation for past block
+    attestation = get_valid_attestation(spec, state, slot=state.slot)
+    assert attestation.data.target.epoch == spec.GENESIS_EPOCH
+    assert spec.compute_epoch_at_slot(spec.get_current_slot(store)) == spec.GENESIS_EPOCH + 2
+
+    run_on_attestation(spec, state, store, attestation, False)
+
+
 @with_all_phases
 @spec_state_test
 def test_on_attestation_target_not_in_store(spec, state):
@@ -77,8 +120,7 @@ def test_on_attestation_future_epoch(spec, state):
     spec.on_block(store, block)
 
     # move state forward but not store
-    attestation_slot = block.slot + spec.SLOTS_PER_EPOCH
-    state.slot = attestation_slot
+    state.slot = block.slot + spec.SLOTS_PER_EPOCH
 
     attestation = get_valid_attestation(spec, state, slot=state.slot)
     run_on_attestation(spec, state, store, attestation, False)