From 7eba8da9d2782eb141cdc7031428b6a208906a94 Mon Sep 17 00:00:00 2001
From: Nishant Das <nishdas93@gmail.com>
Date: Tue, 17 Sep 2019 03:39:16 +0530
Subject: [PATCH] Save Network Keys in Data Directory (#3488)

* change marshalling

* add networkkeys

* gaz

* fix test

* add new function

* resolve comments, rename to datadir
---
 beacon-chain/node/node.go          |  1 +
 beacon-chain/p2p/config.go         |  1 +
 beacon-chain/p2p/discovery_test.go | 17 ++++++++++++-
 beacon-chain/p2p/utils.go          | 41 +++++++++++++++++++++++-------
 4 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/beacon-chain/node/node.go b/beacon-chain/node/node.go
index c6c108753..165e11b28 100644
--- a/beacon-chain/node/node.go
+++ b/beacon-chain/node/node.go
@@ -228,6 +228,7 @@ func (b *BeaconNode) registerP2P(ctx *cli.Context) error {
 		StaticPeers:       sliceutil.SplitCommaSeparated(ctx.GlobalStringSlice(cmd.StaticPeers.Name)),
 		BootstrapNodeAddr: bootnodeENR,
 		RelayNodeAddr:     ctx.GlobalString(cmd.RelayNode.Name),
+		DataDir:           ctx.GlobalString(cmd.DataDirFlag.Name),
 		HostAddress:       ctx.GlobalString(cmd.P2PHost.Name),
 		PrivateKey:        ctx.GlobalString(cmd.P2PPrivKey.Name),
 		TCPPort:           ctx.GlobalUint(cmd.P2PTCPPort.Name),
diff --git a/beacon-chain/p2p/config.go b/beacon-chain/p2p/config.go
index 977374f55..d2fb3d761 100644
--- a/beacon-chain/p2p/config.go
+++ b/beacon-chain/p2p/config.go
@@ -9,6 +9,7 @@ type Config struct {
 	RelayNodeAddr     string
 	HostAddress       string
 	PrivateKey        string
+	DataDir           string
 	TCPPort           uint
 	UDPPort           uint
 	MaxPeers          uint
diff --git a/beacon-chain/p2p/discovery_test.go b/beacon-chain/p2p/discovery_test.go
index 71f659bd0..acbbe6538 100644
--- a/beacon-chain/p2p/discovery_test.go
+++ b/beacon-chain/p2p/discovery_test.go
@@ -3,7 +3,11 @@ package p2p
 import (
 	"crypto/ecdsa"
 	"fmt"
+	"math/rand"
 	"net"
+	"os"
+	"path"
+	"strconv"
 	"testing"
 	"time"
 
@@ -17,13 +21,24 @@ import (
 
 var discoveryWaitTime = 1 * time.Second
 
+func init() {
+	rand.Seed(time.Now().Unix())
+}
+
 func createAddrAndPrivKey(t *testing.T) (net.IP, *ecdsa.PrivateKey) {
 	ip, err := iputils.ExternalIPv4()
 	if err != nil {
 		t.Fatalf("Could not get ip: %v", err)
 	}
 	ipAddr := net.ParseIP(ip)
-	pkey, err := privKey(&Config{Encoding: "ssz"})
+	temp := testutil.TempDir()
+	randNum := rand.Int()
+	tempPath := path.Join(temp, strconv.Itoa(randNum))
+	err = os.Mkdir(tempPath, 0700)
+	if 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)
 	}
diff --git a/beacon-chain/p2p/utils.go b/beacon-chain/p2p/utils.go
index 67ca236fd..694821213 100644
--- a/beacon-chain/p2p/utils.go
+++ b/beacon-chain/p2p/utils.go
@@ -6,6 +6,8 @@ import (
 	"encoding/hex"
 	"io/ioutil"
 	"net"
+	"os"
+	"path"
 
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/libp2p/go-libp2p-core/crypto"
@@ -13,6 +15,8 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/iputils"
 )
 
+const keyPath = "network-keys"
+
 func convertFromInterfacePrivKey(privkey crypto.PrivKey) *ecdsa.PrivateKey {
 	typeAssertedKey := (*ecdsa.PrivateKey)((*btcec.PrivateKey)(privkey.(*crypto.Secp256k1PrivateKey)))
 	return typeAssertedKey
@@ -23,26 +27,46 @@ func convertToInterfacePrivkey(privkey *ecdsa.PrivateKey) crypto.PrivKey {
 	return typeAssertedKey
 }
 
-func convertFromInterfacePubKey(pubkey crypto.PubKey) *ecdsa.PublicKey {
-	typeAssertedKey := (*ecdsa.PublicKey)((*btcec.PublicKey)(pubkey.(*crypto.Secp256k1PublicKey)))
-	return typeAssertedKey
-}
-
 func convertToInterfacePubkey(pubkey *ecdsa.PublicKey) crypto.PubKey {
 	typeAssertedKey := crypto.PubKey((*crypto.Secp256k1PublicKey)((*btcec.PublicKey)(pubkey)))
 	return typeAssertedKey
 }
 
 func privKey(cfg *Config) (*ecdsa.PrivateKey, error) {
-	if cfg.PrivateKey == "" {
+	defaultKeyPath := path.Join(cfg.DataDir, keyPath)
+	privateKeyPath := cfg.PrivateKey
+
+	_, err := os.Stat(defaultKeyPath)
+	defaultKeysExist := !os.IsNotExist(err)
+	if err != nil && defaultKeysExist {
+		return nil, err
+	}
+
+	if privateKeyPath == "" && !defaultKeysExist {
 		priv, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
 		if err != nil {
 			return nil, err
 		}
+		rawbytes, err := priv.Raw()
+		if err != nil {
+			return nil, err
+		}
+		dst := make([]byte, hex.EncodedLen(len(rawbytes)))
+		hex.Encode(dst, rawbytes)
+		if err = ioutil.WriteFile(defaultKeyPath, dst, 0600); err != nil {
+			return nil, err
+		}
 		convertedKey := convertFromInterfacePrivKey(priv)
 		return convertedKey, nil
 	}
-	src, err := ioutil.ReadFile(cfg.PrivateKey)
+	if defaultKeysExist && privateKeyPath == "" {
+		privateKeyPath = defaultKeyPath
+	}
+	return retrievePrivKeyFromFile(privateKeyPath)
+}
+
+func retrievePrivKeyFromFile(path string) (*ecdsa.PrivateKey, error) {
+	src, err := ioutil.ReadFile(path)
 	if err != nil {
 		log.WithError(err).Error("Error reading private key from file")
 		return nil, err
@@ -56,8 +80,7 @@ func privKey(cfg *Config) (*ecdsa.PrivateKey, error) {
 	if err != nil {
 		return nil, err
 	}
-	priv := (*ecdsa.PrivateKey)((*btcec.PrivateKey)(unmarshalledKey.(*crypto.Secp256k1PrivateKey)))
-	return priv, nil
+	return convertFromInterfacePrivKey(unmarshalledKey), nil
 }
 
 func ipAddr(cfg *Config) net.IP {