SSZ-QL: Add endpoints (BeaconState/BeaconBlock)

api/server/structs/endpoints_beacon.go

@@ -296,3 +296,8 @@ type GetBlobsResponse struct {
 	Finalized           bool     `json:"finalized"`
 	Data                []string `json:"data"` //blobs
 }
+
+type QuerySSZRequest struct {
+	Query        string `json:"query"`
+	IncludeProof bool   `json:"include_proof,omitempty"`
+}

beacon-chain/rpc/endpoints.go

@@ -97,7 +97,7 @@ func (s *Service) endpoints(
 	endpoints = append(endpoints, s.beaconEndpoints(ch, stater, blocker, validatorServer, coreService)...)
 	endpoints = append(endpoints, s.configEndpoints()...)
 	endpoints = append(endpoints, s.eventsEndpoints()...)
-	endpoints = append(endpoints, s.prysmBeaconEndpoints(ch, stater, coreService)...)
+	endpoints = append(endpoints, s.prysmBeaconEndpoints(ch, stater, blocker, coreService)...)
 	endpoints = append(endpoints, s.prysmNodeEndpoints()...)
 	endpoints = append(endpoints, s.prysmValidatorEndpoints(stater, coreService)...)
 
@@ -1184,6 +1184,7 @@ func (s *Service) eventsEndpoints() []endpoint {
 func (s *Service) prysmBeaconEndpoints(
 	ch *stategen.CanonicalHistory,
 	stater lookup.Stater,
+	blocker lookup.Blocker,
 	coreService *core.Service,
 ) []endpoint {
 	server := &beaconprysm.Server{
@@ -1194,6 +1195,7 @@ func (s *Service) prysmBeaconEndpoints(
 		CanonicalHistory:      ch,
 		BeaconDB:              s.cfg.BeaconDB,
 		Stater:                stater,
+		Blocker:               blocker,
 		ChainInfoFetcher:      s.cfg.ChainInfoFetcher,
 		FinalizationFetcher:   s.cfg.FinalizationFetcher,
 		CoreService:           coreService,
@@ -1266,6 +1268,26 @@ func (s *Service) prysmBeaconEndpoints(
 			handler: server.PublishBlobs,
 			methods: []string{http.MethodPost},
 		},
+		{
+			template: "/prysm/v1/beacon/states/{state_id}/query",
+			name:     namespace + ".QueryBeaconState",
+			middleware: []middleware.Middleware{
+				middleware.AcceptHeaderHandler([]string{api.JsonMediaType, api.OctetStreamMediaType}),
+				middleware.AcceptEncodingHeaderHandler(),
+			},
+			handler: server.QueryBeaconState,
+			methods: []string{http.MethodPost},
+		},
+		{
+			template: "/prysm/v1/beacon/blocks/{block_id}/query",
+			name:     namespace + ".QueryBeaconBlock",
+			middleware: []middleware.Middleware{
+				middleware.AcceptHeaderHandler([]string{api.JsonMediaType, api.OctetStreamMediaType}),
+				middleware.AcceptEncodingHeaderHandler(),
+			},
+			handler: server.QueryBeaconBlock,
+			methods: []string{http.MethodPost},
+		},
 	}
 }
 

beacon-chain/rpc/endpoints_test.go

@@ -127,6 +127,8 @@ func Test_endpoints(t *testing.T) {
 		"/prysm/v1/beacon/states/{state_id}/validator_count": {http.MethodGet},
 		"/prysm/v1/beacon/chain_head":                        {http.MethodGet},
 		"/prysm/v1/beacon/blobs":                             {http.MethodPost},
+		"/prysm/v1/beacon/states/{state_id}/query":           {http.MethodPost},
+		"/prysm/v1/beacon/blocks/{block_id}/query":           {http.MethodPost},
 	}
 
 	prysmNodeRoutes := map[string][]string{

beacon-chain/rpc/prysm/beacon/BUILD.bazel

@@ -5,11 +5,13 @@ go_library(
     srcs = [
         "handlers.go",
         "server.go",
+        "ssz_query.go",
         "validator_count.go",
     ],
     importpath = "github.com/OffchainLabs/prysm/v6/beacon-chain/rpc/prysm/beacon",
     visibility = ["//visibility:public"],
     deps = [
+        "//api:go_default_library",
         "//api/server/structs:go_default_library",
         "//beacon-chain/blockchain:go_default_library",
         "//beacon-chain/core/helpers:go_default_library",
@@ -27,10 +29,13 @@ go_library(
         "//consensus-types/primitives:go_default_library",
         "//consensus-types/validator:go_default_library",
         "//encoding/bytesutil:go_default_library",
+        "//encoding/ssz/query:go_default_library",
         "//monitoring/tracing/trace:go_default_library",
         "//network/httputil:go_default_library",
         "//proto/eth/v1:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",
+        "//proto/ssz_query:go_default_library",
+        "//runtime/version:go_default_library",
         "//time/slots:go_default_library",
         "@com_github_ethereum_go_ethereum//common/hexutil:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
@@ -41,10 +46,12 @@ go_test(
     name = "go_default_test",
     srcs = [
         "handlers_test.go",
+        "ssz_query_test.go",
         "validator_count_test.go",
     ],
     embed = [":go_default_library"],
     deps = [
+        "//api:go_default_library",
         "//api/server/structs:go_default_library",
         "//beacon-chain/blockchain/testing:go_default_library",
         "//beacon-chain/core/helpers:go_default_library",
@@ -63,10 +70,13 @@ go_test(
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",
+        "//consensus-types/interfaces:go_default_library",
         "//consensus-types/primitives:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//network/httputil:go_default_library",
         "//proto/prysm/v1alpha1:go_default_library",
+        "//proto/ssz_query:go_default_library",
+        "//runtime/version:go_default_library",
         "//testing/assert:go_default_library",
         "//testing/require:go_default_library",
         "//testing/util:go_default_library",

beacon-chain/rpc/prysm/beacon/server.go

@@ -18,6 +18,7 @@ type Server struct {
 	CanonicalHistory      *stategen.CanonicalHistory
 	BeaconDB              beacondb.ReadOnlyDatabase
 	Stater                lookup.Stater
+	Blocker               lookup.Blocker
 	ChainInfoFetcher      blockchain.ChainInfoFetcher
 	FinalizationFetcher   blockchain.FinalizationFetcher
 	CoreService           *core.Service

beacon-chain/rpc/prysm/beacon/ssz_query.go

@@ -0,0 +1,228 @@
+package beacon
+
+import (
+	"encoding/json"
+	"errors"
+	"io"
+	"net/http"
+
+	"github.com/OffchainLabs/prysm/v6/api"
+	"github.com/OffchainLabs/prysm/v6/api/server/structs"
+	"github.com/OffchainLabs/prysm/v6/beacon-chain/rpc/eth/shared"
+	"github.com/OffchainLabs/prysm/v6/beacon-chain/rpc/lookup"
+	"github.com/OffchainLabs/prysm/v6/encoding/ssz/query"
+	"github.com/OffchainLabs/prysm/v6/monitoring/tracing/trace"
+	"github.com/OffchainLabs/prysm/v6/network/httputil"
+	ethpb "github.com/OffchainLabs/prysm/v6/proto/prysm/v1alpha1"
+	sszquerypb "github.com/OffchainLabs/prysm/v6/proto/ssz_query"
+	"github.com/OffchainLabs/prysm/v6/runtime/version"
+)
+
+// QueryBeaconState handles SSZ Query request for BeaconState.
+// Returns as bytes serialized SSZQueryResponse.
+func (s *Server) QueryBeaconState(w http.ResponseWriter, r *http.Request) {
+	ctx, span := trace.StartSpan(r.Context(), "beacon.QueryBeaconState")
+	defer span.End()
+
+	stateID := r.PathValue("state_id")
+	if stateID == "" {
+		httputil.HandleError(w, "state_id is required in URL params", http.StatusBadRequest)
+		return
+	}
+
+	// Validate path before lookup: it might be expensive.
+	var req structs.QuerySSZRequest
+	err := json.NewDecoder(r.Body).Decode(&req)
+	switch {
+	case errors.Is(err, io.EOF):
+		httputil.HandleError(w, "No data submitted", http.StatusBadRequest)
+		return
+	case err != nil:
+		httputil.HandleError(w, "Could not decode request body: "+err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	if len(req.Query) == 0 {
+		httputil.HandleError(w, "Empty query submitted", http.StatusBadRequest)
+		return
+	}
+
+	path, err := query.ParsePath(req.Query)
+	if err != nil {
+		httputil.HandleError(w, "Could not parse path '"+req.Query+"': "+err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	stateRoot, err := s.Stater.StateRoot(ctx, []byte(stateID))
+	if err != nil {
+		var rootNotFoundErr *lookup.StateRootNotFoundError
+		if errors.As(err, &rootNotFoundErr) {
+			httputil.HandleError(w, "State root not found: "+rootNotFoundErr.Error(), http.StatusNotFound)
+			return
+		}
+		httputil.HandleError(w, "Could not get state root: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	st, err := s.Stater.State(ctx, []byte(stateID))
+	if err != nil {
+		shared.WriteStateFetchError(w, err)
+		return
+	}
+
+	var sszObject query.SSZObject
+	switch st.Version() {
+	case version.Phase0:
+		sszObject = st.ToProto().(*ethpb.BeaconState)
+	case version.Altair:
+		sszObject = st.ToProto().(*ethpb.BeaconStateAltair)
+	case version.Bellatrix:
+		sszObject = st.ToProto().(*ethpb.BeaconStateBellatrix)
+	case version.Capella:
+		sszObject = st.ToProto().(*ethpb.BeaconStateCapella)
+	case version.Deneb:
+		sszObject = st.ToProto().(*ethpb.BeaconStateDeneb)
+	case version.Fulu:
+		sszObject = st.ToProto().(*ethpb.BeaconStateFulu)
+	default:
+		httputil.HandleError(w, "Unsupported state version for querying: "+version.String(st.Version()), http.StatusBadRequest)
+		return
+	}
+
+	info, err := query.AnalyzeObject(sszObject)
+	if err != nil {
+		httputil.HandleError(w, "Could not analyze state object: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	encodedState, err := st.MarshalSSZ()
+	if err != nil {
+		httputil.HandleError(w, "Could not marshal state to SSZ: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	_, offset, length, err := query.CalculateOffsetAndLength(info, path)
+	if err != nil {
+		httputil.HandleError(w, "Could not calculate offset and length for path '"+req.Query+"': "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	response := &sszquerypb.SSZQueryResponse{
+		Root:   stateRoot,
+		Result: encodedState[offset : offset+length],
+	}
+
+	responseSsz, err := response.MarshalSSZ()
+	if err != nil {
+		httputil.HandleError(w, "Could not marshal response to SSZ: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set(api.VersionHeader, version.String(st.Version()))
+	httputil.WriteSsz(w, responseSsz)
+}
+
+// QueryBeaconState handles SSZ Query request for BeaconState.
+// Returns as bytes serialized SSZQueryResponse.
+func (s *Server) QueryBeaconBlock(w http.ResponseWriter, r *http.Request) {
+	ctx, span := trace.StartSpan(r.Context(), "beacon.QueryBeaconBlock")
+	defer span.End()
+
+	blockId := r.PathValue("block_id")
+	if blockId == "" {
+		httputil.HandleError(w, "block_id is required in URL params", http.StatusBadRequest)
+		return
+	}
+
+	// Validate path before lookup: it might be expensive.
+	var req structs.QuerySSZRequest
+	err := json.NewDecoder(r.Body).Decode(&req)
+	switch {
+	case errors.Is(err, io.EOF):
+		httputil.HandleError(w, "No data submitted", http.StatusBadRequest)
+		return
+	case err != nil:
+		httputil.HandleError(w, "Could not decode request body: "+err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	if len(req.Query) == 0 {
+		httputil.HandleError(w, "Empty query submitted", http.StatusBadRequest)
+		return
+	}
+
+	path, err := query.ParsePath(req.Query)
+	if err != nil {
+		httputil.HandleError(w, "Could not parse path '"+req.Query+"': "+err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	signedBlock, err := s.Blocker.Block(ctx, []byte(blockId))
+	if !shared.WriteBlockFetchError(w, signedBlock, err) {
+		return
+	}
+
+	encodedBlock, err := signedBlock.Block().MarshalSSZ()
+	if err != nil {
+		httputil.HandleError(w, "Could not marshal block to SSZ: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	protoBlock, err := signedBlock.Block().Proto()
+	if err != nil {
+		httputil.HandleError(w, "Could not convert block to proto: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	var block query.SSZObject
+	switch signedBlock.Version() {
+	case version.Phase0:
+		block = protoBlock.(*ethpb.BeaconBlock)
+	case version.Altair:
+		block = protoBlock.(*ethpb.BeaconBlockAltair)
+	case version.Bellatrix:
+		block = protoBlock.(*ethpb.BeaconBlockBellatrix)
+	case version.Capella:
+		block = protoBlock.(*ethpb.BeaconBlockCapella)
+	case version.Deneb:
+		block = protoBlock.(*ethpb.BeaconBlockDeneb)
+	case version.Electra:
+	case version.Fulu:
+		block = protoBlock.(*ethpb.BeaconBlockElectra)
+	default:
+		httputil.HandleError(w, "Unsupported block version for querying: "+version.String(signedBlock.Version()), http.StatusBadRequest)
+		return
+	}
+
+	blockRoot, err := block.HashTreeRoot()
+	if err != nil {
+		httputil.HandleError(w, "Could not compute block root: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	info, err := query.AnalyzeObject(block)
+	if err != nil {
+		httputil.HandleError(w, "Could not analyze block object: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	_, offset, length, err := query.CalculateOffsetAndLength(info, path)
+	if err != nil {
+		httputil.HandleError(w, "Could not calculate offset and length for path '"+req.Query+"': "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	response := &sszquerypb.SSZQueryResponse{
+		Root:   blockRoot[:],
+		Result: encodedBlock[offset : offset+length],
+	}
+
+	responseSsz, err := response.MarshalSSZ()
+	if err != nil {
+		httputil.HandleError(w, "Could not marshal response to SSZ: "+err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	w.Header().Set(api.VersionHeader, version.String(signedBlock.Version()))
+	httputil.WriteSsz(w, responseSsz)
+}