diff --git a/handler/dns/handler.go b/handler/dns/handler.go index c99436b7..6e42d911 100644 --- a/handler/dns/handler.go +++ b/handler/dns/handler.go @@ -132,6 +132,18 @@ func (h *dnsHandler) Init(md md.Metadata) (err error) { return } +// packResponse serializes a DNS message using a pooled buffer. +// The buffer is returned to the pool if serialization fails. +func (h *dnsHandler) packResponse(mr *dns.Msg) ([]byte, error) { + b := bufpool.Get(h.md.bufferSize) + reply, err := mr.PackBuffer(b) + if err != nil { + bufpool.Put(b) + return nil, err + } + return reply, nil +} + // Forward implements handler.Forwarder. func (h *dnsHandler) Forward(hop hop.Hop) { h.hop = hop @@ -214,6 +226,9 @@ func (h *dnsHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler. } defer bufpool.Put(reply) + if h.md.readTimeout > 0 { + conn.SetWriteDeadline(time.Now().Add(h.md.readTimeout)) + } if _, err = conn.Write(reply); err != nil { log.Error(err) return err @@ -273,15 +288,13 @@ func (h *dnsHandler) request(ctx context.Context, msg []byte, ro *xrecorder.Hand if h.options.Bypass.Contains(context.Background(), "udp", strings.Trim(mq.Question[0].Name, "."), bypass.WithService(h.options.Service)) { log.Debug("bypass: ", mq.Question[0].Name) mr = (&dns.Msg{}).SetReply(&mq) - b := bufpool.Get(h.md.bufferSize) - return mr.PackBuffer(b) + return h.packResponse(mr) } } mr = h.lookupHosts(ctx, &mq, log) if mr != nil { - b := bufpool.Get(h.md.bufferSize) - return mr.PackBuffer(b) + return h.packResponse(mr) } // only cache for single question message. @@ -294,8 +307,7 @@ func (h *dnsHandler) request(ctx context.Context, msg []byte, ro *xrecorder.Hand ro.DNS.Cached = true log.Debugf("message %d (cached): %s", mq.Id, mq.Question[0].String()) - b := bufpool.Get(h.md.bufferSize) - return mr.PackBuffer(b) + return h.packResponse(mr) } } } @@ -307,15 +319,18 @@ func (h *dnsHandler) request(ctx context.Context, msg []byte, ro *xrecorder.Hand ro.Host = ex.String() if mr != nil && h.md.async { - b := bufpool.Get(h.md.bufferSize) - reply, err := mr.PackBuffer(b) + reply, err := h.packResponse(mr) if err != nil { return nil, err } h.cache.RefreshTTL(resolver_util.NewCacheKey(&mq.Question[0])) log.Debugf("exchange message %d (async): %s", mq.Id, mq.Question[0].String()) - go h.exchange(context.WithoutCancel(ctx), ex, &mq) + go func() { + if _, err := h.exchange(context.WithoutCancel(ctx), ex, &mq); err != nil { + log.Debugf("async exchange for %s: %v", mq.Question[0].Name, err) + } + }() return reply, nil } @@ -328,8 +343,7 @@ func (h *dnsHandler) request(ctx context.Context, msg []byte, ro *xrecorder.Hand return nil, err } - b := bufpool.Get(h.md.bufferSize) - return mr.PackBuffer(b) + return h.packResponse(mr) } func (h *dnsHandler) exchange(ctx context.Context, ex exchanger.Exchanger, mq *dns.Msg) (*dns.Msg, error) { @@ -358,7 +372,7 @@ func (h *dnsHandler) exchange(ctx context.Context, ex exchanger.Exchanger, mq *d return mr, nil } -// lookup host mapper +// lookupHosts checks the host mapper for A/AAAA records matching the query. func (h *dnsHandler) lookupHosts(ctx context.Context, r *dns.Msg, log logger.Logger) (m *dns.Msg) { if h.hostMapper == nil || r.Question[0].Qclass != dns.ClassINET || @@ -366,45 +380,31 @@ func (h *dnsHandler) lookupHosts(ctx context.Context, r *dns.Msg, log logger.Log return nil } - m = &dns.Msg{} - m.SetReply(r) - - host := strings.TrimSuffix(r.Question[0].Name, ".") - + var network, rrType string switch r.Question[0].Qtype { case dns.TypeA: - ips, _ := h.hostMapper.Lookup(ctx, "ip4", host) - if len(ips) == 0 { - return nil - } - log.Debugf("hit host mapper: %s -> %s", host, ips) - - for _, ip := range ips { - rr, err := dns.NewRR(fmt.Sprintf("%s IN A %s\n", r.Question[0].Name, ip.String())) - if err != nil { - log.Error(err) - return nil - } - m.Answer = append(m.Answer, rr) - } - + network, rrType = "ip4", "IN A" case dns.TypeAAAA: - ips, _ := h.hostMapper.Lookup(ctx, "ip6", host) - if len(ips) == 0 { - return nil - } - log.Debugf("hit host mapper: %s -> %s", host, ips) - - for _, ip := range ips { - rr, err := dns.NewRR(fmt.Sprintf("%s IN AAAA %s\n", r.Question[0].Name, ip.String())) - if err != nil { - log.Error(err) - return nil - } - m.Answer = append(m.Answer, rr) - } + network, rrType = "ip6", "IN AAAA" } + host := strings.TrimSuffix(r.Question[0].Name, ".") + ips, _ := h.hostMapper.Lookup(ctx, network, host) + if len(ips) == 0 { + return nil + } + log.Debugf("hit host mapper: %s -> %s", host, ips) + + m = new(dns.Msg) + m.SetReply(r) + for _, ip := range ips { + rr, err := dns.NewRR(fmt.Sprintf("%s %s %s\n", r.Question[0].Name, rrType, ip.String())) + if err != nil { + log.Error(err) + return nil + } + m.Answer = append(m.Answer, rr) + } return } diff --git a/handler/dns/handler_test.go b/handler/dns/handler_test.go index ef4dc280..c4be22e7 100644 --- a/handler/dns/handler_test.go +++ b/handler/dns/handler_test.go @@ -916,9 +916,8 @@ func TestHandle_ReadTimeout(t *testing.T) { func TestRequest_AsyncRefresh(t *testing.T) { var ( - mu sync.Mutex - exchanged bool - answerRR = mustNewRR("async.example.com. 300 IN A 10.0.0.1") + answerRR = mustNewRR("async.example.com. 300 IN A 10.0.0.1") + done = make(chan struct{}) ) h := newInitdHandler() @@ -938,9 +937,7 @@ func TestRequest_AsyncRefresh(t *testing.T) { // Set up exchanger for async refresh mockEx := &mockExchanger{ exchangeFn: func(ctx context.Context, msg []byte) ([]byte, error) { - mu.Lock() - exchanged = true - mu.Unlock() + defer close(done) mq := new(dns.Msg) _ = mq.Unpack(msg) @@ -970,13 +967,11 @@ func TestRequest_AsyncRefresh(t *testing.T) { t.Errorf("len(Answer) = %d, want 1 (cached)", len(mr.Answer)) } - // Wait for async goroutine - time.Sleep(100 * time.Millisecond) - mu.Lock() - wasExchanged := exchanged - mu.Unlock() - if !wasExchanged { - t.Error("expected async exchange to occur") + // Wait for async goroutine via channel synchronization. + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("async exchange did not occur within timeout") } } @@ -1019,3 +1014,110 @@ func TestHandle_ConcurrentRequests(t *testing.T) { } wg.Wait() } + +// --------------------------------------------------------------------------- +// Tests: async error handling +// --------------------------------------------------------------------------- + +func TestRequest_AsyncExchangeError(t *testing.T) { + answerRR := mustNewRR("async-err.example.com. 300 IN A 10.0.0.1") + + h := newInitdHandler() + h.md.async = true + h.md.ttl = 30 * time.Second + + // Pre-populate cache with expired entry. + q := new(dns.Msg) + q.SetQuestion("async-err.example.com.", dns.TypeA) + respMsg := q.Copy() + respMsg.Answer = []dns.RR{answerRR} + key := resolver_util.NewCacheKey(&q.Question[0]) + h.cache.Store(context.Background(), key, respMsg, time.Nanosecond) + time.Sleep(5 * time.Millisecond) + + done := make(chan struct{}) + mockEx := &mockExchanger{ + exchangeFn: func(ctx context.Context, msg []byte) ([]byte, error) { + defer close(done) + return nil, errors.New("upstream unreachable") + }, + addr: "udp://8.8.8.8:53", + } + + h.hop = &mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return chain.NewNode("ns1", "udp://8.8.8.8:53") + }, + } + h.exchangers["ns1"] = mockEx + + // The stale response should still be returned despite async error. + reply, err := h.request(context.Background(), packDNSQuery("async-err.example.com.", dns.TypeA), newRecObj(), nopLog()) + if err != nil { + t.Fatalf("request: %v", err) + } + mr := new(dns.Msg) + if err := mr.Unpack(reply); err != nil { + t.Fatalf("unpack: %v", err) + } + if len(mr.Answer) != 1 { + t.Errorf("len(Answer) = %d, want 1 (stale cached)", len(mr.Answer)) + } + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("async exchange did not complete within timeout") + } +} + +// --------------------------------------------------------------------------- +// Tests: write deadline +// --------------------------------------------------------------------------- + +func TestHandle_WriteDeadline(t *testing.T) { + answerRR := mustNewRR("timeout.example.com. 60 IN A 10.0.0.1") + mockEx := &mockExchanger{ + exchangeFn: func(ctx context.Context, msg []byte) ([]byte, error) { + mq := new(dns.Msg) + _ = mq.Unpack(msg) + resp := mq.Copy() + resp.Answer = []dns.RR{answerRR} + return resp.Pack() + }, + addr: "udp://8.8.8.8:53", + } + + h := newTestHandler() + h.md.readTimeout = 50 * time.Millisecond + h.md.timeout = defaultTimeout + h.md.bufferSize = defaultBufferSize + h.cache = resolver_util.NewCache().WithLogger(nopLog()) + h.hop = &mockHop{ + selectFn: func(ctx context.Context, opts ...hop.SelectOption) *chain.Node { + return chain.NewNode("ns1", "udp://8.8.8.8:53") + }, + } + h.exchangers["ns1"] = mockEx + + conn := &blockingWriteConn{ + readBuf: bytes.NewBuffer(packDNSQuery("timeout.example.com.", dns.TypeA)), + } + if err := h.Handle(context.Background(), conn); err == nil { + t.Fatal("expected error from blocked write") + } +} + +// blockingWriteConn is a net.Conn where Write blocks forever. +type blockingWriteConn struct { + readBuf *bytes.Buffer +} + +func (c *blockingWriteConn) Read(b []byte) (int, error) { return c.readBuf.Read(b) } +func (c *blockingWriteConn) Write(b []byte) (int, error) { time.Sleep(200 * time.Millisecond); return 0, errors.New("write deadline exceeded") } +func (c *blockingWriteConn) Close() error { return nil } +func (c *blockingWriteConn) LocalAddr() net.Addr { return &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 53} } +func (c *blockingWriteConn) RemoteAddr() net.Addr { return &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 12345} } +func (c *blockingWriteConn) SetDeadline(t time.Time) error { return nil } +func (c *blockingWriteConn) SetReadDeadline(t time.Time) error { return nil } +func (c *blockingWriteConn) SetWriteDeadline(t time.Time) error { return nil } diff --git a/handler/dns/metadata.go b/handler/dns/metadata.go index f1d116f0..3acf78a2 100644 --- a/handler/dns/metadata.go +++ b/handler/dns/metadata.go @@ -13,7 +13,7 @@ const ( // defaultTimeout is the fallback timeout for DNS exchanges. defaultTimeout = 5 * time.Second // defaultBufferSize is the fallback buffer size for DNS message I/O. - defaultBufferSize = 1024 + defaultBufferSize = 4096 ) // metadata holds parsed DNS handler configuration.