PeerDAS: Generate private key to match subnets.

beacon-chain/core/peerdas/helpers.go

@@ -30,7 +30,7 @@ var (
 )
 
 // CustodyColumnSubnets computes the subnets the node should participate in for custody.
-func CustodyColumnSubnets(nodeId enode.ID, custodySubnetCount uint64) (map[uint64]bool, error) {
+func CustodyColumnSubnets(nodeId enode.ID, custodySubnetCount uint64) ([]uint64, error) {
 	dataColumnSidecarSubnetCount := params.BeaconConfig().DataColumnSidecarSubnetCount
 
 	// Check if the custody subnet count is larger than the data column sidecar subnet count.
@@ -38,11 +38,9 @@ func CustodyColumnSubnets(nodeId enode.ID, custodySubnetCount uint64) (map[uint6
 		return nil, errCustodySubnetCountTooLarge
 	}
 
-	// First, compute the subnet IDs that the node should participate in.
-	subnetIds := make(map[uint64]bool, custodySubnetCount)
-
 	one := uint256.NewInt(1)
 
+	subnetIds, subnetIdsMap := make([]uint64, 0, custodySubnetCount), make(map[uint64]bool, custodySubnetCount)
 	for currentId := new(uint256.Int).SetBytes(nodeId.Bytes()); uint64(len(subnetIds)) < custodySubnetCount; currentId.Add(currentId, one) {
 		// Convert to big endian bytes.
 		currentIdBytesBigEndian := currentId.Bytes32()
@@ -56,8 +54,12 @@ func CustodyColumnSubnets(nodeId enode.ID, custodySubnetCount uint64) (map[uint6
 		// Get the subnet ID.
 		subnetId := binary.LittleEndian.Uint64(hashedCurrentId[:8]) % dataColumnSidecarSubnetCount
 
-		// Add the subnet to the map.
-		subnetIds[subnetId] = true
+		// Add the subnet to the slice.
+		exists := subnetIdsMap[subnetId]
+		if !exists {
+			subnetIds = append(subnetIds, subnetId)
+			subnetIdsMap[subnetId] = true
+		}
 
 		// Overflow prevention.
 		if currentId.Cmp(maxUint256) == 0 {
@@ -85,7 +87,7 @@ func CustodyColumns(nodeId enode.ID, custodySubnetCount uint64) (map[uint64]bool
 	// Columns belonging to the same subnet are contiguous.
 	columnIndices := make(map[uint64]bool, custodySubnetCount*columnsPerSubnet)
 	for i := uint64(0); i < columnsPerSubnet; i++ {
-		for subnetId := range subnetIds {
+		for _, subnetId := range subnetIds {
 			columnIndex := dataColumnSidecarSubnetCount*i + subnetId
 			columnIndices[columnIndex] = true
 		}

beacon-chain/p2p/BUILD.bazel

@@ -145,6 +145,7 @@ go_test(
         "//beacon-chain/blockchain/testing:go_default_library",
         "//beacon-chain/cache:go_default_library",
         "//beacon-chain/core/helpers:go_default_library",
+        "//beacon-chain/core/peerdas:go_default_library",
         "//beacon-chain/core/signing:go_default_library",
         "//beacon-chain/db/testing:go_default_library",
         "//beacon-chain/p2p/encoder:go_default_library",
@@ -162,6 +163,7 @@ go_test(
         "//container/leaky-bucket:go_default_library",
         "//crypto/ecdsa:go_default_library",
         "//crypto/hash:go_default_library",
+        "//crypto/rand:go_default_library",
         "//encoding/bytesutil:go_default_library",
         "//network:go_default_library",
         "//network/forks:go_default_library",

beacon-chain/p2p/subnets.go

@@ -220,17 +220,12 @@ func initializePersistentColumnSubnets(id enode.ID) error {
 	if ok && expTime.After(time.Now()) {
 		return nil
 	}
-	subsMap, err := peerdas.CustodyColumnSubnets(id, params.BeaconConfig().CustodyRequirement)
+	subnetsId, err := peerdas.CustodyColumnSubnets(id, params.BeaconConfig().CustodyRequirement)
 	if err != nil {
 		return err
 	}
 
-	subs := make([]uint64, 0, len(subsMap))
-	for sub := range subsMap {
-		subs = append(subs, sub)
-	}
-
-	cache.ColumnSubnetIDs.AddColumnSubnets(subs)
+	cache.ColumnSubnetIDs.AddColumnSubnets(subnetsId)
 	return nil
 }
 

beacon-chain/p2p/utils.go

@@ -6,6 +6,7 @@ import (
 	"crypto/rand"
 	"encoding/base64"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"net"
 	"os"
@@ -20,7 +21,10 @@ import (
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/go-bitfield"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
+	"github.com/prysmaticlabs/prysm/v5/cmd/beacon-chain/flags"
 	"github.com/prysmaticlabs/prysm/v5/config/features"
+	"github.com/prysmaticlabs/prysm/v5/config/params"
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/wrapper"
 	ecdsaprysm "github.com/prysmaticlabs/prysm/v5/crypto/ecdsa"
 	"github.com/prysmaticlabs/prysm/v5/io/file"
@@ -30,10 +34,13 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
-const keyPath = "network-keys"
-const metaDataPath = "metaData"
+const (
+	keyPath                  = "network-keys"
+	custodyColumnSubnetsPath = "custodyColumnsSubnets.json"
+	metaDataPath             = "metaData"
 
-const dialTimeout = 1 * time.Second
+	dialTimeout = 1 * time.Second
+)
 
 // SerializeENR takes the enr record in its key-value form and serializes it.
 func SerializeENR(record *enr.Record) (string, error) {
@@ -48,22 +55,224 @@ func SerializeENR(record *enr.Record) (string, error) {
 	return enrString, nil
 }
 
-// Determines a private key for p2p networking from the p2p service's
+// randomPrivKeyWithSubnets generates a random private key which, when derived into a node ID, matches expectedSubnets.
+// This is done by brute forcing the generation of a private key until it matches the desired subnets.
+// TODO: Run multiple goroutines to speed up the process.
+func randomPrivKeyWithSubnets(expectedSubnets map[uint64]bool) (crypto.PrivKey, uint64, time.Duration, error) {
+	// Get the current time.
+	start := time.Now()
+
+mainLoop:
+	for i := uint64(1); ; /* No exit condition */ i++ {
+		// Get the subnets count.
+		expectedSubnetsCount := len(expectedSubnets)
+
+		// Generate a random keys pair
+		privKey, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+		if err != nil {
+			return nil, 0, time.Duration(0), errors.Wrap(err, "generate SECP256K1 key")
+		}
+
+		ecdsaPrivKey, err := ecdsaprysm.ConvertFromInterfacePrivKey(privKey)
+		if err != nil {
+			return nil, 0, time.Duration(0), errors.Wrap(err, "convert from interface private key")
+		}
+
+		// Compute the node ID from the public key.
+		nodeID := enode.PubkeyToIDV4(&ecdsaPrivKey.PublicKey)
+
+		// Retrieve the custody column subnets of the node.
+		actualSubnets, err := peerdas.CustodyColumnSubnets(nodeID, uint64(expectedSubnetsCount))
+		if err != nil {
+			return nil, 0, time.Duration(0), errors.Wrap(err, "custody column subnets")
+		}
+
+		// Safe check, just in case.
+		actualSubnetsCount := len(actualSubnets)
+		if actualSubnetsCount != expectedSubnetsCount {
+			return nil, 0, time.Duration(0), errors.Errorf("mismatch counts of custody subnets. Actual %d - Required %d", actualSubnetsCount, expectedSubnetsCount)
+		}
+
+		// Check if the expected subnets are the same as the actual subnets.
+		for _, subnet := range actualSubnets {
+			if !expectedSubnets[subnet] {
+				// At least one subnet does not match, so we need to generate a new key.
+				continue mainLoop
+			}
+		}
+
+		// It's a match, return the private key.
+		return privKey, i, time.Since(start), nil
+	}
+}
+
+// privateKeyWithConstraint reads the subnets from a file and generates a private key that matches the subnets.
+func privateKeyWithConstraint(subnetsPath string) (crypto.PrivKey, error) {
+	// Read the subnets from the file.
+	data, err := file.ReadFileAsBytes(subnetsPath)
+	if err != nil {
+		return nil, errors.Wrapf(err, "read file %s", subnetsPath)
+	}
+
+	var storedSubnets []uint64
+	if err := json.Unmarshal(data, &storedSubnets); err != nil {
+		return nil, errors.Wrapf(err, "unmarshal subnets %s", subnetsPath)
+	}
+
+	storedSubnetsCount := uint64(len(storedSubnets))
+
+	// Retrieve the subnets to custody.
+	custodySubnetsCount := params.BeaconConfig().CustodyRequirement
+	if flags.Get().SubscribeToAllSubnets {
+		custodySubnetsCount = params.BeaconConfig().DataColumnSidecarSubnetCount
+	}
+
+	// Check our subnets count is not greater than the subnet count in the file.
+	// Such a case is possible if the number of subnets increased after the file was created.
+	// This is possible only within a new release. If this happens, we should implement a modification
+	// of the file. At the moment, we raise an error.
+	if custodySubnetsCount > storedSubnetsCount {
+		return nil, errors.Errorf(
+			"subnets count in the file %s (%d) is less than the current subnets count (%d)",
+			subnetsPath,
+			storedSubnetsCount,
+			custodySubnetsCount,
+		)
+	}
+
+	subnetsMap := make(map[uint64]bool, custodySubnetsCount)
+	custodySubnetsMap := make(map[uint64]bool, len(storedSubnets))
+
+	for i, subnet := range storedSubnets {
+		subnetsMap[subnet] = true
+		if uint64(i) < custodySubnetsCount {
+			custodySubnetsMap[subnet] = true
+		}
+	}
+
+	if len(subnetsMap) != len(storedSubnets) {
+		return nil, errors.Errorf("duplicated subnets found in the file %s", subnetsPath)
+	}
+
+	// Generate a private key that matches the subnets.
+	privKey, iterations, duration, err := randomPrivKeyWithSubnets(custodySubnetsMap)
+	log.WithFields(logrus.Fields{
+		"iterations": iterations,
+		"duration":   duration,
+	}).Info("Generated P2P private key")
+
+	return privKey, err
+}
+
+// privateKeyWithoutConstraint generates a private key, computes the subnets and stores them in a file.
+func privateKeyWithoutConstraint(subnetsPath string) (crypto.PrivKey, error) {
+	// Get the total number of subnets.
+	subnetCount := params.BeaconConfig().DataColumnSidecarSubnetCount
+
+	// Generate the private key.
+	privKey, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	if err != nil {
+		return nil, errors.Wrap(err, "generate SECP256K1 key")
+	}
+
+	convertedKey, err := ecdsaprysm.ConvertFromInterfacePrivKey(privKey)
+	if err != nil {
+		return nil, errors.Wrap(err, "convert from interface private key")
+	}
+
+	// Compute the node ID from the public key.
+	nodeID := enode.PubkeyToIDV4(&convertedKey.PublicKey)
+
+	// Retrieve the custody column subnets of the node.
+	subnets, err := peerdas.CustodyColumnSubnets(nodeID, subnetCount)
+	if err != nil {
+		return nil, errors.Wrap(err, "custody column subnets")
+	}
+
+	// Store the subnets in a file.
+	data, err := json.Marshal(subnets)
+	if err != nil {
+		return nil, errors.Wrap(err, "marshal subnets")
+	}
+
+	if err := file.WriteFile(subnetsPath, data); err != nil {
+		return nil, errors.Wrap(err, "write file")
+	}
+
+	return privKey, nil
+}
+
+// storePrivateKey stores a private key to a file.
+func storePrivateKey(privKey crypto.PrivKey, destFilePath string) error {
+	// Get the raw bytes of the private key.
+	rawbytes, err := privKey.Raw()
+	if err != nil {
+		return errors.Wrap(err, "raw")
+	}
+
+	// Encode the raw bytes to hex.
+	dst := make([]byte, hex.EncodedLen(len(rawbytes)))
+	hex.Encode(dst, rawbytes)
+
+	if err := file.WriteFile(destFilePath, dst); err != nil {
+		return errors.Wrapf(err, "write file: %s", destFilePath)
+	}
+
+	return err
+}
+
+// randomPrivKey generates a random private key.
+func randomPrivKey(datadir string) (crypto.PrivKey, error) {
+	if features.Get().EnablePeerDAS {
+		// Check if the file containing the custody column subnets exists.
+		subnetsPath := path.Join(datadir, custodyColumnSubnetsPath)
+		exists, err := file.Exists(subnetsPath, file.Regular)
+		if err != nil {
+			return nil, errors.Wrap(err, "exists")
+		}
+
+		// If the file does not exist, generate a new private key, compute the subnets and store them.
+		if !exists {
+			priv, err := privateKeyWithoutConstraint(subnetsPath)
+			if err != nil {
+				return nil, errors.Wrap(err, "generate private without constraint")
+			}
+
+			return priv, nil
+		}
+
+		// If the file exists, read the subnets and generate a new private key.
+		priv, err := privateKeyWithConstraint(subnetsPath)
+		if err != nil {
+			return nil, errors.Wrap(err, "generate private key with constraint for PeerDAS")
+		}
+
+		return priv, nil
+	}
+
+	privKey, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	if err != nil {
+		return nil, errors.Wrap(err, "generate SECP256K1 key")
+	}
+
+	return privKey, err
+}
+
+// privKey determines a private key for p2p networking from the p2p service's
 // configuration struct. If no key is found, it generates a new one.
 func privKey(cfg *Config) (*ecdsa.PrivateKey, error) {
 	defaultKeyPath := path.Join(cfg.DataDir, keyPath)
 	privateKeyPath := cfg.PrivateKey
 
-	// PrivateKey cli flag takes highest precedence.
+	// PrivateKey CLI flag takes highest precedence.
 	if privateKeyPath != "" {
 		return privKeyFromFile(cfg.PrivateKey)
 	}
 
 	// Default keys have the next highest precedence, if they exist.
-	_, err := os.Stat(defaultKeyPath)
-	defaultKeysExist := !os.IsNotExist(err)
-	if err != nil && defaultKeysExist {
-		return nil, err
+	defaultKeysExist, err := file.Exists(defaultKeyPath, file.Regular)
+	if err != nil {
+		return nil, errors.Wrap(err, "exists")
 	}
 
 	if defaultKeysExist {
@@ -71,34 +280,32 @@ func privKey(cfg *Config) (*ecdsa.PrivateKey, error) {
 		return privKeyFromFile(defaultKeyPath)
 	}
 
-	// There are no keys on the filesystem, so we need to generate one.
-	priv, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	// Generate a new (possibly contrained) random private key.
+	priv, err := randomPrivKey(cfg.DataDir)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "random private key")
 	}
 
-	// If the StaticPeerID flag is not set and if peerDAS is not enabled, return the private key.
-	if !(cfg.StaticPeerID || features.Get().EnablePeerDAS) {
+	// If the StaticPeerID flag is not set, return the private key.
+	if !cfg.StaticPeerID {
 		return ecdsaprysm.ConvertFromInterfacePrivKey(priv)
 	}
 
 	// Save the generated key as the default key, so that it will be used by
 	// default on the next node start.
-	rawbytes, err := priv.Raw()
-	if err != nil {
-		return nil, err
+	log.WithField("file", defaultKeyPath).Info("Wrote network key to")
+	if err := storePrivateKey(priv, defaultKeyPath); err != nil {
+		return nil, errors.Wrap(err, "store private key")
 	}
 
-	dst := make([]byte, hex.EncodedLen(len(rawbytes)))
-	hex.Encode(dst, rawbytes)
-	if err := file.WriteFile(defaultKeyPath, dst); err != nil {
-		return nil, err
-	}
-
-	log.WithField("path", defaultKeyPath).Info("Wrote network key to file")
 	// Read the key from the defaultKeyPath file just written
 	// for the strongest guarantee that the next start will be the same as this one.
-	return privKeyFromFile(defaultKeyPath)
+	privKey, err := privKeyFromFile(defaultKeyPath)
+	if err != nil {
+		return nil, errors.Wrap(err, "private key from file")
+	}
+
+	return privKey, nil
 }
 
 // Retrieves a p2p networking private key from a file path.

beacon-chain/p2p/utils_test.go

@@ -6,12 +6,110 @@ import (
 
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/peerdas"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
+	ecdsaprysm "github.com/prysmaticlabs/prysm/v5/crypto/ecdsa"
+	"github.com/prysmaticlabs/prysm/v5/crypto/rand"
 	"github.com/prysmaticlabs/prysm/v5/testing/assert"
 	"github.com/prysmaticlabs/prysm/v5/testing/require"
 	logTest "github.com/sirupsen/logrus/hooks/test"
 )
 
+// generateRandomSubnets generates a set of `count` random subnets.
+func generateRandomSubnets(requestedCount, totalSubnetsCount uint64) map[uint64]bool {
+	// Populate all the subnets.
+	subnets := make(map[uint64]bool, totalSubnetsCount)
+	for i := uint64(0); i < totalSubnetsCount; i++ {
+		subnets[i] = true
+	}
+
+	// Get a random generator.
+	randGen := rand.NewGenerator()
+
+	// Randomly delete subnets until we have the desired count.
+	for uint64(len(subnets)) > requestedCount {
+		// Get a random subnet.
+		subnet := randGen.Uint64() % totalSubnetsCount
+
+		// Delete the subnet.
+		delete(subnets, subnet)
+	}
+
+	return subnets
+}
+
+func TestRandomPrivKeyWithConstraint(t *testing.T) {
+	// Get the total number of subnets.
+	totalSubnetsCount := params.BeaconConfig().DataColumnSidecarSubnetCount
+
+	// We generate only tests for a low and high number of subnets to minimize computation, as explained here:
+	// https://hackmd.io/@6-HLeMXARN2tdFLKKcqrxw/BJVSxU7VC
+	testCases := []struct {
+		name                 string
+		expectedSubnetsCount uint64
+		expectedError        bool
+	}{
+		{
+			name:                 "0 subnet - n subnets",
+			expectedSubnetsCount: 0,
+		},
+		{
+			name:                 "1 subnet - n-1 subnets",
+			expectedSubnetsCount: 1,
+		},
+		{
+			name:                 "2 subnets - n-2 subnets",
+			expectedSubnetsCount: 2,
+		},
+		{
+			name:                 "3 subnets - n-3 subnets",
+			expectedSubnetsCount: 3,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			expectedSubnetsList := []map[uint64]bool{
+				generateRandomSubnets(tc.expectedSubnetsCount, totalSubnetsCount),
+				generateRandomSubnets(totalSubnetsCount-tc.expectedSubnetsCount, totalSubnetsCount),
+			}
+
+			for _, expectedSubnets := range expectedSubnetsList {
+				// Determine the number of expected subnets.
+				expectedSubnetsCount := uint64(len(expectedSubnets))
+
+				// Determine the private key that matches the expected subnets.
+				privateKey, iterationsCount, _, err := randomPrivKeyWithSubnets(expectedSubnets)
+				require.NoError(t, err)
+
+				// Sanity check the number of iterations.
+				assert.Equal(t, true, iterationsCount > 0)
+
+				// Compute the node ID from the public key.
+				ecdsaPrivKey, err := ecdsaprysm.ConvertFromInterfacePrivKey(privateKey)
+				require.NoError(t, err)
+
+				nodeID := enode.PubkeyToIDV4(&ecdsaPrivKey.PublicKey)
+
+				// Retrieve the subnets from the node ID.
+				actualSubnets, err := peerdas.CustodyColumnSubnets(nodeID, expectedSubnetsCount)
+				require.NoError(t, err)
+
+				// Determine the number of actual subnets.
+				actualSubnetsCounts := uint64(len(actualSubnets))
+
+				// Check the count of the actual subnets against the expected subnets.
+				assert.Equal(t, expectedSubnetsCount, actualSubnetsCounts)
+
+				// Check the actual subnets against the expected subnets.
+				for _, subnet := range actualSubnets {
+					assert.Equal(t, true, expectedSubnets[subnet])
+				}
+			}
+		})
+	}
+}
+
 // Test `verifyConnectivity` function by trying to connect to google.com (successfully)
 // and then by connecting to an unreachable IP and ensuring that a log is emitted
 func TestVerifyConnectivity(t *testing.T) {