From 91e84ffedfe09bad87c8f7190f7269df6ef82df5 Mon Sep 17 00:00:00 2001 From: codeman <68629395+aakejiang@users.noreply.github.com> Date: Thu, 21 Apr 2022 19:57:42 +0800 Subject: [PATCH] fix: not create user root if exists, optimize log & imports (#16554) Signed-off-by: kejiang Co-authored-by: kejiang --- internal/rootcoord/meta_table.go | 5 +-- internal/rootcoord/meta_table_test.go | 4 ++- internal/rootcoord/root_coord.go | 24 ++++++++----- internal/rootcoord/root_coord_test.go | 51 +++++++++++++++++++++++++++ internal/util/crypto/crypto_test.go | 27 ++++++++++++++ 5 files changed, 100 insertions(+), 11 deletions(-) diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index 34bad7fd3d..51a7afd0a8 100644 --- a/internal/rootcoord/meta_table.go +++ b/internal/rootcoord/meta_table.go @@ -18,6 +18,7 @@ package rootcoord import ( "bytes" + "encoding/json" "fmt" "path" "strconv" @@ -1360,7 +1361,7 @@ func (mt *MetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error { return fmt.Errorf("username is empty") } k := fmt.Sprintf("%s/%s", CredentialPrefix, credInfo.Username) - v, err := proto.Marshal(&internalpb.CredentialInfo{EncryptedPassword: credInfo.EncryptedPassword}) + v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: credInfo.EncryptedPassword}) if err != nil { log.Error("MetaTable marshal credential info fail", zap.String("key", k), zap.Error(err)) return fmt.Errorf("metaTable marshal credential info fail key:%s, err:%w", k, err) @@ -1386,7 +1387,7 @@ func (mt *MetaTable) getCredential(username string) (*internalpb.CredentialInfo, } credentialInfo := internalpb.CredentialInfo{} - err = proto.Unmarshal([]byte(v), &credentialInfo) + err = json.Unmarshal([]byte(v), &credentialInfo) if err != nil { return nil, fmt.Errorf("get credential unmarshal err:%w", err) } diff --git a/internal/rootcoord/meta_table_test.go b/internal/rootcoord/meta_table_test.go index b065995ae1..6f5d26ce27 100644 --- a/internal/rootcoord/meta_table_test.go +++ b/internal/rootcoord/meta_table_test.go @@ -1140,7 +1140,9 @@ func TestMetaTable(t *testing.T) { wg.Add(1) t.Run("add credential failed", func(t *testing.T) { defer wg.Done() - + mockTxnKV.save = func(key, value string) error { + return fmt.Errorf("save error") + } err = mt.AddCredential(&internalpb.CredentialInfo{Username: "x", EncryptedPassword: "a\xc5z"}) assert.NotNil(t, err) }) diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go index bd6da15c8a..642e6cd420 100644 --- a/internal/rootcoord/root_coord.go +++ b/internal/rootcoord/root_coord.go @@ -29,11 +29,6 @@ import ( "syscall" "time" - "github.com/milvus-io/milvus/internal/util" - "github.com/milvus-io/milvus/internal/util/crypto" - - "github.com/milvus-io/milvus/internal/util/dependency" - "github.com/golang/protobuf/proto" clientv3 "go.etcd.io/etcd/client/v3" "go.uber.org/zap" @@ -57,6 +52,9 @@ import ( "github.com/milvus-io/milvus/internal/proto/schemapb" "github.com/milvus-io/milvus/internal/tso" "github.com/milvus-io/milvus/internal/types" + "github.com/milvus-io/milvus/internal/util" + "github.com/milvus-io/milvus/internal/util/crypto" + "github.com/milvus-io/milvus/internal/util/dependency" "github.com/milvus-io/milvus/internal/util/metricsinfo" "github.com/milvus-io/milvus/internal/util/paramtable" "github.com/milvus-io/milvus/internal/util/retry" @@ -1122,13 +1120,12 @@ func (c *Core) Init() error { c.impTaskKv, c.CallImportService, ) + // init data - encryptedRootPassword, _ := crypto.PasswordEncrypt(util.DefaultRootPassword) - initError = c.MetaTable.AddCredential(&internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword}) + initError = c.initData() if initError != nil { return } - log.Debug("RootCoord init user root done") }) if initError != nil { log.Debug("RootCoord init error", zap.Error(initError)) @@ -1137,6 +1134,17 @@ func (c *Core) Init() error { return initError } +func (c *Core) initData() error { + credInfo, _ := c.MetaTable.getCredential(util.UserRoot) + if credInfo == nil { + log.Debug("RootCoord init user root") + encryptedRootPassword, _ := crypto.PasswordEncrypt(util.DefaultRootPassword) + err := c.MetaTable.AddCredential(&internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword}) + return err + } + return nil +} + func (c *Core) reSendDdMsg(ctx context.Context, force bool) error { if !force { flag, err := c.MetaTable.txn.Load(DDMsgSendPrefix) diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go index dfff3edf70..623cb45b7b 100644 --- a/internal/rootcoord/root_coord_test.go +++ b/internal/rootcoord/root_coord_test.go @@ -49,6 +49,7 @@ import ( "github.com/milvus-io/milvus/internal/proto/rootcoordpb" "github.com/milvus-io/milvus/internal/proto/schemapb" "github.com/milvus-io/milvus/internal/types" + "github.com/milvus-io/milvus/internal/util" "github.com/milvus-io/milvus/internal/util/dependency" "github.com/milvus-io/milvus/internal/util/etcd" "github.com/milvus-io/milvus/internal/util/funcutil" @@ -626,6 +627,56 @@ func TestRootCoordInit(t *testing.T) { assert.NoError(t, err) } +func TestRootCoordInitData(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + coreFactory := dependency.NewDefaultFactory(true) + Params.Init() + Params.RootCoordCfg.DmlChannelNum = TestDMLChannelNum + + etcdCli, err := etcd.GetEtcdClient(&Params.EtcdCfg) + assert.NoError(t, err) + defer etcdCli.Close() + + core, err := NewCore(ctx, coreFactory) + assert.NoError(t, err) + core.SetEtcdClient(etcdCli) + + randVal := rand.Int() + Params.EtcdCfg.MetaRootPath = fmt.Sprintf("/%d/%s", randVal, Params.EtcdCfg.MetaRootPath) + Params.EtcdCfg.KvRootPath = fmt.Sprintf("/%d/%s", randVal, Params.EtcdCfg.KvRootPath) + + // 1. normal init + err = core.Init() + assert.NoError(t, err) + + // 2. mock init data error + // firstly delete data + err = core.MetaTable.DeleteCredential(util.UserRoot) + assert.NoError(t, err) + + snapshotKV, err := newMetaSnapshot(etcdCli, Params.EtcdCfg.MetaRootPath, TimestampPrefix, 7) + assert.NotNil(t, snapshotKV) + assert.NoError(t, err) + txnKV := etcdkv.NewEtcdKV(etcdCli, Params.EtcdCfg.MetaRootPath) + mt, err := NewMetaTable(txnKV, snapshotKV) + assert.NoError(t, err) + mockTxnKV := &mockTestTxnKV{ + TxnKV: mt.txn, + save: func(key, value string) error { return txnKV.Save(key, value) }, + remove: func(key string) error { return txnKV.Remove(key) }, + } + mt.txn = mockTxnKV + // mock save data error + mockTxnKV.save = func(key, value string) error { + return fmt.Errorf("save error") + } + core.MetaTable = mt + err = core.initData() + assert.Error(t, err) +} + func TestRootCoord_Base(t *testing.T) { const ( dbName = "testDb" diff --git a/internal/util/crypto/crypto_test.go b/internal/util/crypto/crypto_test.go index a045c29276..58c3e24cb7 100644 --- a/internal/util/crypto/crypto_test.go +++ b/internal/util/crypto/crypto_test.go @@ -1,8 +1,13 @@ package crypto import ( + "encoding/json" + "fmt" "testing" + "github.com/golang/protobuf/proto" + "github.com/milvus-io/milvus/internal/proto/internalpb" + "github.com/milvus-io/milvus/internal/util" "github.com/stretchr/testify/assert" ) @@ -13,3 +18,25 @@ func TestPasswordVerify(t *testing.T) { assert.True(t, PasswordVerify(correctPassword, "$2a$10$3H9DLiHyPxJ29bMWRNyueOrGkbzJfE3BAR159ju3UetytAoKk7Ne2")) assert.False(t, PasswordVerify(wrongPassword, hashedPass)) } + +func TestMarshalAndPasswordVerify(t *testing.T) { + encryptedRootPassword, _ := PasswordEncrypt(util.DefaultRootPassword) + credInfo := &internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword} + v, _ := proto.Marshal(&internalpb.CredentialInfo{EncryptedPassword: credInfo.EncryptedPassword}) + fmt.Println(string(v)) + + credentialInfo := internalpb.CredentialInfo{} + proto.Unmarshal(v, &credentialInfo) + assert.True(t, PasswordVerify(util.DefaultRootPassword, credentialInfo.EncryptedPassword)) +} + +func TestJsonMarshalAndPasswordVerify(t *testing.T) { + encryptedRootPassword, _ := PasswordEncrypt(util.DefaultRootPassword) + credInfo := &internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword} + v, _ := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: credInfo.EncryptedPassword}) + fmt.Println(string(v)) + + credentialInfo := internalpb.CredentialInfo{} + json.Unmarshal(v, &credentialInfo) + assert.True(t, PasswordVerify(util.DefaultRootPassword, credentialInfo.EncryptedPassword)) +}