Implement GetBeaconState Endpoint

beacon-chain/rpc/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "//beacon-chain/state/stategen:go_default_library",
         "//beacon-chain/sync:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
+        "//proto/beacon/rpc/v1:go_default_library",
         "//proto/slashing:go_default_library",
         "//shared/featureconfig:go_default_library",
         "//shared/params:go_default_library",

beacon-chain/rpc/beacon/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "config.go",
         "server.go",
         "slashings.go",
+        "state.go",
         "validators.go",
         "validators_stream.go",
     ],
@@ -37,6 +38,7 @@ go_library(
         "//beacon-chain/state/stategen:go_default_library",
         "//beacon-chain/state/stateutil:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
+        "//proto/beacon/rpc/v1:go_default_library",
         "//shared/attestationutil:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/event:go_default_library",
@@ -67,6 +69,7 @@ go_test(
         "committees_test.go",
         "config_test.go",
         "slashings_test.go",
+        "state_test.go",
         "validators_stream_test.go",
         "validators_test.go",
     ],
@@ -92,6 +95,7 @@ go_test(
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/stategen:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
+        "//proto/beacon/rpc/v1:go_default_library",
         "//shared/attestationutil:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/featureconfig:go_default_library",

beacon-chain/rpc/beacon/state.go

@@ -0,0 +1,40 @@
+package beacon
+
+import (
+	"context"
+
+	pbp2p "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	pbrpc "github.com/prysmaticlabs/prysm/proto/beacon/rpc/v1"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+// GetBeaconState retrieves a beacon state
+// from the beacon node by either a slot or block root.
+func (bs *Server) GetBeaconState(
+	ctx context.Context,
+	req *pbrpc.BeaconStateRequest,
+) (*pbp2p.BeaconState, error) {
+	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:
+		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)
+		}
+		return st.CloneInnerState(), nil
+	case *pbrpc.BeaconStateRequest_BlockRoot:
+		st, err := bs.StateGen.StateByRoot(ctx, bytesutil.ToBytes32(q.BlockRoot))
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "could not compute state by block root: %v", err)
+		}
+		return st.CloneInnerState(), nil
+	default:
+		return nil, status.Error(codes.InvalidArgument, "need to specify either a block root or slot to request state")
+	}
+}

beacon-chain/rpc/beacon/state_test.go

@@ -0,0 +1,80 @@
+package beacon
+
+import (
+	"context"
+	"testing"
+
+	"github.com/gogo/protobuf/proto"
+	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/go-ssz"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
+	dbTest "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state/stategen"
+	pbrpc "github.com/prysmaticlabs/prysm/proto/beacon/rpc/v1"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+	"github.com/prysmaticlabs/prysm/shared/testutil"
+)
+
+func TestServer_GetBeaconState(t *testing.T) {
+	resetCfg := featureconfig.InitWithReset(&featureconfig.Flags{NewStateMgmt: true})
+	defer resetCfg()
+
+	db := dbTest.SetupDB(t)
+	defer dbTest.TeardownDB(t, db)
+
+	ctx := context.Background()
+	st := testutil.NewBeaconState()
+	slot := uint64(100)
+	if err := st.SetSlot(slot); err != nil {
+		t.Fatal(err)
+	}
+	b := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{
+		Slot: slot,
+	}}
+	if err := db.SaveBlock(ctx, b); err != nil {
+		t.Fatal(err)
+	}
+	gRoot, err := ssz.HashTreeRoot(b.Block)
+	if err != nil {
+		t.Fatal(err)
+	}
+	gen := stategen.New(db, cache.NewStateSummaryCache())
+	if err := gen.SaveState(ctx, gRoot, st); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.SaveState(ctx, st, gRoot); err != nil {
+		t.Fatal(err)
+	}
+	bs := &Server{
+		BeaconDB: db,
+		StateGen: gen,
+	}
+	if _, err := bs.GetBeaconState(ctx, &pbrpc.BeaconStateRequest{}); err == nil {
+		t.Errorf("Expected error without a query filter, received nil")
+	}
+	req := &pbrpc.BeaconStateRequest{
+		QueryFilter: &pbrpc.BeaconStateRequest_BlockRoot{
+			BlockRoot: gRoot[:],
+		},
+	}
+	res, err := bs.GetBeaconState(ctx, req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	wanted := st.CloneInnerState()
+	if !proto.Equal(wanted, res) {
+		t.Errorf("Wanted %v, received %v", wanted, res)
+	}
+	req = &pbrpc.BeaconStateRequest{
+		QueryFilter: &pbrpc.BeaconStateRequest_Slot{
+			Slot: slot,
+		},
+	}
+	res, err = bs.GetBeaconState(ctx, req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.Equal(wanted, res) {
+		t.Errorf("Wanted %v, received %v", wanted, res)
+	}
+}

beacon-chain/rpc/service.go

@@ -31,6 +31,7 @@ import (
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stategen"
 	"github.com/prysmaticlabs/prysm/beacon-chain/sync"
 	pbp2p "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	pbrpc "github.com/prysmaticlabs/prysm/proto/beacon/rpc/v1"
 	slashpb "github.com/prysmaticlabs/prysm/proto/slashing"
 	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/params"
@@ -269,6 +270,9 @@ func (s *Service) Start() {
 	}
 	ethpb.RegisterNodeServer(s.grpcServer, nodeServer)
 	ethpb.RegisterBeaconChainServer(s.grpcServer, beaconChainServer)
+	if featureconfig.Get().EnableDebugRPCEndpoints {
+		pbrpc.RegisterDebugServer(s.grpcServer, beaconChainServer)
+	}
 	ethpb.RegisterBeaconNodeValidatorServer(s.grpcServer, validatorServer)
 
 	// Register reflection service on gRPC server.

proto/beacon/rpc/v1/debug.proto

@@ -26,8 +26,5 @@ message BeaconStateRequest {
 
         // The block root corresponding to a desired beacon state.
         bytes block_root = 2;
-
-        // The state root corresponding to a desired beacon state.
-        bytes state_root = 3;
     }
 }

shared/featureconfig/config.go

@@ -30,6 +30,7 @@ var log = logrus.WithField("prefix", "flags")
 // Flags is a struct to represent which features the client will perform on runtime.
 type Flags struct {
 	MinimalConfig                              bool // MinimalConfig as defined in the spec.
+	EnableDebugRPCEndpoints                    bool // Enables the debug rpc service, providing utilities such as /v1/beacon/state.
 	WriteSSZStateTransitions                   bool // WriteSSZStateTransitions to tmp directory.
 	InitSyncNoVerify                           bool // InitSyncNoVerify when initial syncing w/o verifying block's contents.
 	DisableDynamicCommitteeSubnets             bool // Disables dynamic attestation committee subnets via p2p.
@@ -132,6 +133,7 @@ func (c *Flags) Copy() *Flags {
 		EnableBlockTreeCache:                       c.EnableBlockTreeCache,
 		KafkaBootstrapServers:                      c.KafkaBootstrapServers,
 		CustomGenesisDelay:                         c.CustomGenesisDelay,
+		EnableDebugRPCEndpoints:                    c.EnableDebugRPCEndpoints,
 	}
 }
 
@@ -253,6 +255,10 @@ func ConfigureBeaconChain(ctx *cli.Context) {
 		log.Warn("Enabling broadcast slashing to p2p network")
 		cfg.BroadcastSlashings = true
 	}
+	if ctx.Bool(enableDebugRPCEndpoints.Name) {
+		log.Warn("Enabling debug RPC endpoints")
+		cfg.EnableDebugRPCEndpoints = true
+	}
 	Init(cfg)
 }
 

shared/featureconfig/flags.go

@@ -150,6 +150,10 @@ var (
 		Name:  "wait-for-synced",
 		Usage: "Uses WaitForSynced for validator startup, to ensure a validator is able to communicate with the beacon node as quick as possible",
 	}
+	enableDebugRPCEndpoints = &cli.BoolFlag{
+		Name:  "enable-debug-rpc-endpoints",
+		Usage: "Enables the debug rpc service, containing utility endpoints such as /v1/beacon/state. Requires --new-state-mgmt",
+	}
 )
 
 // devModeFlags holds list of flags that are set when development mode is on.
@@ -398,6 +402,7 @@ var BeaconChainFlags = append(deprecatedFlags, []cli.Flag{
 	disableInitSyncBatchSaveBlocks,
 	enableStateRefCopy,
 	waitForSyncedFlag,
+	enableDebugRPCEndpoints,
 }...)
 
 // E2EBeaconChainFlags contains a list of the beacon chain feature flags to be tested in E2E.