Certificate authentication support (#494)

Signed-off-by: “gifi-siby” <[email protected]>

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`                                        |

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)

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:

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)

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)
+	}
 }

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)
 }