plugin/transfer: fix go-routine leak (#4380)

PR #4161 is stalled. Tried to cherry pick the code from there, but that
led to conflicts, manually copying over while taking into account the
comments on that PR. Use that code and extend the error checking, don't
modify existing tests and make the badwriter test simpler.

Closes: #4161

Signed-off-by: Miek Gieben <miek@miek.nl>

add tests

Signed-off-by: Miek Gieben <miek@miek.nl>

1 files changed, 31 insertions, 0 deletions

diff --git a/plugin/transfer/failed_write_test.go b/plugin/transfer/failed_write_test.go
new file mode 100644
index 000000000..90b5c4de2
--- /dev/null
+++ b/plugin/transfer/failed_write_test.go
@@ -0,0 +1,31 @@
+package transfer
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/coredns/coredns/plugin/test"
+
+	"github.com/miekg/dns"
+)
+
+type badwriter struct {
+	dns.ResponseWriter
+	count int
+}
+
+func (w *badwriter) WriteMsg(res *dns.Msg) error { return fmt.Errorf("failed to write msg") }
+
+func TestWriteMessageFailed(t *testing.T) {
+	transfer := newTestTransfer()
+	ctx := context.TODO()
+	w := &badwriter{ResponseWriter: &test.ResponseWriter{}}
+	m := &dns.Msg{}
+	m.SetAxfr("example.org.")
+
+	_, err := transfer.ServeDNS(ctx, w, m)
+	if err == nil {
+		t.Error("Expected error, got none")
+	}
+}