From 73373aeb3f96a7a11d7dc7ff8211cc1d691ce1db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 14 Feb 2024 20:42:58 +0800
Subject: [PATCH] Improve DNS truncate behavior

---
 outbound/dns.go | 30 ++----------------------------
 1 file changed, 2 insertions(+), 28 deletions(-)

diff --git a/outbound/dns.go b/outbound/dns.go
index df32a019..2cc11d3a 100644
--- a/outbound/dns.go
+++ b/outbound/dns.go
@@ -111,9 +111,7 @@ func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metada
 			}
 		}
 		if readWaiter, created := bufio.CreatePacketReadWaiter(reader); created {
-			readWaiter.InitializeReadWaiter(N.ReadWaitOptions{
-				MTU: dns.FixedPacketSize,
-			})
+			readWaiter.InitializeReadWaiter(N.ReadWaitOptions{})
 			return d.newPacketConnection(ctx, conn, readWaiter, counters, cachedPackets, metadata)
 		}
 		break
@@ -241,16 +239,11 @@ 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)
-				responseBuffer.Resize(1024, 0)
-				n, err := response.PackBuffer(responseBuffer.FreeBytes())
+				responseBuffer, err := dns.TruncateDNSMessage(&message, response, 1024)
 				if err != nil {
 					cancel(err)
-					responseBuffer.Release()
 					return err
 				}
-				responseBuffer.Truncate(len(n))
 				err = conn.WritePacket(responseBuffer, destination)
 				if err != nil {
 					cancel(err)
@@ -264,22 +257,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
-}