Add p2p support for Electra (prysmaticlabs#13985)

* Add p2p support for Electra * Fix TestGossipTopicMappings_CorrectBlockType
ErnestK · May 19, 2024 · 469a1a9 · 469a1a9
1 parent 0533dea
commit 469a1a9
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 1 deletion.
diff --git a/beacon-chain/p2p/fork_watcher.go b/beacon-chain/p2p/fork_watcher.go
@@ -17,7 +17,8 @@ func (s *Service) forkWatcher() {
 			if currEpoch == params.BeaconConfig().AltairForkEpoch ||
 				currEpoch == params.BeaconConfig().BellatrixForkEpoch ||
 				currEpoch == params.BeaconConfig().CapellaForkEpoch ||
-				currEpoch == params.BeaconConfig().DenebForkEpoch {
+				currEpoch == params.BeaconConfig().DenebForkEpoch ||
+				currEpoch == params.BeaconConfig().ElectraForkEpoch {
 				// If we are in the fork epoch, we update our enr with
 				// the updated fork digest. These repeatedly does
 				// this over the epoch, which might be slightly wasteful

diff --git a/beacon-chain/p2p/gossip_topic_mappings.go b/beacon-chain/p2p/gossip_topic_mappings.go
@@ -28,6 +28,9 @@ var gossipTopicMappings = map[string]proto.Message{
 // versioned by epoch.
 func GossipTopicMappings(topic string, epoch primitives.Epoch) proto.Message {
 	if topic == BlockSubnetTopicFormat {
+		if epoch >= params.BeaconConfig().ElectraForkEpoch {
+			return &ethpb.SignedBeaconBlockElectra{}
+		}
 		if epoch >= params.BeaconConfig().DenebForkEpoch {
 			return &ethpb.SignedBeaconBlockDeneb{}
 		}
@@ -70,4 +73,6 @@ func init() {
 	GossipTypeMapping[reflect.TypeOf(&ethpb.SignedBeaconBlockCapella{})] = BlockSubnetTopicFormat
 	// Specially handle Deneb objects.
 	GossipTypeMapping[reflect.TypeOf(&ethpb.SignedBeaconBlockDeneb{})] = BlockSubnetTopicFormat
+	// Specially handle Electra objects.
+	GossipTypeMapping[reflect.TypeOf(&ethpb.SignedBeaconBlockElectra{})] = BlockSubnetTopicFormat
 }
diff --git a/beacon-chain/p2p/gossip_topic_mappings_test.go b/beacon-chain/p2p/gossip_topic_mappings_test.go
@@ -28,13 +28,19 @@ func TestGossipTopicMappings_CorrectBlockType(t *testing.T) {
 	altairForkEpoch := primitives.Epoch(100)
 	BellatrixForkEpoch := primitives.Epoch(200)
 	CapellaForkEpoch := primitives.Epoch(300)
+	DenebForkEpoch := primitives.Epoch(400)
+	ElectraForkEpoch := primitives.Epoch(500)
 
 	bCfg.AltairForkEpoch = altairForkEpoch
 	bCfg.BellatrixForkEpoch = BellatrixForkEpoch
 	bCfg.CapellaForkEpoch = CapellaForkEpoch
+	bCfg.DenebForkEpoch = DenebForkEpoch
+	bCfg.ElectraForkEpoch = ElectraForkEpoch
 	bCfg.ForkVersionSchedule[bytesutil.ToBytes4(bCfg.AltairForkVersion)] = primitives.Epoch(100)
 	bCfg.ForkVersionSchedule[bytesutil.ToBytes4(bCfg.BellatrixForkVersion)] = primitives.Epoch(200)
 	bCfg.ForkVersionSchedule[bytesutil.ToBytes4(bCfg.CapellaForkVersion)] = primitives.Epoch(300)
+	bCfg.ForkVersionSchedule[bytesutil.ToBytes4(bCfg.DenebForkVersion)] = primitives.Epoch(400)
+	bCfg.ForkVersionSchedule[bytesutil.ToBytes4(bCfg.ElectraForkVersion)] = primitives.Epoch(500)
 	params.OverrideBeaconConfig(bCfg)
 
 	// Phase 0
@@ -56,4 +62,14 @@ func TestGossipTopicMappings_CorrectBlockType(t *testing.T) {
 	pMessage = GossipTopicMappings(BlockSubnetTopicFormat, CapellaForkEpoch)
 	_, ok = pMessage.(*ethpb.SignedBeaconBlockCapella)
 	assert.Equal(t, true, ok)
+
+	// Deneb Fork
+	pMessage = GossipTopicMappings(BlockSubnetTopicFormat, DenebForkEpoch)
+	_, ok = pMessage.(*ethpb.SignedBeaconBlockDeneb)
+	assert.Equal(t, true, ok)
+
+	// Electra Fork
+	pMessage = GossipTopicMappings(BlockSubnetTopicFormat, ElectraForkEpoch)
+	_, ok = pMessage.(*ethpb.SignedBeaconBlockElectra)
+	assert.Equal(t, true, ok)
 }
diff --git a/beacon-chain/p2p/pubsub_filter.go b/beacon-chain/p2p/pubsub_filter.go
@@ -62,12 +62,18 @@ func (s *Service) CanSubscribe(topic string) bool {
 		log.WithError(err).Error("Could not determine Deneb fork digest")
 		return false
 	}
+	electraForkDigest, err := forks.ForkDigestFromEpoch(params.BeaconConfig().ElectraForkEpoch, s.genesisValidatorsRoot)
+	if err != nil {
+		log.WithError(err).Error("Could not determine Electra fork digest")
+		return false
+	}
 	switch parts[2] {
 	case fmt.Sprintf("%x", phase0ForkDigest):
 	case fmt.Sprintf("%x", altairForkDigest):
 	case fmt.Sprintf("%x", bellatrixForkDigest):
 	case fmt.Sprintf("%x", capellaForkDigest):
 	case fmt.Sprintf("%x", denebForkDigest):
+	case fmt.Sprintf("%x", electraForkDigest):
 	default:
 		return false
 	}

diff --git a/beacon-chain/p2p/types/object_mapping.go b/beacon-chain/p2p/types/object_mapping.go
@@ -59,6 +59,11 @@ func InitializeDataMaps() {
 				&ethpb.SignedBeaconBlockDeneb{Block: &ethpb.BeaconBlockDeneb{Body: &ethpb.BeaconBlockBodyDeneb{ExecutionPayload: &enginev1.ExecutionPayloadDeneb{}}}},
 			)
 		},
+		bytesutil.ToBytes4(params.BeaconConfig().ElectraForkVersion): func() (interfaces.ReadOnlySignedBeaconBlock, error) {
+			return blocks.NewSignedBeaconBlock(
+				&ethpb.SignedBeaconBlockElectra{Block: &ethpb.BeaconBlockElectra{Body: &ethpb.BeaconBlockBodyElectra{ExecutionPayload: &enginev1.ExecutionPayloadElectra{}}}},
+			)
+		},
 	}
 
 	// Reset our metadata map.
@@ -78,5 +83,8 @@ func InitializeDataMaps() {
 		bytesutil.ToBytes4(params.BeaconConfig().DenebForkVersion): func() metadata.Metadata {
 			return wrapper.WrappedMetadataV1(&ethpb.MetaDataV1{})
 		},
+		bytesutil.ToBytes4(params.BeaconConfig().ElectraForkVersion): func() metadata.Metadata {
+			return wrapper.WrappedMetadataV1(&ethpb.MetaDataV1{})
+		},
 	}
 }