From be2b06f8e16a0c55e891fc6b2f1714f713b3775f Mon Sep 17 00:00:00 2001
From: database64128 <free122448@hotmail.com>
Date: Mon, 30 Sep 2024 11:28:12 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=A1=20conn:=20add=20option=20for=20pro?=
 =?UTF-8?q?bing=20UDP=20GSO=20support?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 conn/conn.go                               |  6 ++++++
 conn/conn_linux.go                         | 10 ++++++++++
 conn/{conn_udpgro.go => conn_udpgsogro.go} | 10 ++++++++++
 conn/conn_windows.go                       | 11 +++++++++++
 4 files changed, 37 insertions(+)
 rename conn/{conn_udpgro.go => conn_udpgsogro.go} (54%)

diff --git a/conn/conn.go b/conn/conn.go
index 0c6046d..0daffd4 100644
--- a/conn/conn.go
+++ b/conn/conn.go
@@ -168,6 +168,12 @@ type ListenerSocketOptions struct {
 	// Available on platforms supported by Go std's MPTCP implementation.
 	MultipathTCP bool
 
+	// ProbeUDPGSOSupport enables best-effort probing of
+	// UDP Generic Segmentation Offload (GSO) support on the listener.
+	//
+	// Available on Linux and Windows.
+	ProbeUDPGSOSupport bool
+
 	// UDPGenericReceiveOffload enables UDP Generic Receive Offload (GRO) on the listener.
 	//
 	// Available on Linux and Windows.
diff --git a/conn/conn_linux.go b/conn/conn_linux.go
index 0dc2975..78efde6 100644
--- a/conn/conn_linux.go
+++ b/conn/conn_linux.go
@@ -67,6 +67,15 @@ func setTCPUserTimeout(fd, msecs int) error {
 	return nil
 }
 
+func probeUDPGSOSupport(fd int, info *SocketInfo) {
+	if err := unix.SetsockoptInt(fd, unix.IPPROTO_UDP, unix.UDP_SEGMENT, 0); err == nil {
+		// UDP_MAX_SEGMENTS as defined in linux/udp.h was originally 64.
+		// It got bumped to 128 in Linux 6.9: https://github.com/torvalds/linux/commit/1382e3b6a3500c245e5278c66d210c02926f804f
+		// The receive path still only supports 64 segments, so 64 it is.
+		info.MaxUDPGSOSegments = 64
+	}
+}
+
 func setUDPGenericReceiveOffload(fd int, info *SocketInfo) {
 	if err := unix.SetsockoptInt(fd, unix.IPPROTO_UDP, unix.UDP_GRO, 1); err == nil {
 		info.UDPGenericReceiveOffload = true
@@ -190,6 +199,7 @@ func (lso ListenerSocketOptions) buildSetFns() setFuncSlice {
 		appendSetReusePortFunc(lso.ReusePort).
 		appendSetTransparentFunc(lso.Transparent).
 		appendSetPMTUDFunc(lso.PathMTUDiscovery).
+		appendProbeUDPGSOSupportFunc(lso.ProbeUDPGSOSupport).
 		appendSetUDPGenericReceiveOffloadFunc(lso.UDPGenericReceiveOffload).
 		appendSetRecvPktinfoFunc(lso.ReceivePacketInfo).
 		appendSetRecvOrigDstAddrFunc(lso.ReceiveOriginalDestAddr)
diff --git a/conn/conn_udpgro.go b/conn/conn_udpgsogro.go
similarity index 54%
rename from conn/conn_udpgro.go
rename to conn/conn_udpgsogro.go
index 646c4b0..91438f4 100644
--- a/conn/conn_udpgro.go
+++ b/conn/conn_udpgsogro.go
@@ -2,6 +2,16 @@
 
 package conn
 
+func (fns setFuncSlice) appendProbeUDPGSOSupportFunc(probeUDPGSO bool) setFuncSlice {
+	if probeUDPGSO {
+		return append(fns, func(fd int, _ string, info *SocketInfo) error {
+			probeUDPGSOSupport(fd, info)
+			return nil
+		})
+	}
+	return fns
+}
+
 func (fns setFuncSlice) appendSetUDPGenericReceiveOffloadFunc(gro bool) setFuncSlice {
 	if gro {
 		return append(fns, func(fd int, _ string, info *SocketInfo) error {
diff --git a/conn/conn_windows.go b/conn/conn_windows.go
index 3427456..bf172b0 100644
--- a/conn/conn_windows.go
+++ b/conn/conn_windows.go
@@ -53,6 +53,16 @@ func setPMTUD(fd int, network string) error {
 	return nil
 }
 
+// Implementation inspired by:
+// https://github.com/quinn-rs/quinn/blob/main/quinn-udp/src/windows.rs
+
+func probeUDPGSOSupport(fd int, info *SocketInfo) {
+	if err := windows.SetsockoptInt(windows.Handle(fd), windows.IPPROTO_UDP, windows.UDP_SEND_MSG_SIZE, 0); err == nil {
+		// As "empirically found on Windows 11 x64" by quinn.
+		info.MaxUDPGSOSegments = 512
+	}
+}
+
 func setUDPGenericReceiveOffload(fd int, info *SocketInfo) {
 	// Both quinn and msquic set this to 65535.
 	if err := windows.SetsockoptInt(windows.Handle(fd), windows.IPPROTO_UDP, windows.UDP_RECV_MAX_COALESCED_SIZE, 65535); err == nil {
@@ -84,6 +94,7 @@ func (lso ListenerSocketOptions) buildSetFns() setFuncSlice {
 		appendSetSendBufferSize(lso.SendBufferSize).
 		appendSetRecvBufferSize(lso.ReceiveBufferSize).
 		appendSetPMTUDFunc(lso.PathMTUDiscovery).
+		appendProbeUDPGSOSupportFunc(lso.ProbeUDPGSOSupport).
 		appendSetUDPGenericReceiveOffloadFunc(lso.UDPGenericReceiveOffload).
 		appendSetRecvPktinfoFunc(lso.ReceivePacketInfo)
 }