From bd546084738abe65cb980b737958f60f651831ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Tue, 9 May 2023 21:59:44 +0800
Subject: [PATCH] Fix DNS outbound

---
 outbound/dns.go | 49 +++++++++++++++++++++++++++++++++----------------
 1 file changed, 33 insertions(+), 16 deletions(-)

diff --git a/outbound/dns.go b/outbound/dns.go
index 07773ab5..dff5978d 100644
--- a/outbound/dns.go
+++ b/outbound/dns.go
@@ -119,30 +119,47 @@ func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metada
 		}
 		break
 	}
-
 	ctx = adapter.WithContext(ctx, &metadata)
 	fastClose, cancel := common.ContextWithCancelCause(ctx)
 	timeout := canceler.New(fastClose, cancel, C.DNSTimeout)
 	var group task.Group
 	group.Append0(func(ctx context.Context) error {
-		_buffer := buf.StackNewSize(dns.FixedPacketSize)
-		defer common.KeepAlive(_buffer)
-		buffer := common.Dup(_buffer)
-		defer buffer.Release()
 		for {
-			buffer.FullReset()
-			destination, err := conn.ReadPacket(buffer)
-			if err != nil {
-				cancel(err)
-				return err
-			}
 			var message mDNS.Msg
-			err = message.Unpack(buffer.Bytes())
-			if err != nil {
-				cancel(err)
-				return err
+			var destination M.Socksaddr
+			var err error
+			if len(cachedPackets) > 0 {
+				packet := cachedPackets[0]
+				cachedPackets = cachedPackets[1:]
+				for _, counter := range counters {
+					counter(int64(packet.Buffer.Len()))
+				}
+				err = message.Unpack(packet.Buffer.Bytes())
+				packet.Buffer.Release()
+				if err != nil {
+					cancel(err)
+					return err
+				}
+				destination = packet.Destination
+			} else {
+				buffer := buf.NewPacket()
+				destination, err = conn.ReadPacket(buffer)
+				if err != nil {
+					buffer.Release()
+					cancel(err)
+					return err
+				}
+				for _, counter := range counters {
+					counter(int64(buffer.Len()))
+				}
+				err = message.Unpack(buffer.Bytes())
+				buffer.Release()
+				if err != nil {
+					cancel(err)
+					return err
+				}
+				timeout.Update()
 			}
-			timeout.Update()
 			metadataInQuery := metadata
 			go func() error {
 				response, err := d.router.Exchange(adapter.WithContext(ctx, &metadataInQuery), &message)