From e68c0d088b412bb4211997f21f964a6fd65b4c70 Mon Sep 17 00:00:00 2001
From: Comzyh <comzyh@users.noreply.github.com>
Date: Fri, 10 Jan 2020 14:13:44 +0800
Subject: [PATCH] Fix: upstream dns ExchangeContext workaround (#468)

---
 common/picker/picker.go | 1 +
 dns/client.go           | 1 +
 dns/resolver.go         | 4 ++--
 dns/util.go             | 1 +
 4 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/common/picker/picker.go b/common/picker/picker.go
index a30202d..49a58f0 100644
--- a/common/picker/picker.go
+++ b/common/picker/picker.go
@@ -59,6 +59,7 @@ func (p *Picker) Wait() interface{} {
 }
 
 // WaitWithoutCancel blocks until the first result return, if timeout will return nil.
+// The return of this function will not wait for the cancel of context.
 func (p *Picker) WaitWithoutCancel() interface{} {
 	select {
 	case <-p.firstDone:
diff --git a/dns/client.go b/dns/client.go
index 881fcde..b79a997 100644
--- a/dns/client.go
+++ b/dns/client.go
@@ -16,6 +16,7 @@ func (c *client) Exchange(m *D.Msg) (msg *D.Msg, err error) {
 }
 
 func (c *client) ExchangeContext(ctx context.Context, m *D.Msg) (msg *D.Msg, err error) {
+	// Please note that miekg/dns ExchangeContext doesn't respond to context cancel.
 	msg, _, err = c.Client.ExchangeContext(ctx, m, c.Address)
 	return
 }
diff --git a/dns/resolver.go b/dns/resolver.go
index db4d7d0..6a9eeeb 100644
--- a/dns/resolver.go
+++ b/dns/resolver.go
@@ -180,7 +180,7 @@ func (r *Resolver) IsFakeIP(ip net.IP) bool {
 }
 
 func (r *Resolver) batchExchange(clients []resolver, m *D.Msg) (msg *D.Msg, err error) {
-	fast, ctx := picker.WithTimeout(context.Background(), time.Second * 5)
+	fast, ctx := picker.WithTimeout(context.Background(), time.Second*5)
 	for _, client := range clients {
 		r := client
 		fast.Go(func() (interface{}, error) {
@@ -192,7 +192,7 @@ func (r *Resolver) batchExchange(clients []resolver, m *D.Msg) (msg *D.Msg, err
 		})
 	}
 
-	elm := fast.Wait()
+	elm := fast.WaitWithoutCancel()
 	if elm == nil {
 		return nil, errors.New("All DNS requests failed")
 	}
diff --git a/dns/util.go b/dns/util.go
index cb6a1c4..d5cb216 100644
--- a/dns/util.go
+++ b/dns/util.go
@@ -134,6 +134,7 @@ func transform(servers []NameServer) []resolver {
 					NextProtos: []string{"dns"},
 				},
 				UDPSize: 4096,
+				Timeout: 5 * time.Second,
 			},
 			Address: s.Addr,
 		})