1 files changed, 25 insertions, 27 deletions
diff --git a/plugin/dnstap/setup.go b/plugin/dnstap/setup.go
index a863639ad..4324087dd 100644
--- a/plugin/dnstap/setup.go
+++ b/plugin/dnstap/setup.go
@@ -6,64 +6,62 @@ import (
 	"github.com/coredns/caddy"
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin"
-	"github.com/coredns/coredns/plugin/dnstap/dnstapio"
+	clog "github.com/coredns/coredns/plugin/pkg/log"
 	"github.com/coredns/coredns/plugin/pkg/parse"
 )
 
-func init() { plugin.Register("dnstap", setup) }
+var log = clog.NewWithPlugin("dnstap")
 
-type config struct {
-	proto  string
-	target string
-	full   bool
-}
+func init() { plugin.Register("dnstap", setup) }
 
-func parseConfig(d *caddy.Controller) (c config, err error) {
-	d.Next() // directive name
+func parseConfig(c *caddy.Controller) (Dnstap, error) {
+	c.Next() // directive name
+	d := Dnstap{}
+	endpoint := ""
 
-	if !d.Args(&c.target) {
-		return c, d.ArgErr()
+	if !c.Args(&endpoint) {
+		return d, c.ArgErr()
 	}
 
-	if strings.HasPrefix(c.target, "tcp://") {
+	if strings.HasPrefix(endpoint, "tcp://") {
 		// remote IP endpoint
-		servers, err := parse.HostPortOrFile(c.target[6:])
+		servers, err := parse.HostPortOrFile(endpoint[6:])
 		if err != nil {
-			return c, d.ArgErr()
+			return d, c.ArgErr()
 		}
-		c.target = servers[0]
-		c.proto = "tcp"
+		dio := newIO("tcp", servers[0])
+		d = Dnstap{io: dio}
 	} else {
-		c.target = strings.TrimPrefix(c.target, "unix://")
-		c.proto = "unix"
+		endpoint = strings.TrimPrefix(endpoint, "unix://")
+		dio := newIO("unix", endpoint)
+		d = Dnstap{io: dio}
 	}
 
-	c.full = d.NextArg() && d.Val() == "full"
+	d.IncludeRawMessage = c.NextArg() && c.Val() == "full"
 
-	return
+	return d, nil
 }
 
 func setup(c *caddy.Controller) error {
-	conf, err := parseConfig(c)
+	dnstap, err := parseConfig(c)
 	if err != nil {
 		return plugin.Error("dnstap", err)
 	}
 
-	dio := dnstapio.New(conf.proto, conf.target)
-	dnstap := Dnstap{io: dio, IncludeRawMessage: conf.full}
-
 	c.OnStartup(func() error {
-		dio.Connect()
+		if err := dnstap.io.(*dio).connect(); err != nil {
+			log.Errorf("No connection to dnstap endpoint: %s", err)
+		}
 		return nil
 	})
 
 	c.OnRestart(func() error {
-		dio.Close()
+		dnstap.io.(*dio).close()
 		return nil
 	})
 
 	c.OnFinalShutdown(func() error {
-		dio.Close()
+		dnstap.io.(*dio).close()
 		return nil
 	})