From 2f9c42d82ec5e2edd475e8166547ad7a89ca4f1d Mon Sep 17 00:00:00 2001 From: Thong Huynh Date: Fri, 29 Sep 2017 13:38:01 -0700 Subject: [PATCH] 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 --- plugin/dnstap/handler.go | 18 ++++++++-- plugin/dnstap/taprw/writer.go | 45 ++++++++++++++++--------- plugin/dnstap/taprw/writer_test.go | 54 ++++++++++++++++++++++++++++++ 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 + } +}