From 0455febc34676b3f8de445ab042fa6de42f31235 Mon Sep 17 00:00:00 2001
From: Ruslan Drozhdzh <30860269+rdrozhdzh@users.noreply.github.com>
Date: Tue, 8 May 2018 20:35:47 +0300
Subject: [PATCH] fix TestNewServer() hanging (#1786)

---
 plugin/pkg/dnstest/server.go      | 11 +++++++++--
 plugin/pkg/dnstest/server_test.go |  9 +++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/plugin/pkg/dnstest/server.go b/plugin/pkg/dnstest/server.go
index 72437cf8b..8e9a2c76a 100644
--- a/plugin/pkg/dnstest/server.go
+++ b/plugin/pkg/dnstest/server.go
@@ -23,8 +23,15 @@ func NewServer(f dns.HandlerFunc) *Server {
 	ch1 := make(chan bool)
 	ch2 := make(chan bool)
 
-	p, _ := net.ListenPacket("udp", ":0")
-	l, _ := net.Listen("tcp", p.LocalAddr().String())
+	l, _ := net.Listen("tcp", ":0")
+	if l == nil {
+		return nil
+	}
+	p, _ := net.ListenPacket("udp", l.Addr().String())
+	if p == nil {
+		l.Close()
+		return nil // yes, this may crash some test, but this is better than hanging
+	}
 
 	s1 := &dns.Server{PacketConn: p}
 	s2 := &dns.Server{Listener: l}
diff --git a/plugin/pkg/dnstest/server_test.go b/plugin/pkg/dnstest/server_test.go
index 3f23c340b..41450e49e 100644
--- a/plugin/pkg/dnstest/server_test.go
+++ b/plugin/pkg/dnstest/server_test.go
@@ -25,4 +25,13 @@ func TestNewServer(t *testing.T) {
 	if ret.Id != m.Id {
 		t.Fatalf("Msg ID's should match, expected %d, got %d", m.Id, ret.Id)
 	}
+
+	c.Net = "udp"
+	ret, _, err = c.Exchange(m, s.Addr)
+	if err != nil {
+		t.Fatalf("Could not send message to dnstest.Server: %s", err)
+	}
+	if ret.Id != m.Id {
+		t.Fatalf("Msg ID's should match, expected %d, got %d", m.Id, ret.Id)
+	}
 }