From aea7a8d291edc14f04bcafc1bd45fe0e97054aae Mon Sep 17 00:00:00 2001
From: terence tsao <terence@prysmaticlabs.com>
Date: Tue, 5 May 2020 18:42:11 -0700
Subject: [PATCH] Future slot check for state end point (#5755)

---
 beacon-chain/rpc/beacon/state.go      | 12 ++++++++++++
 beacon-chain/rpc/beacon/state_test.go | 23 +++++++++++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/beacon-chain/rpc/beacon/state.go b/beacon-chain/rpc/beacon/state.go
index 7b32c1de9095..d6d9e5f7c029 100644
--- a/beacon-chain/rpc/beacon/state.go
+++ b/beacon-chain/rpc/beacon/state.go
@@ -20,8 +20,20 @@ func (bs *Server) GetBeaconState(
 	if !featureconfig.Get().NewStateMgmt {
 		return nil, status.Error(codes.FailedPrecondition, "requires --enable-new-state-mgmt to function")
 	}
+
 	switch q := req.QueryFilter.(type) {
 	case *pbrpc.BeaconStateRequest_Slot:
+		currentSlot := bs.GenesisTimeFetcher.CurrentSlot()
+		requestedSlot := q.Slot
+		if requestedSlot > currentSlot {
+			return nil, status.Errorf(
+				codes.InvalidArgument,
+				"Cannot retrieve information about a slot in the future, current slot %d, requested slot %d",
+				currentSlot,
+				requestedSlot,
+			)
+		}
+
 		st, err := bs.StateGen.StateBySlot(ctx, q.Slot)
 		if err != nil {
 			return nil, status.Errorf(codes.Internal, "could not compute state by slot: %v", err)
diff --git a/beacon-chain/rpc/beacon/state_test.go b/beacon-chain/rpc/beacon/state_test.go
index 6b2e4fdd6afb..0fd57c816bf1 100644
--- a/beacon-chain/rpc/beacon/state_test.go
+++ b/beacon-chain/rpc/beacon/state_test.go
@@ -2,10 +2,12 @@ package beacon
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/gogo/protobuf/proto"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	mock "github.com/prysmaticlabs/prysm/beacon-chain/blockchain/testing"
 	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	dbTest "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stategen"
@@ -44,8 +46,9 @@ func TestServer_GetBeaconState(t *testing.T) {
 		t.Fatal(err)
 	}
 	bs := &Server{
-		BeaconDB: db,
-		StateGen: gen,
+		BeaconDB:           db,
+		StateGen:           gen,
+		GenesisTimeFetcher: &mock.ChainService{},
 	}
 	if _, err := bs.GetBeaconState(ctx, &pbrpc.BeaconStateRequest{}); err == nil {
 		t.Errorf("Expected error without a query filter, received nil")
@@ -76,3 +79,19 @@ func TestServer_GetBeaconState(t *testing.T) {
 		t.Errorf("Wanted %v, received %v", wanted, res)
 	}
 }
+
+func TestServer_GetBeaconState_RequestFutureSlot(t *testing.T) {
+	resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{NewStateMgmt: true})
+	defer resetCfg()
+
+	bs := &Server{GenesisTimeFetcher: &mock.ChainService{}}
+	req := &pbrpc.BeaconStateRequest{
+		QueryFilter: &pbrpc.BeaconStateRequest_Slot{
+			Slot: bs.GenesisTimeFetcher.CurrentSlot() + 1,
+		},
+	}
+	wanted := "Cannot retrieve information about a slot in the future"
+	if _, err := bs.GetBeaconState(context.Background(), req); err != nil && !strings.Contains(err.Error(), wanted) {
+		t.Errorf("Expected error %v, received %v", wanted, err)
+	}
+}