From 5b1e38e2c0ff6b335a872ebf373a4eedf7ac9cf0 Mon Sep 17 00:00:00 2001
From: wenyifan <yifan.wen@eisgroup.com>
Date: Sat, 4 Jul 2026 13:48:22 +0800
Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9Ednst=E9=9A=A7=E9=81=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/util/dns/conn.go      | 34 ++++++++++++++++++++++++
 internal/util/dns/conn_test.go | 48 ++++++++++++++++++++++++++++++++++
 listener/dns/tunnel.go         | 25 ++++++++++++++++++
 3 files changed, 107 insertions(+)

diff --git a/internal/util/dns/conn.go b/internal/util/dns/conn.go
index ebb2037f..b4891a1f 100644
--- a/internal/util/dns/conn.go
+++ b/internal/util/dns/conn.go
@@ -21,6 +21,7 @@ const (
 var (
 	ErrInvalidPacket = errors.New("dns: invalid packet")
 	ErrShortBuffer   = errors.New("dns: short buffer")
+	ErrUnsupported   = errors.New("dns: unsupported operation")
 )
 
 type options struct {
@@ -28,6 +29,11 @@ type options struct {
 	udpSize uint16
 }
 
+type setBuffer interface {
+	SetReadBuffer(int) error
+	SetWriteBuffer(int) error
+}
+
 // Option configures a DNS packet wrapper.
 type Option func(*options)
 
@@ -117,6 +123,20 @@ func (c *clientConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	return len(b), nil
 }
 
+func (c *clientConn) SetReadBuffer(n int) error {
+	if pc, ok := c.PacketConn.(setBuffer); ok {
+		return pc.SetReadBuffer(n)
+	}
+	return ErrUnsupported
+}
+
+func (c *clientConn) SetWriteBuffer(n int) error {
+	if pc, ok := c.PacketConn.(setBuffer); ok {
+		return pc.SetWriteBuffer(n)
+	}
+	return ErrUnsupported
+}
+
 type queryState struct {
 	id       uint16
 	question []mdns.Question
@@ -184,6 +204,20 @@ func (c *serverConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	return len(b), nil
 }
 
+func (c *serverConn) SetReadBuffer(n int) error {
+	if pc, ok := c.PacketConn.(setBuffer); ok {
+		return pc.SetReadBuffer(n)
+	}
+	return ErrUnsupported
+}
+
+func (c *serverConn) SetWriteBuffer(n int) error {
+	if pc, ok := c.PacketConn.(setBuffer); ok {
+		return pc.SetWriteBuffer(n)
+	}
+	return ErrUnsupported
+}
+
 func (c *serverConn) storeQuery(addr net.Addr, m *mdns.Msg) {
 	if addr == nil {
 		return
diff --git a/internal/util/dns/conn_test.go b/internal/util/dns/conn_test.go
index 44120749..91c55e09 100644
--- a/internal/util/dns/conn_test.go
+++ b/internal/util/dns/conn_test.go
@@ -59,3 +59,51 @@ func TestPacketConnRoundTrip(t *testing.T) {
 		t.Fatalf("client ReadFrom got %q, want %q", got, responsePayload)
 	}
 }
+
+func TestPacketConnSetBuffer(t *testing.T) {
+	base := &bufferPacketConn{}
+
+	client := ClientConn(base)
+	if err := client.(setBuffer).SetReadBuffer(1024); err != nil {
+		t.Fatal(err)
+	}
+	if err := client.(setBuffer).SetWriteBuffer(2048); err != nil {
+		t.Fatal(err)
+	}
+	if base.readBuffer != 1024 {
+		t.Fatalf("read buffer = %d, want 1024", base.readBuffer)
+	}
+	if base.writeBuffer != 2048 {
+		t.Fatalf("write buffer = %d, want 2048", base.writeBuffer)
+	}
+
+	server := ServerConn(base)
+	if err := server.(setBuffer).SetReadBuffer(4096); err != nil {
+		t.Fatal(err)
+	}
+	if err := server.(setBuffer).SetWriteBuffer(8192); err != nil {
+		t.Fatal(err)
+	}
+	if base.readBuffer != 4096 {
+		t.Fatalf("read buffer = %d, want 4096", base.readBuffer)
+	}
+	if base.writeBuffer != 8192 {
+		t.Fatalf("write buffer = %d, want 8192", base.writeBuffer)
+	}
+}
+
+type bufferPacketConn struct {
+	net.PacketConn
+	readBuffer  int
+	writeBuffer int
+}
+
+func (c *bufferPacketConn) SetReadBuffer(n int) error {
+	c.readBuffer = n
+	return nil
+}
+
+func (c *bufferPacketConn) SetWriteBuffer(n int) error {
+	c.writeBuffer = n
+	return nil
+}
diff --git a/listener/dns/tunnel.go b/listener/dns/tunnel.go
index 41b7e70f..a6371d8e 100644
--- a/listener/dns/tunnel.go
+++ b/listener/dns/tunnel.go
@@ -3,6 +3,7 @@ package dns
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"net"
 	"strings"
 
@@ -65,6 +66,7 @@ func (l *dnsTunnelListener) Init(md md.Metadata) (err error) {
 	if err != nil {
 		return
 	}
+	bufferConn := conn
 
 	conn = dns_util.ServerConn(
 		conn,
@@ -82,6 +84,10 @@ func (l *dnsTunnelListener) Init(md md.Metadata) (err error) {
 		limiter.ServiceOption(l.options.Service),
 		limiter.NetworkOption(conn.LocalAddr().Network()),
 	)
+	conn = &setBufferPacketConn{
+		PacketConn: conn,
+		bufferConn: bufferConn,
+	}
 
 	config := &quic.Config{
 		KeepAlivePeriod:      l.md.keepAlivePeriod,
@@ -203,3 +209,22 @@ func (c *dnsTunnelConn) LocalAddr() net.Addr {
 func (c *dnsTunnelConn) RemoteAddr() net.Addr {
 	return c.raddr
 }
+
+type setBufferPacketConn struct {
+	net.PacketConn
+	bufferConn net.PacketConn
+}
+
+func (c *setBufferPacketConn) SetReadBuffer(n int) error {
+	if pc, ok := c.bufferConn.(interface{ SetReadBuffer(int) error }); ok {
+		return pc.SetReadBuffer(n)
+	}
+	return errors.New("dns: set read buffer unsupported")
+}
+
+func (c *setBufferPacketConn) SetWriteBuffer(n int) error {
+	if pc, ok := c.bufferConn.(interface{ SetWriteBuffer(int) error }); ok {
+		return pc.SetWriteBuffer(n)
+	}
+	return errors.New("dns: set write buffer unsupported")
+}