Merge pull request #191 from SebastianJ/feature-generate-multi-bls

Add support for generating multiple bls keys for a given shard and node

cmd/subcommands/keys.go

@@ -33,6 +33,8 @@ var (
 	passphraseFilePath     string
 	passphrase             string
 	blsFilePath            string
+	blsShardID             uint32
+	blsCount               uint32
 	ppPrompt               = fmt.Sprintf(
 		"prompt for passphrase, otherwise use default passphrase: \"`%s`\"", c.DefaultPassphrase,
 	)
@@ -313,7 +315,13 @@ func keysSub() []*cobra.Command {
 			if err != nil {
 				return err
 			}
-			return keys.GenBlsKeys(passphrase, blsFilePath)
+
+			blsKey := &keys.BlsKey{
+				Passphrase: passphrase,
+				FilePath:   blsFilePath,
+			}
+
+			return keys.GenBlsKeys(blsKey)
 		},
 	}
 	cmdGenerateBlsKey.Flags().StringVar(&blsFilePath, "bls-file-path", "",
@@ -321,6 +329,41 @@ func keysSub() []*cobra.Command {
 	cmdGenerateBlsKey.Flags().BoolVar(&userProvidesPassphrase, "passphrase", false, ppPrompt)
 	cmdGenerateBlsKey.Flags().StringVar(&passphraseFilePath, "passphrase-file", "", "path to a file containing the passphrase")
 
+	cmdGenerateMultiBlsKeys := &cobra.Command{
+		Use:   "generate-bls-keys",
+		Short: "Generates multiple bls keys for a given shard network configuration and then encrypts and saves the private key with a requested passphrase",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			blsKeys := []*keys.BlsKey{}
+
+			for i := uint32(0); i < blsCount; i++ {
+				keyFilePath := blsFilePath
+				if blsFilePath != "" {
+					fmt.Printf("Enter absolute path for key #%d:\n", i+1)
+					fmt.Scanln(&keyFilePath)
+				}
+
+				passphrase, err := getPassphraseWithConfirm()
+				if err != nil {
+					return err
+				}
+
+				blsKey := &keys.BlsKey{
+					Passphrase: passphrase,
+					FilePath:   keyFilePath,
+				}
+				blsKeys = append(blsKeys, blsKey)
+			}
+
+			return keys.GenMultiBlsKeys(blsKeys, node, blsCount, blsShardID)
+		},
+	}
+	cmdGenerateMultiBlsKeys.Flags().StringVar(&blsFilePath, "bls-file-path", "",
+		"absolute path of where to save encrypted bls private keys")
+	cmdGenerateMultiBlsKeys.Flags().BoolVar(&userProvidesPassphrase, "passphrase", false, ppPrompt)
+	cmdGenerateMultiBlsKeys.Flags().StringVar(&passphraseFilePath, "passphrase-file", "", "path to a file containing the passphrase")
+	cmdGenerateMultiBlsKeys.Flags().Uint32Var(&blsShardID, "shard", 0, "which shard to create bls keys for")
+	cmdGenerateMultiBlsKeys.Flags().Uint32Var(&blsCount, "count", 1, "how many bls keys to generate")
+
 	cmdRecoverBlsKey := &cobra.Command{
 		Use:   "recover-bls-key <ABSOLUTE_PATH_BLS_KEY>",
 		Short: "Recover bls keys from an encrypted bls key file",
@@ -363,7 +406,7 @@ func keysSub() []*cobra.Command {
 	}
 
 	return []*cobra.Command{cmdList, cmdLocation, cmdAdd, cmdRemove, cmdMnemonic, cmdRecoverMnemonic, cmdImportKS, cmdImportPK,
-		cmdExportKS, cmdExportPK, cmdGenerateBlsKey, cmdRecoverBlsKey, cmdSaveBlsKey, GetPublicBlsKey}
+		cmdExportKS, cmdExportPK, cmdGenerateBlsKey, cmdGenerateMultiBlsKeys, cmdRecoverBlsKey, cmdSaveBlsKey, GetPublicBlsKey}
 }
 
 func init() {

pkg/keys/bls.go

@@ -11,12 +11,15 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"math/big"
 	"os"
 	"path"
 	"strings"
 
 	ffiBls "github.com/harmony-one/bls/ffi/go/bls"
 	"github.com/harmony-one/go-sdk/pkg/common"
+	"github.com/harmony-one/go-sdk/pkg/sharding"
+	"github.com/harmony-one/go-sdk/pkg/validation"
 	"github.com/harmony-one/harmony/crypto/bls"
 	"github.com/harmony-one/harmony/crypto/hash"
 	"github.com/harmony-one/harmony/shard"
@@ -24,31 +27,67 @@ import (
 	"golang.org/x/crypto/ssh/terminal"
 )
 
-func GenBlsKeys(passphrase, filePath string) error {
-	privateKey := bls.RandPrivateKey()
-	publicKey := privateKey.GetPublicKey()
-	publicKeyHex := publicKey.SerializeToHexStr()
-	privateKeyHex := privateKey.SerializeToHexStr()
+//BlsKey - struct to represent bls key data
+type BlsKey struct {
+	PrivateKey     *ffiBls.SecretKey
+	PublicKey      *ffiBls.PublicKey
+	PublicKeyHex   string
+	PrivateKeyHex  string
+	Passphrase     string
+	FilePath       string
+	ShardPublicKey *shard.BlsPublicKey
+}
 
-	if filePath == "" {
-		cwd, _ := os.Getwd()
-		filePath = fmt.Sprintf("%s/%s.key", cwd, publicKeyHex)
+//Initialize - initialize a bls key and assign a random private bls key if not already done
+func (blsKey *BlsKey) Initialize() {
+	if blsKey.PrivateKey == nil {
+		blsKey.PrivateKey = bls.RandPrivateKey()
+		blsKey.PrivateKeyHex = blsKey.PrivateKey.SerializeToHexStr()
+		blsKey.PublicKey = blsKey.PrivateKey.GetPublicKey()
+		blsKey.PublicKeyHex = blsKey.PublicKey.SerializeToHexStr()
 	}
-	if !path.IsAbs(filePath) {
-		return common.ErrNotAbsPath
-	}
-	encryptedPrivateKeyStr, err := encrypt([]byte(privateKeyHex), passphrase)
+}
+
+//Reset - resets the currently assigned private and public key fields
+func (blsKey *BlsKey) Reset() {
+	blsKey.PrivateKey = nil
+	blsKey.PrivateKeyHex = ""
+	blsKey.PublicKey = nil
+	blsKey.PublicKeyHex = ""
+}
+
+// GenBlsKeys - generate a random bls key using the supplied passphrase, write it to disk at the given filePath
+func GenBlsKeys(blsKey *BlsKey) error {
+	blsKey.Initialize()
+	out, err := writeBlsKeyToFile(blsKey)
 	if err != nil {
 		return err
 	}
-	err = writeToFile(filePath, encryptedPrivateKeyStr)
+	fmt.Println(common.JSONPrettyFormat(out))
+	return nil
+}
+
+// GenMultiBlsKeys - generate multiple BLS keys for a given shard and node/network
+func GenMultiBlsKeys(blsKeys []*BlsKey, node string, count uint32, shardID uint32) error {
+	blsKeys, _, err := generateMultipleBlsKeys(blsKeys, node, count, shardID)
 	if err != nil {
 		return err
 	}
-	out := fmt.Sprintf(`
-{"public-key" : "%s", "private-key" : "%s", "encrypted-private-key-path" : "%s"}`,
-		publicKeyHex, privateKeyHex, filePath)
-	fmt.Println(common.JSONPrettyFormat(out))
+
+	outputs := []string{}
+	for _, blsKey := range blsKeys {
+		out, err := writeBlsKeyToFile(blsKey)
+		if err != nil {
+			return err
+		}
+
+		outputs = append(outputs, out)
+	}
+
+	if len(outputs) > 0 {
+		fmt.Println(common.JSONPrettyFormat(fmt.Sprintf("[%s]", strings.Join(outputs[:], ","))))
+	}
+
 	return nil
 }
 
@@ -191,6 +230,29 @@ func getBlsKey(privateKeyHex string) (*ffiBls.SecretKey, error) {
 	return privateKey, nil
 }
 
+func writeBlsKeyToFile(blsKey *BlsKey) (string, error) {
+	if blsKey.FilePath == "" {
+		cwd, _ := os.Getwd()
+		blsKey.FilePath = fmt.Sprintf("%s/%s.key", cwd, blsKey.PublicKeyHex)
+	}
+	if !path.IsAbs(blsKey.FilePath) {
+		return "", common.ErrNotAbsPath
+	}
+	encryptedPrivateKeyStr, err := encrypt([]byte(blsKey.PrivateKeyHex), blsKey.Passphrase)
+	if err != nil {
+		return "", err
+	}
+	err = writeToFile(blsKey.FilePath, encryptedPrivateKeyStr)
+	if err != nil {
+		return "", err
+	}
+	out := fmt.Sprintf(`
+{"public-key" : "%s", "private-key" : "%s", "encrypted-private-key-path" : "%s"}`,
+		blsKey.PublicKeyHex, blsKey.PrivateKeyHex, blsKey.FilePath)
+
+	return out, nil
+}
+
 func writeToFile(filename string, data string) error {
 	file, err := os.Create(filename)
 	if err != nil {
@@ -257,3 +319,40 @@ func decryptRaw(data []byte, passphrase string) ([]byte, error) {
 	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
 	return plaintext, err
 }
+
+func generateMultipleBlsKeys(blsKeys []*BlsKey, node string, count uint32, shardID uint32) ([]*BlsKey, int, error) {
+	shardingStructure, err := sharding.Structure(node)
+	if err != nil {
+		return blsKeys, -1, err
+	}
+	shardCount := len(shardingStructure)
+
+	if !validation.ValidShardID(shardID, uint32(shardCount)) {
+		return blsKeys, shardCount, fmt.Errorf("node %s only supports a total of %d shards - supplied shard id %d isn't valid", node, shardCount, shardID)
+	}
+
+	for _, blsKey := range blsKeys {
+		for {
+			blsKey.Initialize()
+			shardPubKey := new(shard.BlsPublicKey)
+			if err = shardPubKey.FromLibBLSPublicKey(blsKey.PublicKey); err != nil {
+				return blsKeys, shardCount, err
+			}
+
+			if blsKeyMatchesShardID(shardPubKey, shardID, shardCount) {
+				blsKey.ShardPublicKey = shardPubKey
+				break
+			} else {
+				blsKey.Reset()
+			}
+		}
+	}
+
+	return blsKeys, shardCount, nil
+}
+
+func blsKeyMatchesShardID(pubKey *shard.BlsPublicKey, shardID uint32, shardCount int) bool {
+	bigShardCount := big.NewInt(int64(shardCount))
+	resolvedShardID := int(new(big.Int).Mod(pubKey.Big(), bigShardCount).Int64())
+	return int(shardID) == resolvedShardID
+}

pkg/keys/bls_test.go

@@ -0,0 +1,91 @@
+package keys
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestBlsKeyGeneration(t *testing.T) {
+	valid := true
+	passphrase := ""
+	folderPath := "TestBlsKeyGeneration"
+	absFolderPath, err := filepath.Abs(fmt.Sprintf("./%s", folderPath))
+	if err != nil {
+		t.Errorf("TestBlsKeyGeneration - failed to convert relative path to absolute path")
+	}
+
+	absFilePath := fmt.Sprintf("%s/TestBlsKeyGeneration.key", absFolderPath)
+
+	err = os.MkdirAll(absFolderPath, 0755)
+	if err != nil {
+		t.Errorf("TestBlsKeyGeneration - failed to make test key folder")
+	}
+
+	if err := GenBlsKeys(&BlsKey{Passphrase: passphrase, FilePath: absFilePath}); err != nil {
+		t.Errorf("TestBlsKeyGeneration - failed to generate bls key using passphrase %s and path %s", passphrase, absFilePath)
+	}
+
+	if _, err = os.Stat(absFilePath); err != nil {
+		t.Errorf("TestBlsKeyGeneration - failed to check if file %s exists", absFilePath)
+	}
+
+	valid = !os.IsNotExist(err)
+
+	if !valid {
+		t.Errorf("GenBlsKeys - failed to generate a bls key using passphrase %s", "")
+	}
+
+	os.RemoveAll(absFolderPath)
+}
+
+func TestMultiBlsKeyGeneration(t *testing.T) {
+	tests := []struct {
+		node     string
+		count    uint32
+		shardID  uint32
+		filePath string
+		expected bool
+	}{
+		{node: "https://api.s0.os.hmny.io", count: 3, shardID: 0, expected: true},
+		{node: "https://api.s0.ps.hmny.io", count: 3, shardID: 0, expected: true},
+		{node: "https://api.s0.stn.hmny.io", count: 3, shardID: 0, expected: true},
+		{node: "https://api.s0.b.hmny.io", count: 3, shardID: 0, expected: true},
+		{node: "https://api.s0.t.hmny.io", count: 3, shardID: 0, expected: true},
+
+		{node: "https://api.s0.os.hmny.io", count: 3, shardID: 4, expected: false},
+		{node: "https://api.s0.ps.hmny.io", count: 3, shardID: 4, expected: false},
+		{node: "https://api.s0.stn.hmny.io", count: 3, shardID: 4, expected: false},
+		{node: "https://api.s0.b.hmny.io", count: 3, shardID: 4, expected: false},
+		{node: "https://api.s0.t.hmny.io", count: 3, shardID: 4, expected: false},
+	}
+
+	for _, test := range tests {
+		valid := false
+		blsKeys := []*BlsKey{}
+		for i := uint32(0); i < test.count; i++ {
+			blsKeys = append(blsKeys, &BlsKey{Passphrase: "", FilePath: ""})
+		}
+
+		blsKeys, shardCount, err := generateMultipleBlsKeys(blsKeys, test.node, test.count, test.shardID)
+		if err != nil {
+			valid = false
+		}
+
+		successCount := 0
+
+		for _, blsKey := range blsKeys {
+			success := (blsKey.ShardPublicKey != nil && blsKeyMatchesShardID(blsKey.ShardPublicKey, test.shardID, shardCount))
+			if success {
+				successCount++
+			}
+		}
+
+		valid = (successCount == int(test.count))
+
+		if valid != test.expected {
+			t.Errorf("generateMultipleBlsKeys - failed to generate %d keys for shard %d using node %s", test.count, test.shardID, test.node)
+		}
+	}
+}