Plugin dnstap: add support for "extra" field in payload (#6226)

* dnstap: add 'extra' field

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* dnstap: add setup_test for 'extra' field

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* udnstap: update document and test

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* dnstap: update setup_test for more coverage

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* dnstap: add TapMessageWithMetadata function to Dnstap

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* dnstap: adapt dnstap and forward plugins to use TapMessageWithMetadata

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* change TapMessageWithMetadata function

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* tab inconsistency fix

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* fix replacer to support empty state

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* add replacer test for empty status parameter

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* dnstap: update unit test for 'extra' field

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* clean up code

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* gofmt fix & static analysis fix

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

* dnstap: refactor

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

---------

Signed-off-by: chenyuheng <chenyuheng99@qq.com>

1 files changed, 62 insertions, 10 deletions

diff --git a/plugin/dnstap/handler_test.go b/plugin/dnstap/handler_test.go
index 2c54f70e6..cb492ac10 100644
--- a/plugin/dnstap/handler_test.go
+++ b/plugin/dnstap/handler_test.go
@@ -6,13 +6,14 @@ import (
 	"testing"
 
 	"github.com/coredns/coredns/plugin/dnstap/msg"
+	"github.com/coredns/coredns/plugin/metadata"
 	test "github.com/coredns/coredns/plugin/test"
 
 	tap "github.com/dnstap/golang-dnstap"
 	"github.com/miekg/dns"
 )
 
-func testCase(t *testing.T, tapq, tapr *tap.Message, q, r *dns.Msg) {
+func testCase(t *testing.T, tapq, tapr *tap.Dnstap, q, r *dns.Msg, extraFormat string) {
 	w := writer{t: t}
 	w.queue = append(w.queue, tapq, tapr)
 	h := Dnstap{
@@ -20,9 +21,17 @@ func testCase(t *testing.T, tapq, tapr *tap.Message, q, r *dns.Msg) {
 			w dns.ResponseWriter, _ *dns.Msg) (int, error) {
 			return 0, w.WriteMsg(r)
 		}),
-		io: &w,
+		io:          &w,
+		ExtraFormat: extraFormat,
 	}
-	_, err := h.ServeDNS(context.TODO(), &test.ResponseWriter{}, q)
+	ctx := metadata.ContextWithMetadata(context.TODO())
+	ok := metadata.SetValueFunc(ctx, "metadata/test", func() string {
+		return "MetadataValue"
+	})
+	if !ok {
+		t.Fatal("Failed to set metadata")
+	}
+	_, err := h.ServeDNS(ctx, &test.ResponseWriter{}, q)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -30,7 +39,7 @@ func testCase(t *testing.T, tapq, tapr *tap.Message, q, r *dns.Msg) {
 
 type writer struct {
 	t     *testing.T
-	queue []*tap.Message
+	queue []*tap.Dnstap
 }
 
 func (w *writer) Dnstap(e *tap.Dnstap) {
@@ -38,7 +47,7 @@ func (w *writer) Dnstap(e *tap.Dnstap) {
 		w.t.Error("Message not expected")
 	}
 
-	ex := w.queue[0]
+	ex := w.queue[0].Message
 	got := e.Message
 
 	if string(ex.QueryAddress) != string(got.QueryAddress) {
@@ -53,6 +62,9 @@ func (w *writer) Dnstap(e *tap.Dnstap) {
 	if *ex.SocketFamily != *got.SocketFamily {
 		w.t.Errorf("Expected socket family %d, got %d", *ex.SocketFamily, *got.SocketFamily)
 	}
+	if string(w.queue[0].Extra) != string(e.Extra) {
+		w.t.Errorf("Expected extra %s, got %s", w.queue[0].Extra, e.Extra)
+	}
 	w.queue = w.queue[1:]
 }
 
@@ -64,11 +76,29 @@ func TestDnstap(t *testing.T) {
 			test.A("example.org. 3600	IN	A 10.0.0.1"),
 		},
 	}.Msg()
-	tapq := testMessage() // leave type unset for deepEqual
-	msg.SetType(tapq, tap.Message_CLIENT_QUERY)
-	tapr := testMessage()
-	msg.SetType(tapr, tap.Message_CLIENT_RESPONSE)
-	testCase(t, tapq, tapr, q, r)
+
+	tapq := &tap.Dnstap{
+		Message: testMessage(),
+	}
+	msg.SetType(tapq.Message, tap.Message_CLIENT_QUERY)
+	tapr := &tap.Dnstap{
+		Message: testMessage(),
+	}
+	msg.SetType(tapr.Message, tap.Message_CLIENT_RESPONSE)
+	testCase(t, tapq, tapr, q, r, "")
+
+	tapq_with_extra := &tap.Dnstap{
+		Message: testMessage(), // leave type unset for deepEqual
+		Extra:   []byte("extra_field_MetadataValue_A_example.org._IN_udp_29_10.240.0.1_40212_127.0.0.1"),
+	}
+	msg.SetType(tapq_with_extra.Message, tap.Message_CLIENT_QUERY)
+	tapr_with_extra := &tap.Dnstap{
+		Message: testMessage(),
+		Extra:   []byte("extra_field_MetadataValue_A_example.org._IN_udp_29_10.240.0.1_40212_127.0.0.1"),
+	}
+	msg.SetType(tapr_with_extra.Message, tap.Message_CLIENT_RESPONSE)
+	extraFormat := "extra_field_{/metadata/test}_{type}_{name}_{class}_{proto}_{size}_{remote}_{port}_{local}"
+	testCase(t, tapq_with_extra, tapr_with_extra, q, r, extraFormat)
 }
 
 func testMessage() *tap.Message {
@@ -82,3 +112,25 @@ func testMessage() *tap.Message {
 		QueryPort:      &port,
 	}
 }
+
+func TestTapMessage(t *testing.T) {
+	extraFormat := "extra_field_no_replacement_{/metadata/test}_{type}_{name}_{class}_{proto}_{size}_{remote}_{port}_{local}"
+	tapq := &tap.Dnstap{
+		Message: testMessage(),
+		// extra field would not be replaced, since TapMessage won't pass context
+		Extra: []byte(extraFormat),
+	}
+	msg.SetType(tapq.Message, tap.Message_CLIENT_QUERY)
+
+	w := writer{t: t}
+	w.queue = append(w.queue, tapq)
+	h := Dnstap{
+		Next: test.HandlerFunc(func(_ context.Context,
+			w dns.ResponseWriter, r *dns.Msg) (int, error) {
+			return 0, w.WriteMsg(r)
+		}),
+		io:          &w,
+		ExtraFormat: extraFormat,
+	}
+	h.TapMessage(tapq.Message)
+}