plugin/forward: fix broken tap plugins when dnstap plugins specified (#5890)

* plugin/forward: fix broken tap plugins when dnstap plugins specified

---------

Signed-off-by: Gerhard Tan <gwohau.tan@gmail.com>

5 files changed, 60 insertions, 13 deletions

diff --git a/plugin/dnstap/README.md b/plugin/dnstap/README.md
index 7405ddd77..fc59033ac 100644
--- a/plugin/dnstap/README.md
+++ b/plugin/dnstap/README.md
@@ -102,8 +102,8 @@ x :=  &ExamplePlugin{}
 
 c.OnStartup(func() error {
     if taph := dnsserver.GetConfig(c).Handler("dnstap"); taph != nil {
-        if tapPlugin, ok := taph.(dnstap.Dnstap); ok {
-            x.tapPlugins = append(x.tapPlugins, &tapPlugin)
+        for tapPlugin, ok := taph.(*dnstap.Dnstap); ok; tapPlugin, ok = tapPlugin.Next.(*dnstap.Dnstap) {
+            x.tapPlugins = append(x.tapPlugins, tapPlugin)
         }
     }
     return nil
diff --git a/plugin/forward/dnstap.go b/plugin/forward/dnstap.go
index 95ef90acf..edbee8715 100644
--- a/plugin/forward/dnstap.go
+++ b/plugin/forward/dnstap.go
@@ -14,9 +14,6 @@ import (
 
 // toDnstap will send the forward and received message to the dnstap plugin.
 func toDnstap(f *Forward, host string, state request.Request, opts options, reply *dns.Msg, start time.Time) {
-	// Query
-	q := new(tap.Message)
-	msg.SetQueryTime(q, start)
 	h, p, _ := net.SplitHostPort(host)      // this is preparsed and can't err here
 	port, _ := strconv.ParseUint(p, 10, 32) // same here
 	ip := net.ParseIP(h)
@@ -34,12 +31,14 @@ func toDnstap(f *Forward, host string, state request.Request, opts options, repl
 		ta = &net.TCPAddr{IP: ip, Port: int(port)}
 	}
 
-	// Forwarder dnstap messages are from the perspective of the downstream server
-	// (upstream is the forward server)
-	msg.SetQueryAddress(q, state.W.RemoteAddr())
-	msg.SetResponseAddress(q, ta)
-
 	for _, t := range f.tapPlugins {
+		// Query
+		q := new(tap.Message)
+		msg.SetQueryTime(q, start)
+		// Forwarder dnstap messages are from the perspective of the downstream server
+		// (upstream is the forward server)
+		msg.SetQueryAddress(q, state.W.RemoteAddr())
+		msg.SetResponseAddress(q, ta)
 		if t.IncludeRawMessage {
 			buf, _ := state.Req.Pack()
 			q.QueryMessage = buf
diff --git a/plugin/forward/forward.go b/plugin/forward/forward.go
index b32e8b3b4..223d7e398 100644
--- a/plugin/forward/forward.go
+++ b/plugin/forward/forward.go
@@ -66,6 +66,14 @@ func (f *Forward) SetProxy(p *Proxy) {
 	p.start(f.hcInterval)
 }
 
+// SetTapPlugin appends one or more dnstap plugins to the tap plugin list.
+func (f *Forward) SetTapPlugin(tapPlugin *dnstap.Dnstap) {
+	f.tapPlugins = append(f.tapPlugins, tapPlugin)
+	if nextPlugin, ok := tapPlugin.Next.(*dnstap.Dnstap); ok {
+		f.SetTapPlugin(nextPlugin)
+	}
+}
+
 // Len returns the number of configured proxies.
 func (f *Forward) Len() int { return len(f.proxies) }
 
diff --git a/plugin/forward/forward_test.go b/plugin/forward/forward_test.go
index b0ef47ba9..b50f4ff22 100644
--- a/plugin/forward/forward_test.go
+++ b/plugin/forward/forward_test.go
@@ -1,7 +1,13 @@
 package forward
 
 import (
+	"strings"
 	"testing"
+
+	"github.com/coredns/caddy"
+	"github.com/coredns/caddy/caddyfile"
+	"github.com/coredns/coredns/core/dnsserver"
+	"github.com/coredns/coredns/plugin/dnstap"
 )
 
 func TestList(t *testing.T) {
@@ -22,3 +28,39 @@ func TestList(t *testing.T) {
 		}
 	}
 }
+
+func TestSetTapPlugin(t *testing.T) {
+	input := `forward . 127.0.0.1
+	dnstap /tmp/dnstap.sock full
+	dnstap tcp://example.com:6000
+	`
+	stanzas := strings.Split(input, "\n")
+	c := caddy.NewTestController("dns", strings.Join(stanzas[1:], "\n"))
+	dnstapSetup, err := caddy.DirectiveAction("dns", "dnstap")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = dnstapSetup(c); err != nil {
+		t.Fatal(err)
+	}
+	c.Dispenser = caddyfile.NewDispenser("", strings.NewReader(stanzas[0]))
+	if err = setup(c); err != nil {
+		t.Fatal(err)
+	}
+	dnsserver.NewServer("", []*dnsserver.Config{dnsserver.GetConfig(c)})
+	f, ok := dnsserver.GetConfig(c).Handler("forward").(*Forward)
+	if !ok {
+		t.Fatal("Expected a forward plugin")
+	}
+	tap, ok := dnsserver.GetConfig(c).Handler("dnstap").(*dnstap.Dnstap)
+	if !ok {
+		t.Fatal("Expected a dnstap plugin")
+	}
+	f.SetTapPlugin(tap)
+	if len(f.tapPlugins) != 2 {
+		t.Fatalf("Expected: 2 results, got: %v", len(f.tapPlugins))
+	}
+	if f.tapPlugins[0] != tap || tap.Next != f.tapPlugins[1] {
+		t.Error("Unexpected order of dnstap plugins")
+	}
+}
diff --git a/plugin/forward/setup.go b/plugin/forward/setup.go
index 39709301e..7ca24df4d 100644
--- a/plugin/forward/setup.go
+++ b/plugin/forward/setup.go
@@ -51,9 +51,7 @@ func setup(c *caddy.Controller) error {
 		})
 		c.OnStartup(func() error {
 			if taph := dnsserver.GetConfig(c).Handler("dnstap"); taph != nil {
-				if tapPlugin, ok := taph.(dnstap.Dnstap); ok {
-					f.tapPlugins = append(f.tapPlugins, &tapPlugin)
-				}
+				f.SetTapPlugin(taph.(*dnstap.Dnstap))
 			}
 			return nil
 		})