Use TCP when appropriate for the given query type/response

This commit is contained in:
Wesley Squasher
2019-05-12 08:03:48 +00:00
parent 2332d9b7c1
commit f6b52a653a

View File

@@ -263,19 +263,38 @@ func (s *Server) patchRootRD(req *DNSRequest) *DNSRequest {
return req
}
// Return the position index for the question of qtype from a DNS msg, otherwise return -1
func (s *Server) indexQuestionType(msg *dns.Msg, qtype uint16) int {
for i, question := range msg.Question {
if question.Qtype == qtype {
return i
}
}
return -1
}
func (s *Server) doDNSQuery(ctx context.Context, req *DNSRequest) (resp *DNSRequest, err error) {
// TODO(m13253): Make ctx work. Waiting for a patch for ExchangeContext from miekg/dns.
numServers := len(s.conf.Upstream)
for i := uint(0); i < s.conf.Tries; i++ {
req.currentUpstream = s.conf.Upstream[rand.Intn(numServers)]
if !s.conf.TCPOnly {
// Use TCP if always configured to or if the Query type dictates it (AXFR)
if s.conf.TCPOnly || (s.indexQuestionType(req.request, dns.TypeAXFR) > -1) {
req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
} else {
req.response, _, err = s.udpClient.Exchange(req.request, req.currentUpstream)
if err == nil && req.response != nil && req.response.Truncated {
log.Println(err)
req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
}
} else {
req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
// Retry with TCP if this was an IXFR request and we only received an SOA
if (s.indexQuestionType(req.request, dns.TypeIXFR) > -1) &&
(len(req.response.Answer) == 1) &&
(req.response.Answer[0].Header().Rrtype == dns.TypeSOA) {
req.response, _, err = s.tcpClient.Exchange(req.request, req.currentUpstream)
}
}
if err == nil {
return req, nil