fix(handler/dns): buffer pool leak, write deadline, async error logging, EDNS0 buffer

Fix bufpool leak when PackBuffer fails by extracting packResponse helper.
Add write deadline before conn.Write to prevent goroutine hangs on slow
clients. Log async exchange errors instead of silently discarding them.
Raise defaultBufferSize from 1024 to 4096 to match EDNS0 standard.
Deduplicate lookupHosts TypeA/TypeAAAA branches.

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 h.packResponse(mr)
-			return mr.PackBuffer(b)
 		}
 	}
 
 	mr = h.lookupHosts(ctx, &mq, log)
 	if mr != nil {
-		b := bufpool.Get(h.md.bufferSize)
+		return h.packResponse(mr)
-		return mr.PackBuffer(b)
 	}
 
 	// 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 h.packResponse(mr)
-				return mr.PackBuffer(b)
 			}
 		}
 	}
@@ -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 := h.packResponse(mr)
-		reply, err := mr.PackBuffer(b)
 		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 h.packResponse(mr)
-	return mr.PackBuffer(b)
 }
 
 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{}
+	var network, rrType string
-	m.SetReply(r)
-
-	host := strings.TrimSuffix(r.Question[0].Name, ".")
-
 	switch r.Question[0].Qtype {
 	case dns.TypeA:
-		ips, _ := h.hostMapper.Lookup(ctx, "ip4", host)
+		network, rrType = "ip4", "IN A"
-		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)
-		}
-
 	case dns.TypeAAAA:
-		ips, _ := h.hostMapper.Lookup(ctx, "ip6", host)
+		network, rrType = "ip6", "IN AAAA"
-		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)
-		}
 	}
 
+	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
 }
 

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
+		answerRR = mustNewRR("async.example.com. 300 IN A 10.0.0.1")
-		exchanged  bool
+		done     = make(chan struct{})
-		answerRR   = mustNewRR("async.example.com. 300 IN A 10.0.0.1")
 	)
 
 	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()
+			defer close(done)
-			exchanged = true
-			mu.Unlock()
 
 			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
+	// Wait for async goroutine via channel synchronization.
-	time.Sleep(100 * time.Millisecond)
+	select {
-	mu.Lock()
+	case <-done:
-	wasExchanged := exchanged
+	case <-time.After(2 * time.Second):
-	mu.Unlock()
+		t.Fatal("async exchange did not occur within timeout")
-	if !wasExchanged {
-		t.Error("expected async exchange to occur")
 	}
 }
 
@@ -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 }

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.