From 6b55d33ea2b1e2286b52cca0b823cd4125f08450 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rados=C5=82aw=20Kapka?= <rkapka@wp.pl>
Date: Wed, 22 Jun 2022 01:13:19 +0200
Subject: [PATCH] Return error when requesting future state (#10915)

* Return error when requesting future state

* fix genesis test case

* fix test cases
---
 beacon-chain/rpc/statefetcher/BUILD.bazel     |  1 +
 beacon-chain/rpc/statefetcher/fetcher.go      |  7 ++++
 beacon-chain/rpc/statefetcher/fetcher_test.go | 38 ++++++++++++-------
 3 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/beacon-chain/rpc/statefetcher/BUILD.bazel b/beacon-chain/rpc/statefetcher/BUILD.bazel
index abc7ba2aa..d80cc6974 100644
--- a/beacon-chain/rpc/statefetcher/BUILD.bazel
+++ b/beacon-chain/rpc/statefetcher/BUILD.bazel
@@ -26,6 +26,7 @@ go_test(
     deps = [
         "//beacon-chain/blockchain/testing:go_default_library",
         "//beacon-chain/db/testing:go_default_library",
+        "//beacon-chain/state/state-native:go_default_library",
         "//beacon-chain/state/stategen:go_default_library",
         "//beacon-chain/state/stategen/mock:go_default_library",
         "//config/params:go_default_library",
diff --git a/beacon-chain/rpc/statefetcher/fetcher.go b/beacon-chain/rpc/statefetcher/fetcher.go
index 602d9a455..9d30d376f 100644
--- a/beacon-chain/rpc/statefetcher/fetcher.go
+++ b/beacon-chain/rpc/statefetcher/fetcher.go
@@ -206,6 +206,13 @@ func (p *StateProvider) StateBySlot(ctx context.Context, target types.Slot) (sta
 	ctx, span := trace.StartSpan(ctx, "statefetcher.StateBySlot")
 	defer span.End()
 
+	if target > p.GenesisTimeFetcher.CurrentSlot() {
+		return nil, errors.New("requested slot is in the future")
+	}
+	if target > p.ChainInfoFetcher.HeadSlot() {
+		return nil, errors.New("requested slot number is higher than head slot number")
+	}
+
 	st, err := p.ReplayerBuilder.ReplayerForSlot(target).ReplayBlocks(ctx)
 	if err != nil {
 		msg := fmt.Sprintf("error while replaying history to slot=%d", target)
diff --git a/beacon-chain/rpc/statefetcher/fetcher_test.go b/beacon-chain/rpc/statefetcher/fetcher_test.go
index f14d5a884..587cc95bb 100644
--- a/beacon-chain/rpc/statefetcher/fetcher_test.go
+++ b/beacon-chain/rpc/statefetcher/fetcher_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	statenative "github.com/prysmaticlabs/prysm/beacon-chain/state/state-native"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stategen"
 
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -78,9 +79,12 @@ func TestGetState(t *testing.T) {
 		cc := &mockstategen.MockCanonicalChecker{Is: true}
 		cs := &mockstategen.MockCurrentSlotter{Slot: bs.Slot() + 1}
 		ch := stategen.NewCanonicalHistory(db, cc, cs)
+		currentSlot := types.Slot(0)
 		p := StateProvider{
-			BeaconDB:        db,
-			ReplayerBuilder: ch,
+			BeaconDB:           db,
+			ReplayerBuilder:    ch,
+			GenesisTimeFetcher: &chainMock.ChainService{Slot: &currentSlot},
+			ChainInfoFetcher:   &chainMock.ChainService{State: bs},
 		}
 
 		s, err := p.State(ctx, []byte("genesis"))
@@ -165,6 +169,7 @@ func TestGetState(t *testing.T) {
 				CanonicalRoots: map[[32]byte]bool{
 					bytesutil.ToBytes32(newBeaconState.LatestBlockHeader().ParentRoot): true,
 				},
+				State: newBeaconState,
 			},
 			ReplayerBuilder: mockstategen.NewMockReplayerBuilder(mockstategen.WithMockState(newBeaconState)),
 		}
@@ -176,18 +181,6 @@ func TestGetState(t *testing.T) {
 		assert.Equal(t, stateRoot, sRoot)
 	})
 
-	rb := mockstategen.NewMockReplayerBuilder(mockstategen.WithStateError(1, stategen.ErrFutureSlotRequested))
-	t.Run("slot_too_big", func(t *testing.T) {
-		p := StateProvider{
-			GenesisTimeFetcher: &chainMock.ChainService{
-				Genesis: time.Now(),
-			},
-			ReplayerBuilder: rb,
-		}
-		_, err := p.State(ctx, []byte(strconv.FormatUint(1, 10)))
-		assert.ErrorContains(t, "cannot replay to future slots", err)
-	})
-
 	t.Run("invalid_state", func(t *testing.T) {
 		p := StateProvider{}
 		_, err := p.State(ctx, []byte("foo"))
@@ -394,3 +387,20 @@ func TestNewStateNotFoundError(t *testing.T) {
 	e := NewStateNotFoundError(100)
 	assert.Equal(t, "state not found in the last 100 state roots", e.message)
 }
+
+func TestStateBySlot_FutureSlot(t *testing.T) {
+	slot := types.Slot(100)
+	p := StateProvider{GenesisTimeFetcher: &chainMock.ChainService{Slot: &slot}}
+	_, err := p.StateBySlot(context.Background(), 101)
+	assert.ErrorContains(t, "requested slot is in the future", err)
+}
+
+func TestStateBySlot_AfterHeadSlot(t *testing.T) {
+	st, err := statenative.InitializeFromProtoPhase0(&ethpb.BeaconState{Slot: 100})
+	require.NoError(t, err)
+	currentSlot := types.Slot(102)
+	mock := &chainMock.ChainService{State: st, Slot: &currentSlot}
+	p := StateProvider{ChainInfoFetcher: mock, GenesisTimeFetcher: mock}
+	_, err = p.StateBySlot(context.Background(), 101)
+	assert.ErrorContains(t, "requested slot number is higher than head slot number", err)
+}