Enable dnstap plugin to insert other plugin's specific data into extra field of tap.Dnstap message (#1101)

* Add custom data into dnstap context

* Fix error and fix UT compile errors

* Add UTs

* Change as per review comments.  Use boolean to indicate which Dnstap message to send out

* Merge with master and fix lint warning

* Remove newline

* Fix review comments

3 files changed, 99 insertions, 18 deletions

diff --git a/plugin/dnstap/handler.go b/plugin/dnstap/handler.go
index b6a8afbe7..63838492c 100644
--- a/plugin/dnstap/handler.go
+++ b/plugin/dnstap/handler.go
@@ -36,6 +36,14 @@ type (
 	}
 )
 
+// ContextKey defines the type of key that is used to save data into the context
+type ContextKey string
+
+const (
+	// DnstapSendOption specifies the Dnstap message to be send.  Default is sent all.
+	DnstapSendOption ContextKey = "dnstap-send-option"
+)
+
 // TapperFromContext will return a Tapper if the dnstap plugin is enabled.
 func TapperFromContext(ctx context.Context) (t Tapper) {
 	t, _ = ctx.(Tapper)
@@ -64,10 +72,16 @@ func (h Dnstap) TapBuilder() msg.Builder {
 
 // ServeDNS logs the client query and response to dnstap and passes the dnstap Context.
 func (h Dnstap) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
-	rw := &taprw.ResponseWriter{ResponseWriter: w, Tapper: &h, Query: r}
+
+	// Add send option into context so other plugin can decide on which DNSTap
+	// message to be sent out
+	sendOption := taprw.SendOption{Cq: true, Cr: true}
+	newCtx := context.WithValue(ctx, DnstapSendOption, &sendOption)
+
+	rw := &taprw.ResponseWriter{ResponseWriter: w, Tapper: &h, Query: r, Send: &sendOption}
 	rw.QueryEpoch()
 
-	code, err := plugin.NextOrFailure(h.Name(), h.Next, tapContext{ctx, h}, rw, r)
+	code, err := plugin.NextOrFailure(h.Name(), h.Next, tapContext{newCtx, h}, rw, r)
 	if err != nil {
 		// ignore dnstap errors
 		return code, err
diff --git a/plugin/dnstap/taprw/writer.go b/plugin/dnstap/taprw/writer.go
index ae9965411..14a899443 100644
--- a/plugin/dnstap/taprw/writer.go
+++ b/plugin/dnstap/taprw/writer.go
@@ -11,6 +11,13 @@ import (
 	"github.com/miekg/dns"
 )
 
+// SendOption stores the flag to indicate whether a certain DNSTap message to
+// be sent out or not.
+type SendOption struct {
+	Cq bool
+	Cr bool
+}
+
 // Tapper is what ResponseWriter needs to log to dnstap.
 type Tapper interface {
 	TapMessage(m *tap.Message) error
@@ -19,12 +26,14 @@ type Tapper interface {
 
 // ResponseWriter captures the client response and logs the query to dnstap.
 // Single request use.
+// SendOption configures Dnstap to selectively send Dnstap messages. Default is send all.
 type ResponseWriter struct {
 	queryEpoch uint64
 	Query      *dns.Msg
 	dns.ResponseWriter
 	Tapper
-	err error
+	err  error
+	Send *SendOption
 }
 
 // DnstapError check if a dnstap error occurred during Write and returns it.
@@ -46,28 +55,32 @@ func (w *ResponseWriter) WriteMsg(resp *dns.Msg) (writeErr error) {
 
 	b := w.TapBuilder()
 	b.TimeSec = w.queryEpoch
-	if err := func() (err error) {
-		err = b.AddrMsg(w.ResponseWriter.RemoteAddr(), w.Query)
-		if err != nil {
-			return
-		}
-		return w.TapMessage(b.ToClientQuery())
-	}(); err != nil {
-		w.err = fmt.Errorf("client query: %s", err)
-		// don't forget to call DnstapError later
-	}
 
-	if writeErr == nil {
+	if w.Send == nil || w.Send.Cq {
 		if err := func() (err error) {
-			b.TimeSec = writeEpoch
-			if err = b.Msg(resp); err != nil {
+			err = b.AddrMsg(w.ResponseWriter.RemoteAddr(), w.Query)
+			if err != nil {
 				return
 			}
-			return w.TapMessage(b.ToClientResponse())
+			return w.TapMessage(b.ToClientQuery())
 		}(); err != nil {
-			w.err = fmt.Errorf("client response: %s", err)
+			w.err = fmt.Errorf("client query: %s", err)
+			// don't forget to call DnstapError later
 		}
 	}
 
+	if w.Send == nil || w.Send.Cr {
+		if writeErr == nil {
+			if err := func() (err error) {
+				b.TimeSec = writeEpoch
+				if err = b.Msg(resp); err != nil {
+					return
+				}
+				return w.TapMessage(b.ToClientResponse())
+			}(); err != nil {
+				w.err = fmt.Errorf("client response: %s", err)
+			}
+		}
+	}
 	return
 }
diff --git a/plugin/dnstap/taprw/writer_test.go b/plugin/dnstap/taprw/writer_test.go
index 6969dc515..740b44770 100644
--- a/plugin/dnstap/taprw/writer_test.go
+++ b/plugin/dnstap/taprw/writer_test.go
@@ -80,3 +80,57 @@ func TestClientQueryResponse(t *testing.T) {
 		t.Fatalf("response: want: %v\nhave: %v", want, have)
 	}
 }
+
+func TestClientQueryResponseWithSendOption(t *testing.T) {
+	trapper := test.TrapTapper{Full: true}
+	m := testingMsg()
+	rw := ResponseWriter{
+		Query:          m,
+		Tapper:         &trapper,
+		ResponseWriter: &mwtest.ResponseWriter{},
+	}
+	d := test.TestingData()
+	bin, err := m.Pack()
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	d.Packed = bin
+
+	// Do not send both CQ and CR
+	o := SendOption{Cq: false, Cr: false}
+	rw.Send = &o
+
+	if err := rw.WriteMsg(m); err != nil {
+		t.Fatal(err)
+		return
+	}
+	if l := len(trapper.Trap); l != 0 {
+		t.Fatalf("%d msg trapped", l)
+		return
+	}
+
+	//Send CQ
+	o.Cq = true
+	if err := rw.WriteMsg(m); err != nil {
+		t.Fatal(err)
+		return
+	}
+	if l := len(trapper.Trap); l != 1 {
+		t.Fatalf("%d msg trapped", l)
+		return
+	}
+
+	//Send CR
+	trapper.Trap = trapper.Trap[:0]
+	o.Cq = false
+	o.Cr = true
+	if err := rw.WriteMsg(m); err != nil {
+		t.Fatal(err)
+		return
+	}
+	if l := len(trapper.Trap); l != 1 {
+		t.Fatalf("%d msg trapped", l)
+		return
+	}
+}