[blockchain] change private key loading logic (#3621)

iotexproject · Sep 2, 2022 · ba84d70 · ba84d70
1 parent e5e9d11
commit ba84d70
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 50 deletions.
diff --git a/blockchain/config.go b/blockchain/config.go
@@ -36,7 +36,7 @@ type (
 		EVMNetworkID           uint32           `yaml:"evmNetworkID"`
 		Address                string           `yaml:"address"`
 		ProducerPrivKey        string           `yaml:"producerPrivKey"`
-		PrivKeyConfigFile      string           `yaml:"privKeyConfigFile"`
+		ProducerPrivKeySchema  string           `yaml:"producerPrivKeySchema"`
 		SignatureScheme        []string         `yaml:"signatureScheme"`
 		EmptyGenesis           bool             `yaml:"emptyGenesis"`
 		GravityChainDB         db.Config        `yaml:"gravityChainDB"`
@@ -137,39 +137,32 @@ func (cfg *Config) ProducerPrivateKey() crypto.PrivateKey {
 
 // SetProducerPrivKey set producer privKey by PrivKeyConfigFile info
 func (cfg *Config) SetProducerPrivKey() error {
-	if cfg.PrivKeyConfigFile == "" {
-		return nil
-	}
-
-	yaml, err := config.NewYAML(config.Expand(os.LookupEnv), config.File(cfg.PrivKeyConfigFile))
-	if err != nil {
-		return errors.Wrap(err, "failed to init private key config")
-	}
-	pc := &privKeyConfig{}
-	if err := yaml.Get(config.Root).Populate(pc); err != nil {
-		return errors.Wrap(err, "failed to unmarshal YAML config to privKeyConfig struct")
-	}
-
-	var loader privKeyLoader
-	switch pc.Method {
+	switch cfg.ProducerPrivKeySchema {
+	case "hex", "":
+		// do nothing
 	case "hashiCorpVault":
-		cli, err := newVaultClient(&pc.VaultConfig)
+		yaml, err := config.NewYAML(config.Expand(os.LookupEnv), config.File(cfg.ProducerPrivKey))
+		if err != nil {
+			return errors.Wrap(err, "failed to init private key config")
+		}
+		hcv := &hashiCorpVault{}
+		if err := yaml.Get(config.Root).Populate(hcv); err != nil {
+			return errors.Wrap(err, "failed to unmarshal YAML config to privKeyConfig struct")
+		}
+
+		loader, err := newVaultPrivKeyLoader(hcv)
 		if err != nil {
 			return errors.Wrap(err, "failed to new vault client")
 		}
-		loader = &vaultPrivKeyLoader{
-			cfg:         &pc.VaultConfig,
-			vaultClient: cli,
+		key, err := loader.load()
+		if err != nil {
+			return errors.Wrap(err, "failed to load producer private key")
 		}
+		cfg.ProducerPrivKey = key
 	default:
-		return errors.Wrap(ErrConfig, "invalid private key method")
+		return errors.Wrap(ErrConfig, "invalid private key schema")
 	}
 
-	key, err := loader.load()
-	if err != nil {
-		return errors.Wrap(err, "failed to load producer private key")
-	}
-	cfg.ProducerPrivKey = key
 	return nil
 }
 

diff --git a/blockchain/config_privatekey.go b/blockchain/config_privatekey.go
@@ -26,15 +26,6 @@ type (
 		Key     string `yaml:"key"`
 	}
 
-	privKeyConfig struct {
-		Method      string         `yaml:"method"`
-		VaultConfig hashiCorpVault `yaml:"hashiCorpVault"`
-	}
-
-	privKeyLoader interface {
-		load() (string, error)
-	}
-
 	vaultPrivKeyLoader struct {
 		cfg *hashiCorpVault
 		*vaultClient
@@ -73,7 +64,7 @@ func (l *vaultPrivKeyLoader) load() (string, error) {
 	return v, nil
 }
 
-func newVaultClient(cfg *hashiCorpVault) (*vaultClient, error) {
+func newVaultPrivKeyLoader(cfg *hashiCorpVault) (*vaultPrivKeyLoader, error) {
 	conf := api.DefaultConfig()
 	conf.Address = cfg.Address
 	conf.Timeout = defaultHTTPTimeout
@@ -83,7 +74,8 @@ func newVaultClient(cfg *hashiCorpVault) (*vaultClient, error) {
 	}
 	cli.SetToken(cfg.Token)
 
-	return &vaultClient{
-		cli: cli.Logical(),
+	return &vaultPrivKeyLoader{
+		vaultClient: &vaultClient{cli: cli.Logical()},
+		cfg:         cfg,
 	}, nil
 }
diff --git a/blockchain/config_privatekey_test.go b/blockchain/config_privatekey_test.go
@@ -19,12 +19,10 @@ import (
 
 const (
 	hashiCorpVaultTestCfg = `
-method: hashiCorpVault
-hashiCorpVault:
-    address: http://127.0.0.1:8200
-    token: secret/data/test
-    path: secret/data/test
-    key: my key
+address: http://127.0.0.1:8200
+token: secret/data/test
+path: secret/data/test
+key: my key
 `
 
 	vaultTestKey   = "my key"
@@ -47,8 +45,8 @@ func TestVault(t *testing.T) {
 		vaultClient: &vaultClient{reader},
 	}
 
-	t.Run("NewVaultClientSuccess", func(t *testing.T) {
-		_, err := newVaultClient(cfg)
+	t.Run("NewVaultPrivKeyLoaderSuccess", func(t *testing.T) {
+		_, err := newVaultPrivKeyLoader(cfg)
 		r.NoError(err)
 	})
 	t.Run("VaultSuccess", func(t *testing.T) {
@@ -111,14 +109,15 @@ func TestSetProducerPrivKey(t *testing.T) {
 		r.NoError(err)
 		r.Equal(key, cfg.ProducerPrivKey)
 	})
-	t.Run("PrivateConfigFileIsEmpty", func(t *testing.T) {
+	t.Run("PrivateConfigUnknownSchema", func(t *testing.T) {
 		cfg := DefaultConfig
 		tmp, err := os.CreateTemp("", testfile)
 		r.NoError(err)
 		defer os.Remove(tmp.Name())
-		cfg.PrivKeyConfigFile = tmp.Name()
+		cfg.ProducerPrivKey = tmp.Name()
+		cfg.ProducerPrivKeySchema = "unknown"
 		err = cfg.SetProducerPrivKey()
-		r.Contains(err.Error(), "invalid private key method")
+		r.Contains(err.Error(), "invalid private key schema")
 	})
 	t.Run("PrivateConfigFileHasHashiCorpVault", func(t *testing.T) {
 		cfg := DefaultConfig
@@ -130,7 +129,8 @@ func TestSetProducerPrivKey(t *testing.T) {
 		r.NoError(err)
 		err = tmp.Close()
 		r.NoError(err)
-		cfg.PrivKeyConfigFile = tmp.Name()
+		cfg.ProducerPrivKey = tmp.Name()
+		cfg.ProducerPrivKeySchema = "hashiCorpVault"
 		err = cfg.SetProducerPrivKey()
 		r.Contains(err.Error(), "dial tcp 127.0.0.1:8200: connect: connection refused")
 	})