From adbf107973480e82f536b3a62d8ec4c74dd58052 Mon Sep 17 00:00:00 2001
From: jiaxuanzhou <a3721659a@hotmail.com>
Date: Tue, 29 Aug 2017 23:35:34 +0800
Subject: [PATCH] etcdctl: add discovery-srv global flag for v3

---
 etcdctl/ctlv3/command/global.go | 33 +++++++++++++++++++++++----------
 etcdctl/ctlv3/ctl.go            |  1 +
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/etcdctl/ctlv3/command/global.go b/etcdctl/ctlv3/command/global.go
index 616d32e21ef7..581d1da41749 100644
--- a/etcdctl/ctlv3/command/global.go
+++ b/etcdctl/ctlv3/command/global.go
@@ -51,10 +51,10 @@ type GlobalFlags struct {
 }
 
 type secureCfg struct {
-	cert   string
-	key    string
-	cacert string
-
+	cert               string
+	key                string
+	cacert             string
+	serverName         string
 	insecureTransport  bool
 	insecureSkipVerify bool
 }
@@ -137,6 +137,11 @@ func newClientCfg(endpoints []string, dialTimeout time.Duration, scfg *secureCfg
 		cfgtls = &tlsinfo
 	}
 
+	if scfg.serverName != "" {
+		tlsinfo.ServerName = scfg.serverName
+		cfgtls = &tlsinfo
+	}
+
 	cfg := &clientv3.Config{
 		Endpoints:   endpoints,
 		DialTimeout: dialTimeout,
@@ -189,14 +194,15 @@ func dialTimeoutFromCmd(cmd *cobra.Command) time.Duration {
 }
 
 func secureCfgFromCmd(cmd *cobra.Command) *secureCfg {
-	cert, key, cacert := keyAndCertFromCmd(cmd)
+	cert, key, cacert, serverName := keyAndCertFromCmd(cmd)
 	insecureTr := insecureTransportFromCmd(cmd)
 	skipVerify := insecureSkipVerifyFromCmd(cmd)
 
 	return &secureCfg{
-		cert:   cert,
-		key:    key,
-		cacert: cacert,
+		cert:       cert,
+		key:        key,
+		cacert:     cacert,
+		serverName: serverName,
 
 		insecureTransport:  insecureTr,
 		insecureSkipVerify: skipVerify,
@@ -219,7 +225,7 @@ func insecureSkipVerifyFromCmd(cmd *cobra.Command) bool {
 	return skipVerify
 }
 
-func keyAndCertFromCmd(cmd *cobra.Command) (cert, key, cacert string) {
+func keyAndCertFromCmd(cmd *cobra.Command) (cert, key, cacert, serverName string) {
 	var err error
 	if cert, err = cmd.Flags().GetString("cert"); err != nil {
 		ExitWithError(ExitBadArgs, err)
@@ -239,7 +245,14 @@ func keyAndCertFromCmd(cmd *cobra.Command) (cert, key, cacert string) {
 		ExitWithError(ExitBadArgs, errors.New("empty string is passed to --cacert option"))
 	}
 
-	return cert, key, cacert
+	// TODO: reconstruct the duplicated code
+	if serverName, err = cmd.Flags().GetString("discovery-srv"); err != nil {
+		ExitWithError(ExitBadArgs, err)
+	} else if serverName == "" && cmd.Flags().Changed("discovery-srv") {
+		ExitWithError(ExitBadArgs, errors.New("empty string is passed to --discovery-srv option"))
+	}
+
+	return cert, key, cacert, serverName
 }
 
 func authCfgFromCmd(cmd *cobra.Command) *authCfg {
diff --git a/etcdctl/ctlv3/ctl.go b/etcdctl/ctlv3/ctl.go
index 2ca2b409622d..430c1ca49ff3 100644
--- a/etcdctl/ctlv3/ctl.go
+++ b/etcdctl/ctlv3/ctl.go
@@ -59,6 +59,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVar(&globalFlags.TLS.KeyFile, "key", "", "identify secure client using this TLS key file")
 	rootCmd.PersistentFlags().StringVar(&globalFlags.TLS.CAFile, "cacert", "", "verify certificates of TLS-enabled secure servers using this CA bundle")
 	rootCmd.PersistentFlags().StringVar(&globalFlags.User, "user", "", "username[:password] for authentication (prompt if password is not supplied)")
+	rootCmd.PersistentFlags().StringVar(&globalFlags.TLS.ServerName, "discovery-srv", "", "domain name to query for SRV records describing cluster endpoints")
 
 	rootCmd.AddCommand(
 		command.NewGetCommand(),