From 17aebc56c1b72696aa379a4806a6b73eb4fa3f85 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Thu, 1 Feb 2024 12:10:39 +0800
Subject: [PATCH] Fix UDP DNS response not truncated

---
 outbound/dns.go | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/outbound/dns.go b/outbound/dns.go
index fcb67d45..0f003377 100644
--- a/outbound/dns.go
+++ b/outbound/dns.go
@@ -241,7 +241,8 @@ func (d *DNS) newPacketConnection(ctx context.Context, conn N.PacketConn, readWa
 					return err
 				}
 				timeout.Update()
-				responseBuffer := buf.NewPacket()
+				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())
 				if err != nil {
@@ -263,3 +264,21 @@ 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
+	}
+	response = response.Copy()
+	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
+}