From f6b52a653abd0026ea4349a73b9e7360c48e20bf Mon Sep 17 00:00:00 2001 From: Wesley Squasher Date: Sun, 12 May 2019 08:03:48 +0000 Subject: [PATCH] Use TCP when appropriate for the given query type/response --- doh-server/server.go | 25 ++++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/doh-server/server.go b/doh-server/server.go index e758767..b8d61d5 100644 --- a/doh-server/server.go +++ b/doh-server/server.go @@ -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