Fix Fork Data in Local Node

beacon-chain/p2p/BUILD.bazel

@@ -36,6 +36,7 @@ go_library(
         "//beacon-chain/cache:go_default_library",
         "//beacon-chain/core/feed:go_default_library",
         "//beacon-chain/core/feed/state:go_default_library",
+        "//beacon-chain/core/helpers:go_default_library",
         "//beacon-chain/p2p/connmgr:go_default_library",
         "//beacon-chain/p2p/encoder:go_default_library",
         "//beacon-chain/p2p/peers:go_default_library",

beacon-chain/p2p/fork.go

@@ -4,15 +4,18 @@ import (
 	"bytes"
 	"encoding/base64"
 	"fmt"
+	"math"
 	"time"
 
 	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/go-ssz"
+	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/p2putils"
 	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/roughtime"
 	"github.com/sirupsen/logrus"
 )
 
@@ -86,10 +89,25 @@ func addForkEntry(
 	if err != nil {
 		return nil, err
 	}
+	currentSlot := helpers.SlotsSince(genesisTime)
+	currentEpoch := helpers.SlotToEpoch(currentSlot)
+	if roughtime.Now().Before(genesisTime) {
+		currentSlot, currentEpoch = 0, 0
+	}
+	fork, err := p2putils.Fork(currentEpoch)
+	if err != nil {
+		return nil, err
+	}
+
 	nextForkEpoch := params.BeaconConfig().NextForkEpoch
+	nextForkVersion := params.BeaconConfig().NextForkVersion
+	// Set to the current fork version if our next fork is not planned.
+	if nextForkEpoch == math.MaxUint64 {
+		nextForkVersion = fork.CurrentVersion
+	}
 	enrForkID := &pb.ENRForkID{
 		CurrentForkDigest: digest[:],
-		NextForkVersion:   params.BeaconConfig().NextForkVersion,
+		NextForkVersion:   nextForkVersion,
 		NextForkEpoch:     nextForkEpoch,
 	}
 	enc, err := ssz.Marshal(enrForkID)

beacon-chain/p2p/fork_test.go

@@ -258,3 +258,33 @@ func TestDiscv5_AddRetrieveForkEntryENR(t *testing.T) {
 		t.Errorf("Wanted next for epoch: %d, received: %d", nextForkEpoch, resp.NextForkEpoch)
 	}
 }
+
+func TestAddForkEntry_Genesis(t *testing.T) {
+	temp := testutil.TempDir()
+	randNum := rand.Int()
+	tempPath := path.Join(temp, strconv.Itoa(randNum))
+	if err := os.Mkdir(tempPath, 0700); err != nil {
+		t.Fatal(err)
+	}
+	pkey, err := privKey(&Config{Encoding: "ssz", DataDir: tempPath})
+	if err != nil {
+		t.Fatalf("Could not get private key: %v", err)
+	}
+	db, err := enode.OpenDB("")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	localNode := enode.NewLocalNode(db, pkey)
+	localNode, err = addForkEntry(localNode, time.Now().Add(10*time.Second), []byte{'A', 'B', 'C', 'D'})
+	if err != nil {
+		t.Fatal(err)
+	}
+	forkEntry, err := retrieveForkEntry(localNode.Node().Record())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(forkEntry.NextForkVersion, params.BeaconConfig().GenesisForkVersion) {
+		t.Errorf("Wanted Next Fork Version to be equal to genesis fork version, instead got %#x", forkEntry.NextForkVersion)
+	}
+}

shared/p2putils/fork.go

@@ -26,19 +26,12 @@ func CreateForkDigest(
 	currentSlot := helpers.SlotsSince(genesisTime)
 	currentEpoch := helpers.SlotToEpoch(currentSlot)
 
-	// We retrieve a list of scheduled forks by epoch.
-	// We loop through the keys in this map to determine the current
-	// fork version based on the current, time-based epoch number
-	// since the genesis time.
-	currentForkVersion := params.BeaconConfig().GenesisForkVersion
-	scheduledForks := params.BeaconConfig().ForkVersionSchedule
-	for epoch, forkVersion := range scheduledForks {
-		if epoch <= currentEpoch {
-			currentForkVersion = forkVersion
-		}
+	forkData, err := Fork(currentEpoch)
+	if err != nil {
+		return [4]byte{}, err
 	}
 
-	digest, err := helpers.ComputeForkDigest(currentForkVersion, genesisValidatorsRoot)
+	digest, err := helpers.ComputeForkDigest(forkData.CurrentVersion, genesisValidatorsRoot)
 	if err != nil {
 		return [4]byte{}, err
 	}
@@ -62,10 +55,8 @@ func Fork(
 			previousForkVersion = retrievedForkVersion
 			retrievedForkVersion = forkVersion
 			forkEpoch = epoch
-
 		}
 	}
-
 	return &pb.Fork{
 		PreviousVersion: previousForkVersion,
 		CurrentVersion:  retrievedForkVersion,