From 2bc1dfd9216fda881d757617d83dee306441fc72 Mon Sep 17 00:00:00 2001
From: fanmin shi <fanmin.shi@coreos.com>
Date: Wed, 5 Apr 2017 15:25:22 -0700
Subject: [PATCH] etcdmain: support SRV discovery for gRPC proxy

FIX #7562
---
 etcdmain/gateway.go    | 25 +++------------------
 etcdmain/grpc_proxy.go | 18 ++++++++++-----
 etcdmain/util.go       | 50 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 27 deletions(-)
 create mode 100644 etcdmain/util.go

diff --git a/etcdmain/gateway.go b/etcdmain/gateway.go
index 4a64a19ae2b..1a72bddcf08 100644
--- a/etcdmain/gateway.go
+++ b/etcdmain/gateway.go
@@ -21,8 +21,6 @@ import (
 	"os"
 	"time"
 
-	"github.com/coreos/etcd/client"
-	"github.com/coreos/etcd/pkg/transport"
 	"github.com/coreos/etcd/proxy/tcpproxy"
 
 	"github.com/spf13/cobra"
@@ -95,26 +93,9 @@ func stripSchema(eps []string) []string {
 }
 func startGateway(cmd *cobra.Command, args []string) {
 	endpoints := gatewayEndpoints
-	if gatewayDNSCluster != "" {
-		eps, err := client.NewSRVDiscover().Discover(gatewayDNSCluster)
-		if err != nil {
-			fmt.Fprintln(os.Stderr, err)
-			os.Exit(1)
-		}
-		plog.Infof("discovered the cluster %s from %s", eps, gatewayDNSCluster)
-		// confirm TLS connections are good
-		if !gatewayInsecureDiscovery {
-			tlsInfo := transport.TLSInfo{
-				TrustedCAFile: gatewayCA,
-				ServerName:    gatewayDNSCluster,
-			}
-			plog.Infof("validating discovered endpoints %v", eps)
-			endpoints, err = transport.ValidateSecureEndpoints(tlsInfo, eps)
-			if err != nil {
-				plog.Warningf("%v", err)
-			}
-			plog.Infof("using discovered endpoints %v", endpoints)
-		}
+
+	if eps := discoverEndpoints(gatewayDNSCluster, gatewayCA, gatewayInsecureDiscovery); len(eps) != 0 {
+		endpoints = eps
 	}
 
 	// Strip the schema from the endpoints because we start just a TCP proxy
diff --git a/etcdmain/grpc_proxy.go b/etcdmain/grpc_proxy.go
index 068da2315c9..1f701ba1297 100644
--- a/etcdmain/grpc_proxy.go
+++ b/etcdmain/grpc_proxy.go
@@ -37,11 +37,13 @@ import (
 )
 
 var (
-	grpcProxyListenAddr string
-	grpcProxyEndpoints  []string
-	grpcProxyCert       string
-	grpcProxyKey        string
-	grpcProxyCA         string
+	grpcProxyListenAddr        string
+	grpcProxyEndpoints         []string
+	grpcProxyDNSCluster        string
+	grpcProxyInsecureDiscovery bool
+	grpcProxyCert              string
+	grpcProxyKey               string
+	grpcProxyCA                string
 
 	grpcProxyAdvertiseClientURL string
 	grpcProxyResolverPrefix     string
@@ -75,6 +77,8 @@ func newGRPCProxyStartCommand() *cobra.Command {
 	}
 
 	cmd.Flags().StringVar(&grpcProxyListenAddr, "listen-addr", "127.0.0.1:23790", "listen address")
+	cmd.Flags().StringVar(&grpcProxyDNSCluster, "discovery-srv", "", "DNS domain used to bootstrap initial cluster")
+	cmd.Flags().BoolVar(&grpcProxyInsecureDiscovery, "insecure-discovery", false, "accept insecure SRV records")
 	cmd.Flags().StringSliceVar(&grpcProxyEndpoints, "endpoints", []string{"127.0.0.1:2379"}, "comma separated etcd cluster endpoints")
 	cmd.Flags().StringVar(&grpcProxyCert, "cert", "", "identify secure connections with etcd servers using this TLS certificate file")
 	cmd.Flags().StringVar(&grpcProxyKey, "key", "", "identify secure connections with etcd servers using this TLS key file")
@@ -102,6 +106,10 @@ func startGRPCProxy(cmd *cobra.Command, args []string) {
 		os.Exit(1)
 	}
 
+	if eps := discoverEndpoints(grpcProxyDNSCluster, grpcProxyCA, grpcProxyInsecureDiscovery); len(eps) != 0 {
+		grpcProxyEndpoints = eps
+	}
+
 	l, err := net.Listen("tcp", grpcProxyListenAddr)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
diff --git a/etcdmain/util.go b/etcdmain/util.go
new file mode 100644
index 00000000000..23e19b44057
--- /dev/null
+++ b/etcdmain/util.go
@@ -0,0 +1,50 @@
+// Copyright 2017 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package etcdmain
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/coreos/etcd/client"
+	"github.com/coreos/etcd/pkg/transport"
+)
+
+func discoverEndpoints(dns string, ca string, insecure bool) (endpoints []string) {
+	if dns == "" {
+		return nil
+	}
+	endpoints, err := client.NewSRVDiscover().Discover(dns)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	plog.Infof("discovered the cluster %s from %s", endpoints, dns)
+	if insecure {
+		return endpoints
+	}
+	// confirm TLS connections are good
+	tlsInfo := transport.TLSInfo{
+		TrustedCAFile: ca,
+		ServerName:    dns,
+	}
+	plog.Infof("validating discovered endpoints %v", endpoints)
+	endpoints, err = transport.ValidateSecureEndpoints(tlsInfo, endpoints)
+	if err != nil {
+		plog.Warningf("%v", err)
+	}
+	plog.Infof("using discovered endpoints %v", endpoints)
+	return endpoints
+}