From c3de4b8392d06b27199cbf13ccde584b8027fb22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=A3=8E=E6=89=87=E6=BB=91=E7=BF=94=E7=BF=BC?=
 <Fangliding.fshxy@outlook.com>
Date: Tue, 18 Mar 2025 19:47:02 +0000
Subject: [PATCH] DNS: Retry with EDNS0 when response is truncated

---
 app/dns/dnscommon.go      | 11 ++++++++---
 app/dns/nameserver_udp.go | 21 +++++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/app/dns/dnscommon.go b/app/dns/dnscommon.go
index 42f9da2c..5f3d08ac 100644
--- a/app/dns/dnscommon.go
+++ b/app/dns/dnscommon.go
@@ -35,6 +35,9 @@ type IPRecord struct {
 	IP     []net.Address
 	Expire time.Time
 	RCode  dnsmessage.RCode
+
+	// Truncated is for udp dns to indicates if the response is truncated and needs to be retried
+	Truncated bool
 }
 
 func (r *IPRecord) getIPs() ([]net.Address, error) {
@@ -65,6 +68,7 @@ type dnsRequest struct {
 	start   time.Time
 	expire  time.Time
 	msg     *dnsmessage.Message
+	ctx     context.Context
 }
 
 func genEDNS0Options(clientIP net.IP) *dnsmessage.Resource {
@@ -179,9 +183,10 @@ func parseResponse(payload []byte) (*IPRecord, error) {
 
 	now := time.Now()
 	ipRecord := &IPRecord{
-		ReqID:  h.ID,
-		RCode:  h.RCode,
-		Expire: now.Add(time.Second * 600),
+		ReqID:     h.ID,
+		RCode:     h.RCode,
+		Expire:    now.Add(time.Second * 600),
+		Truncated: h.Truncated,
 	}
 
 L:
diff --git a/app/dns/nameserver_udp.go b/app/dns/nameserver_udp.go
index 47d70435..a088eef4 100644
--- a/app/dns/nameserver_udp.go
+++ b/app/dns/nameserver_udp.go
@@ -128,6 +128,26 @@ func (s *ClassicNameServer) HandleResponse(ctx context.Context, packet *udp_prot
 		return
 	}
 
+	// if truncated, retry with EDNS0 option(udp payload size: 1350)
+	if ipRec.Truncated {
+		// if already has EDNS0 option, no need to retry
+		if ok && len(req.msg.Additionals) == 0 {
+			// copy necessary meta data from original request
+			// and add EDNS0 option
+			opt := new(dnsmessage.Resource)
+			common.Must(opt.Header.SetEDNS0(1350, 0xfe00, true))
+			newMsg := *req.msg
+			newReq := *req
+			newMsg.Additionals = append(newMsg.Additionals, *opt)
+			newMsg.ID = s.newReqID()
+			newReq.msg = &newMsg
+			s.addPendingRequest(&newReq)
+			b, _ := dns.PackMessage(req.msg)
+			s.udpServer.Dispatch(toDnsContext(newReq.ctx, s.address.String()), *s.address, b)
+			return
+		}
+	}
+
 	var rec record
 	switch req.reqType {
 	case dnsmessage.TypeA:
@@ -194,6 +214,7 @@ func (s *ClassicNameServer) sendQuery(ctx context.Context, domain string, client
 	reqs := buildReqMsgs(domain, option, s.newReqID, genEDNS0Options(clientIP))
 
 	for _, req := range reqs {
+		req.ctx = ctx
 		s.addPendingRequest(req)
 		b, _ := dns.PackMessage(req.msg)
 		s.udpServer.Dispatch(toDnsContext(ctx, s.address.String()), *s.address, b)