more updates

Signed-off-by: Miek Gieben <miek@miek.nl>

plugin/traffic/setup.go

@@ -1,6 +1,7 @@
 package traffic
 
 import (
+	"crypto/tls"
 	"fmt"
 	"math/rand"
 	"strings"
@@ -10,10 +11,13 @@ import (
 	"github.com/coredns/coredns/plugin"
 	clog "github.com/coredns/coredns/plugin/pkg/log"
 	"github.com/coredns/coredns/plugin/pkg/parse"
+	pkgtls "github.com/coredns/coredns/plugin/pkg/tls"
 	"github.com/coredns/coredns/plugin/pkg/transport"
 	"github.com/coredns/coredns/plugin/traffic/xds"
 
 	"github.com/caddyserver/caddy"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
 )
 
 var log = clog.NewWithPlugin("traffic")
@@ -32,23 +36,11 @@ func setup(c *caddy.Controller) error {
 		return t
 	})
 
-	stream, err := t.c.Run()
-	if err != nil {
-		return plugin.Error("traffic", err)
-	}
-
-	if err := t.c.ClusterDiscovery(stream, "", "", []string{}); err != nil {
-		log.Error(err)
-	}
-
-	go func() {
-		err = t.c.Receive(stream)
-		if err != nil {
-			// can't do log debug in setup functions
-			log.Debug(err)
-		}
-	}()
-
+	c.OnStartup(func() error {
+		go t.c.Run()
+		return nil
+	})
+	c.OnShutdown(func() error { return t.c.Stop() })
 	return nil
 }
 
@@ -56,7 +48,11 @@ func parseTraffic(c *caddy.Controller) (*Traffic, error) {
 	node := "coredns"
 	toHosts := []string{}
 	t := &Traffic{}
-	var err error
+	var (
+		err           error
+		tlsConfig     *tls.Config
+		tlsServerName string
+	)
 
 	t.origins = make([]string, len(c.ServerBlockKeys))
 	for i := range c.ServerBlockKeys {
@@ -88,14 +84,37 @@ func parseTraffic(c *caddy.Controller) (*Traffic, error) {
 					return nil, c.ArgErr()
 				}
 				node = args[0]
+			case "tls":
+				args := c.RemainingArgs()
+				if len(args) > 3 {
+					return nil, c.ArgErr()
+				}
+
+				tlsConfig, err = pkgtls.NewTLSConfigFromArgs(args...)
+				if err != nil {
+					return nil, err
+				}
+			case "tls_servername":
+				if !c.NextArg() {
+					return nil, c.ArgErr()
+				}
+				tlsServerName = c.Val()
 			default:
 				return nil, c.Errf("unknown property '%s'", c.Val())
 			}
 		}
 	}
 
+	opts := []grpc.DialOption{grpc.WithInsecure()}
+	if tlsConfig != nil {
+		if tlsServerName != "" {
+			tlsConfig.ServerName = tlsServerName
+		}
+		opts = []grpc.DialOption{grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))}
+	}
+
 	// TODO: only the first host is used, need to figure out how to reconcile multiple upstream providers.
-	if t.c, err = xds.New(toHosts[0], node); err != nil {
+	if t.c, err = xds.New(toHosts[0], node, opts...); err != nil {
 		return nil, err
 	}