From e70b37865fe00b26d0e06001dff3258f656ba396 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Sat, 24 Feb 2024 21:30:19 +0800
Subject: [PATCH] Improve DNS truncate behavior

---
 outbound/dns.go | 24 +++---------------------
 1 file changed, 3 insertions(+), 21 deletions(-)

diff --git a/outbound/dns.go b/outbound/dns.go
index df32a019..426552ac 100644
--- a/outbound/dns.go
+++ b/outbound/dns.go
@@ -241,8 +241,9 @@ func (d *DNS) newPacketConnection(ctx context.Context, conn N.PacketConn, readWa
 					return err
 				}
 				timeout.Update()
-				response = truncateDNSMessage(response, 512) // TODO: add an option to custom UDP buffer size
-				responseBuffer := buf.NewSize(dns.FixedPacketSize)
+				var responseLen int
+				response, responseLen = dns.TruncateDNSMessage(&message, response)
+				responseBuffer := buf.NewSize(1024 + responseLen)
 				responseBuffer.Resize(1024, 0)
 				n, err := response.PackBuffer(responseBuffer.FreeBytes())
 				if err != nil {
@@ -264,22 +265,3 @@ func (d *DNS) newPacketConnection(ctx context.Context, conn N.PacketConn, readWa
 	})
 	return group.Run(fastClose)
 }
-
-func truncateDNSMessage(response *mDNS.Msg, maxLen int) *mDNS.Msg {
-	responseLen := response.Len()
-	if responseLen <= maxLen {
-		return response
-	}
-	newResponse := *response
-	response = &newResponse
-	for len(response.Answer) > 0 && responseLen > maxLen {
-		response.Answer = response.Answer[:len(response.Answer)-1]
-		response.Truncated = true
-		responseLen = response.Len()
-	}
-	if responseLen > maxLen {
-		response.Ns = nil
-		response.Extra = nil
-	}
-	return response
-}