From 727f9e5654cc0afaae4b7e71ba67dc7784be6dd5 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 5 May 2023 12:53:55 +0300
Subject: [PATCH] introduce a OptimizeConn function to manually enable UDP
 optimizations

---
 sys_conn.go         | 19 ++++++++++++++++++-
 sys_conn_no_oob.go  |  2 +-
 sys_conn_windows.go |  2 +-
 transport.go        | 21 +++++++++++++--------
 4 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/sys_conn.go b/sys_conn.go
index 59b730fab..72f84b0a4 100644
--- a/sys_conn.go
+++ b/sys_conn.go
@@ -15,13 +15,30 @@ import (
 type OOBCapablePacketConn interface {
 	net.PacketConn
 	SyscallConn() (syscall.RawConn, error)
+	SetReadBuffer(int) error
 	ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
 	WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
 }
 
 var _ OOBCapablePacketConn = &net.UDPConn{}
 
-func wrapConn(pc net.PacketConn) (rawConn, error) {
+// OptimizeConn takes a net.PacketConn and attempts to enable various optimizations that will improve QUIC performance:
+//  1. It enables the Don't Fragment (DF) bit on the IP header.
+//     This allows us to do DPLPMTUD (Path MTU Discovery).
+//  2. It enables reading of the ECN bits from the IP header.
+//     This allows the remote node to speed up its loss detection and recovery.
+//  3. It uses batched syscalls (recvmmsg) to more efficiently receive packets from the socket.
+//
+// For this to work, the connection needs to implement the OOBCapablePacketConn interface (as a *net.UDPConn does).
+func OptimizeConn(c net.PacketConn) (net.PacketConn, error) {
+	return wrapConn(c)
+}
+
+func wrapConn(pc net.PacketConn) (interface {
+	net.PacketConn
+	rawConn
+}, error,
+) {
 	conn, ok := pc.(interface {
 		SyscallConn() (syscall.RawConn, error)
 	})
diff --git a/sys_conn_no_oob.go b/sys_conn_no_oob.go
index 997775c91..106ca975a 100644
--- a/sys_conn_no_oob.go
+++ b/sys_conn_no_oob.go
@@ -4,7 +4,7 @@ package quic
 
 import "net"
 
-func newConn(c net.PacketConn, supportsDF bool) (rawConn, error) {
+func newConn(c net.PacketConn, supportsDF bool) (*basicConn, error) {
 	return &basicConn{PacketConn: c, supportsDF: supportsDF}, nil
 }
 
diff --git a/sys_conn_windows.go b/sys_conn_windows.go
index 8a583fdaa..2ee9490cf 100644
--- a/sys_conn_windows.go
+++ b/sys_conn_windows.go
@@ -8,7 +8,7 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-func newConn(c OOBCapablePacketConn, supportsDF bool) (rawConn, error) {
+func newConn(c OOBCapablePacketConn, supportsDF bool) (*basicConn, error) {
 	return &basicConn{PacketConn: c, supportsDF: supportsDF}, nil
 }
 
diff --git a/transport.go b/transport.go
index dee524f29..fcd63e62d 100644
--- a/transport.go
+++ b/transport.go
@@ -30,10 +30,8 @@ type Transport struct {
 	// A single net.PacketConn can only be handled by one Transport.
 	// Bad things will happen if passed to multiple Transports.
 	//
-	// If the connection satisfies the OOBCapablePacketConn interface
-	// (as a net.UDPConn does), ECN and packet info support will be enabled.
-	// In this case, optimized syscalls might be used, skipping the
-	// ReadFrom and WriteTo calls to read / write packets.
+	// If not done by the user, the connection is passed through OptimizeConn to enable a number of optimizations.
+	// After passing the connection to the Transport, its invalid to call ReadFrom and WriteTo.
 	Conn net.PacketConn
 
 	// The length of the connection ID in bytes.
@@ -180,11 +178,18 @@ func (t *Transport) init(isServer bool) error {
 	t.initOnce.Do(func() {
 		getMultiplexer().AddConn(t.Conn)
 
-		conn, err := wrapConn(t.Conn)
-		if err != nil {
-			t.initErr = err
-			return
+		var conn rawConn
+		if c, ok := t.Conn.(rawConn); ok {
+			conn = c
+		} else {
+			var err error
+			conn, err = wrapConn(t.Conn)
+			if err != nil {
+				t.initErr = err
+				return
+			}
 		}
+		t.conn = conn
 
 		t.logger = utils.DefaultLogger // TODO: make this configurable
 		t.conn = conn