diff --git a/README.md b/README.md index 87b0d62..7b59369 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,8 @@ Below is a summary of the configurations supported in `backup.yaml`: | | `tlsMode` | TLS mode (0: none, 1: one-way, 2: two-way). | `0` | | | `user` | Username for Milvus. | `root` | | | `password` | Password for Milvus. | `Milvus` | +| | `tlsCertPath` | Path to your certificate file | `/path/to/certificate` | +| | `serverName ` | Server name | `localhost` | | `minio` | `storageType` | Storage type for Milvus (e.g., `local`, `minio`, `s3`, `aws`, `gcp`, `ali(aliyun)`, `azure`, `tc(tencent)`). | `minio` | | | `address` | MinIO/S3 address. | `localhost` | | | `port` | MinIO/S3 port. | `9000` | diff --git a/cmd/backup_yaml.go b/cmd/backup_yaml.go index 20e4492..e63f7c5 100644 --- a/cmd/backup_yaml.go +++ b/cmd/backup_yaml.go @@ -3,6 +3,7 @@ package cmd import ( "fmt" "strings" + "github.com/spf13/cobra" "github.com/zilliztech/milvus-backup/core/paramtable" "gopkg.in/yaml.v3" @@ -39,6 +40,8 @@ type YAMLConFig struct { TlsMode int `yaml:"tlsMode"` User string `yaml:"user"` Password string `yaml:"password"` + TlsCertPath string `yaml:"tlsCertPath"` + ServerName string `yaml:"serverName"` } `yaml:"milvus"` Minio struct { Address string `yaml:"address"` @@ -77,6 +80,8 @@ func printParams(base *paramtable.BackupParams) { yml.Milvus.TlsMode = base.ParseIntWithDefault("milvus.tlsMode", 0) yml.Milvus.User = base.BaseTable.LoadWithDefault("milvus.user", "") yml.Milvus.Password = base.BaseTable.LoadWithDefault("milvus.password", "") + yml.Milvus.TlsCertPath = base.BaseTable.LoadWithDefault("milvus.tlsCertPath", "") + yml.Milvus.ServerName = base.BaseTable.LoadWithDefault("milvus.serverName", "localhost") yml.Minio.Address = base.LoadWithDefault("minio.address", "localhost") yml.Minio.Port = base.ParseIntWithDefault("minio.port", 9000) diff --git a/configs/backup.yaml b/configs/backup.yaml index ae5a4bb..3aa23ed 100644 --- a/configs/backup.yaml +++ b/configs/backup.yaml @@ -18,6 +18,8 @@ milvus: tlsMode: 0 user: "root" password: "Milvus" + tlsCertPath: "" + serverName: "" # Related configuration of minio, which is responsible for data persistence for Milvus. minio: diff --git a/core/backup_context.go b/core/backup_context.go index 7c76e95..4db0581 100644 --- a/core/backup_context.go +++ b/core/backup_context.go @@ -19,6 +19,9 @@ import ( "github.com/zilliztech/milvus-backup/core/utils" "github.com/zilliztech/milvus-backup/internal/common" "github.com/zilliztech/milvus-backup/internal/log" + + grpc "google.golang.org/grpc" + "google.golang.org/grpc/credentials" ) const ( @@ -69,13 +72,35 @@ func CreateMilvusClient(ctx context.Context, params paramtable.BackupParams) (go var c gomilvus.Client var err error if params.MilvusCfg.AuthorizationEnabled && params.MilvusCfg.User != "" && params.MilvusCfg.Password != "" { - if params.MilvusCfg.TLSMode == 0 { + switch params.MilvusCfg.TLSMode { + case 0: c, err = gomilvus.NewDefaultGrpcClientWithAuth(ctx, milvusEndpoint, params.MilvusCfg.User, params.MilvusCfg.Password) - } else if params.MilvusCfg.TLSMode == 1 || params.MilvusCfg.TLSMode == 2 { + case 1: + if params.MilvusCfg.TLSCertPath != "" { + var creds credentials.TransportCredentials + creds, err = credentials.NewClientTLSFromFile(params.MilvusCfg.TLSCertPath, params.MilvusCfg.ServerName) + if err != nil { + log.Error("failed to create client from the certificate", zap.Error(err)) + return nil, err + } + opts := []grpc.DialOption{ + grpc.WithTransportCredentials(creds), + } + c, err = gomilvus.NewClient(ctx, gomilvus.Config{ + Address: milvusEndpoint, + Username: params.MilvusCfg.User, + Password: params.MilvusCfg.Password, + EnableTLSAuth: true, + DialOptions: opts, + }) + } else { + c, err = gomilvus.NewDefaultGrpcClientWithTLSAuth(ctx, milvusEndpoint, params.MilvusCfg.User, params.MilvusCfg.Password) + } + case 2: c, err = gomilvus.NewDefaultGrpcClientWithTLSAuth(ctx, milvusEndpoint, params.MilvusCfg.User, params.MilvusCfg.Password) - } else { - log.Error("milvus.TLSMode is not illegal, support value 0, 1, 2") - return nil, errors.New("milvus.TLSMode is not illegal, support value 0, 1, 2") + default: + log.Error("milvus.TLSMode is illegal, support value 0, 1, 2") + return nil, errors.New("milvus.TLSMode is illegal, support value 0, 1, 2") } } else { c, err = gomilvus.NewGrpcClient(ctx, milvusEndpoint) diff --git a/core/paramtable/base_table.go b/core/paramtable/base_table.go index 23087b9..7fb2e88 100644 --- a/core/paramtable/base_table.go +++ b/core/paramtable/base_table.go @@ -57,6 +57,8 @@ const ( DefaultMilvusTlsMode = "0" DefaultMilvusUser = "root" DefaultMilvusPassword = "Milvus" + DefaultMilvusTLSCertPath = "" + DefaultMilvusServerName = "" ) var defaultYaml = DefaultBackupYaml @@ -529,4 +531,14 @@ func (gp *BaseTable) loadMilvusConfig() { if milvusPassword != "" { _ = gp.Save("milvus.password", milvusPassword) } + + milvusTLSCertPath := os.Getenv("MILVUS_TLS_CERTPATH") + if milvusTLSCertPath != "" { + _ = gp.Save("milvus.tlsCertPath", milvusTLSCertPath) + } + + milvusServerName := os.Getenv("MILVUS_SERVER_NAME") + if milvusServerName != "" { + _ = gp.Save("milvus.serverName", milvusServerName) + } } diff --git a/core/paramtable/params.go b/core/paramtable/params.go index 210974c..5a07682 100644 --- a/core/paramtable/params.go +++ b/core/paramtable/params.go @@ -109,6 +109,8 @@ type MilvusConfig struct { Port string User string Password string + TLSCertPath string + ServerName string AuthorizationEnabled bool TLSMode int } @@ -120,6 +122,8 @@ func (p *MilvusConfig) init(base *BaseTable) { p.initPort() p.initUser() p.initPassword() + p.initTLSCertPath() + p.initServerName() p.initAuthorizationEnabled() p.initTLSMode() } @@ -156,6 +160,22 @@ func (p *MilvusConfig) initPassword() { p.Password = password } +func (p *MilvusConfig) initTLSCertPath() { + tlsCertPath, err := p.Base.Load("milvus.tlsCertPath") + if err != nil { + panic(err) + } + p.TLSCertPath = tlsCertPath +} + +func (p *MilvusConfig) initServerName() { + serverName, err := p.Base.Load("milvus.serverName") + if err != nil { + panic(err) + } + p.ServerName = serverName +} + func (p *MilvusConfig) initAuthorizationEnabled() { p.AuthorizationEnabled = p.Base.ParseBool("milvus.authorizationEnabled", false) }