From eb6e2f569b1d2ae737222608d9d123457056f5af Mon Sep 17 00:00:00 2001
From: Leo Chen <leo@harmony.one>
Date: Mon, 28 Jan 2019 23:11:22 +0000
Subject: [PATCH] add private key save and load functions

add test functions

Signed-off-by: Leo Chen <leo@harmony.one>
---
 internal/utils/utils.go      | 41 ++++++++++++++++++++++++++++++++++++
 internal/utils/utils_test.go | 38 +++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/internal/utils/utils.go b/internal/utils/utils.go
index a50444048..1084bdcba 100644
--- a/internal/utils/utils.go
+++ b/internal/utils/utils.go
@@ -23,6 +23,11 @@ import (
 
 var lock sync.Mutex
 
+// PrivKeyStore is used to persist private key to/from file
+type PrivKeyStore struct {
+	Key string `json:"key"`
+}
+
 func init() {
 	bls.Init(bls.BLS12_381)
 }
@@ -183,3 +188,39 @@ func SavePrivateKey(key p2p_crypto.PrivKey) (string, error) {
 	}
 	return "", fmt.Errorf("key is nil")
 }
+
+// SaveKeyToFile save private key to keyfile
+func SaveKeyToFile(keyfile string, key p2p_crypto.PrivKey) (err error) {
+	str, err := SavePrivateKey(key)
+	if err != nil {
+		return
+	}
+
+	keyStruct := PrivKeyStore{Key: str}
+
+	err = Save(keyfile, &keyStruct)
+	return
+}
+
+// LoadKeyFromFile load private key from keyfile
+// If the private key is not loadable or no file, it will generate
+// a new random private key
+func LoadKeyFromFile(keyfile string) (key p2p_crypto.PrivKey, err error) {
+	var keyStruct PrivKeyStore
+	err = Load(keyfile, &keyStruct)
+	if err != nil {
+		log.Print("No priviate key can be loaded from file", "keyfile", keyfile)
+		log.Print("Using random private key")
+		key, _, err = GenKeyP2PRand()
+		if err != nil {
+			log.Panic("LoadKeyFromFile", "GenKeyP2PRand Error", err)
+		}
+		err = SaveKeyToFile(keyfile, key)
+		if err != nil {
+			log.Print("LoadKeyFromFile", "failed to save key to keyfile", err)
+		}
+		return key, nil
+	}
+	key, err = LoadPrivateKey(keyStruct.Key)
+	return key, err
+}
diff --git a/internal/utils/utils_test.go b/internal/utils/utils_test.go
index d7cea4454..ad0838126 100644
--- a/internal/utils/utils_test.go
+++ b/internal/utils/utils_test.go
@@ -3,6 +3,7 @@ package utils
 import (
 	"bytes"
 	"encoding/hex"
+	"os"
 	"testing"
 
 	"github.com/harmony-one/harmony/p2p"
@@ -124,3 +125,40 @@ func TestSaveLoadPrivateKey(t *testing.T) {
 		t.Errorf("got pk1: %v\n", b2)
 	}
 }
+
+func TestSaveLoadKeyFile(t *testing.T) {
+	filename := "/tmp/keystore"
+	nonexist := "/tmp/please_ignore_the_non-exist_file"
+
+	key, _, err := GenKeyP2PRand()
+	if err != nil {
+		t.Fatalf("failed to generate random p2p key: %v", err)
+	}
+
+	err = SaveKeyToFile(filename, key)
+	if err != nil {
+		t.Fatalf("failed to save key to file: %v", err)
+	}
+
+	key1, err := LoadKeyFromFile(filename)
+	if err != nil {
+		t.Fatalf("failed to load key from file (%s): %v", filename, err)
+	}
+
+	if !crypto.KeyEqual(key, key1) {
+		t.Fatalf("loaded key is not equal to the saved one")
+	}
+
+	key2, err := LoadKeyFromFile(nonexist)
+
+	if err != nil {
+		t.Fatalf("failed to load key from non-exist file: %v", err)
+	}
+
+	if crypto.KeyEqual(key1, key2) {
+		t.Fatalf("new random key can't equal to existing one, something is wrong!")
+	}
+
+	os.Remove(filename)
+	os.Remove(nonexist)
+}