Fix: use the fastest whether the result is successful

dns/client.go

@@ -20,7 +20,22 @@ 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) {
 	c.Client.Dialer = dialer.Dialer()
 
-	// Please note that miekg/dns ExchangeContext doesn't respond to context cancel.
-	msg, _, err = c.Client.ExchangeContext(ctx, m, c.Address)
-	return
+	// miekg/dns ExchangeContext doesn't respond to context cancel.
+	// this is a workaround
+	type result struct {
+		msg *D.Msg
+		err error
+	}
+	ch := make(chan result, 1)
+	go func() {
+		msg, _, err := c.Client.ExchangeContext(ctx, m, c.Address)
+		ch <- result{msg, err}
+	}()
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case ret := <-ch:
+		return ret.msg, ret.err
+	}
 }

dns/middleware.go

@@ -39,7 +39,8 @@ func withFakeIP(fakePool *fakeip.Pool) middleware {
 			msg.Answer = []D.RR{rr}
 
 			setMsgTTL(msg, 1)
-			msg.SetReply(r)
+			msg.SetRcode(r, msg.Rcode)
+			msg.Authoritative = true
 			w.WriteMsg(msg)
 			return
 		}
@@ -55,7 +56,8 @@ func withResolver(resolver *Resolver) handler {
 			D.HandleFailed(w, r)
 			return
 		}
-		msg.SetReply(r)
+		msg.SetRcode(r, msg.Rcode)
+		msg.Authoritative = true
 		w.WriteMsg(msg)
 		return
 	}

dns/resolver.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
+	"math/rand"
 	"net"
 	"strings"
 	"time"
@@ -178,15 +179,11 @@ func (r *Resolver) batchExchange(clients []resolver, m *D.Msg) (msg *D.Msg, err
 	for _, client := range clients {
 		r := client
 		fast.Go(func() (interface{}, error) {
-			msg, err := r.ExchangeContext(ctx, m)
-			if err != nil || msg.Rcode != D.RcodeSuccess {
-				return nil, errors.New("resolve error")
-			}
-			return msg, nil
+			return r.ExchangeContext(ctx, m)
 		})
 	}
 
-	elm := fast.WaitWithoutCancel()
+	elm := fast.Wait()
 	if elm == nil {
 		return nil, errors.New("All DNS requests failed")
 	}
@@ -239,11 +236,12 @@ func (r *Resolver) resolveIP(host string, dnsType uint16) (ip net.IP, err error)
 	}
 
 	ips := r.msgToIP(msg)
-	if len(ips) == 0 {
+	ipLength := len(ips)
+	if ipLength == 0 {
 		return nil, errIPNotFound
 	}
 
-	ip = ips[0]
+	ip = ips[rand.Intn(ipLength)]
 	return
 }