gso and gro with uring on send/receive for udp

udp/config.go

@@ -0,0 +1,16 @@
+package udp
+
+import "sync/atomic"
+
+var disableUDPCsum atomic.Bool
+
+// SetDisableUDPCsum controls whether IPv4 UDP sockets opt out of kernel
+// checksum calculation via SO_NO_CHECK. Only applicable on platforms that
+// support the option (Linux). IPv6 always keeps the checksum enabled.
+func SetDisableUDPCsum(disable bool) {
+	disableUDPCsum.Store(disable)
+}
+
+func udpChecksumDisabled() bool {
+	return disableUDPCsum.Load()
+}

udp/conn.go

@@ -11,6 +11,7 @@ const MTU = 9001
 type EncReader func(
 	addr netip.AddrPort,
 	payload []byte,
+	release func(),
 )
 
 type Conn interface {

udp/msghdr_helper_linux_32.go

@@ -0,0 +1,25 @@
+//go:build linux && (386 || amd64p32 || arm || mips || mipsle) && !android && !e2e_testing
+// +build linux
+// +build 386 amd64p32 arm mips mipsle
+// +build !android
+// +build !e2e_testing
+
+package udp
+
+import "golang.org/x/sys/unix"
+
+func controllen(n int) uint32 {
+	return uint32(n)
+}
+
+func setCmsgLen(h *unix.Cmsghdr, n int) {
+	h.Len = uint32(unix.CmsgLen(n))
+}
+
+func setIovecLen(v *unix.Iovec, n int) {
+	v.Len = uint32(n)
+}
+
+func setMsghdrIovlen(m *unix.Msghdr, n int) {
+	m.Iovlen = uint32(n)
+}

udp/msghdr_helper_linux_64.go

@@ -0,0 +1,25 @@
+//go:build linux && (amd64 || arm64 || ppc64 || ppc64le || mips64 || mips64le || s390x || riscv64 || loong64) && !android && !e2e_testing
+// +build linux
+// +build amd64 arm64 ppc64 ppc64le mips64 mips64le s390x riscv64 loong64
+// +build !android
+// +build !e2e_testing
+
+package udp
+
+import "golang.org/x/sys/unix"
+
+func controllen(n int) uint64 {
+	return uint64(n)
+}
+
+func setCmsgLen(h *unix.Cmsghdr, n int) {
+	h.Len = uint64(unix.CmsgLen(n))
+}
+
+func setIovecLen(v *unix.Iovec, n int) {
+	v.Len = uint64(n)
+}
+
+func setMsghdrIovlen(m *unix.Msghdr, n int) {
+	m.Iovlen = uint64(n)
+}

udp/sendmmsg_linux_32.go

@@ -0,0 +1,25 @@
+//go:build linux && (386 || amd64p32 || arm || mips || mipsle) && !android && !e2e_testing
+
+package udp
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+type linuxMmsgHdr struct {
+	Hdr unix.Msghdr
+	Len uint32
+}
+
+func sendmmsg(fd int, hdrs []linuxMmsgHdr, flags int) (int, error) {
+	if len(hdrs) == 0 {
+		return 0, nil
+	}
+	n, _, errno := unix.Syscall6(unix.SYS_SENDMMSG, uintptr(fd), uintptr(unsafe.Pointer(&hdrs[0])), uintptr(len(hdrs)), uintptr(flags), 0, 0)
+	if errno != 0 {
+		return int(n), errno
+	}
+	return int(n), nil
+}

udp/sendmmsg_linux_64.go

@@ -0,0 +1,26 @@
+//go:build linux && (amd64 || arm64 || ppc64 || ppc64le || mips64 || mips64le || s390x || riscv64 || loong64) && !android && !e2e_testing
+
+package udp
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/unix"
+)
+
+type linuxMmsgHdr struct {
+	Hdr unix.Msghdr
+	Len uint32
+	_   uint32
+}
+
+func sendmmsg(fd int, hdrs []linuxMmsgHdr, flags int) (int, error) {
+	if len(hdrs) == 0 {
+		return 0, nil
+	}
+	n, _, errno := unix.Syscall6(unix.SYS_SENDMMSG, uintptr(fd), uintptr(unsafe.Pointer(&hdrs[0])), uintptr(len(hdrs)), uintptr(flags), 0, 0)
+	if errno != 0 {
+		return int(n), errno
+	}
+	return int(n), nil
+}

udp/udp_darwin.go

@@ -180,7 +180,7 @@ func (u *StdConn) ListenOut(r EncReader) {
 			u.l.WithError(err).Error("unexpected udp socket receive error")
 		}
 
-		r(netip.AddrPortFrom(rua.Addr().Unmap(), rua.Port()), buffer[:n])
+		r(netip.AddrPortFrom(rua.Addr().Unmap(), rua.Port()), buffer[:n], nil)
 	}
 }
 

udp/udp_generic.go

@@ -82,6 +82,6 @@ func (u *GenericConn) ListenOut(r EncReader) {
 			return
 		}
 
-		r(netip.AddrPortFrom(rua.Addr().Unmap(), rua.Port()), buffer[:n])
+		r(netip.AddrPortFrom(rua.Addr().Unmap(), rua.Port()), buffer[:n], nil)
 	}
 }

udp/udp_linux_32.go

@@ -7,6 +7,9 @@
 package udp
 
 import (
+	"errors"
+	"fmt"
+
 	"golang.org/x/sys/unix"
 )
 
@@ -30,17 +33,29 @@ type rawMessage struct {
 	Len uint32
 }
 
-func (u *StdConn) PrepareRawMessages(n int) ([]rawMessage, [][]byte, [][]byte) {
+func (u *StdConn) PrepareRawMessages(n int) ([]rawMessage, [][]byte, [][]byte, [][]byte) {
+	controlLen := int(u.controlLen.Load())
+
 	msgs := make([]rawMessage, n)
 	buffers := make([][]byte, n)
 	names := make([][]byte, n)
 
+	var controls [][]byte
+	if controlLen > 0 {
+		controls = make([][]byte, n)
+	}
+
 	for i := range msgs {
-		buffers[i] = make([]byte, MTU)
+		size := int(u.groBufSize.Load())
+		if size < MTU {
+			size = MTU
+		}
+		buf := u.borrowRxBuffer(size)
+		buffers[i] = buf
 		names[i] = make([]byte, unix.SizeofSockaddrInet6)
 
 		vs := []iovec{
-			{Base: &buffers[i][0], Len: uint32(len(buffers[i]))},
+			{Base: &buf[0], Len: uint32(len(buf))},
 		}
 
 		msgs[i].Hdr.Iov = &vs[0]
@@ -48,7 +63,71 @@ func (u *StdConn) PrepareRawMessages(n int) ([]rawMessage, [][]byte, [][]byte) {
 
 		msgs[i].Hdr.Name = &names[i][0]
 		msgs[i].Hdr.Namelen = uint32(len(names[i]))
+
+		if controlLen > 0 {
+			controls[i] = make([]byte, controlLen)
+			msgs[i].Hdr.Control = &controls[i][0]
+			msgs[i].Hdr.Controllen = controllen(len(controls[i]))
+		} else {
+			msgs[i].Hdr.Control = nil
+			msgs[i].Hdr.Controllen = controllen(0)
+		}
 	}
 
-	return msgs, buffers, names
+	return msgs, buffers, names, controls
+}
+
+func setIovecBase(msg *rawMessage, buf []byte) {
+	iov := (*iovec)(msg.Hdr.Iov)
+	iov.Base = &buf[0]
+	iov.Len = uint32(len(buf))
+}
+
+func rawMessageToUnixMsghdr(msg *rawMessage) (unix.Msghdr, unix.Iovec, error) {
+	var hdr unix.Msghdr
+	var iov unix.Iovec
+	if msg == nil {
+		return hdr, iov, errors.New("nil rawMessage")
+	}
+	if msg.Hdr.Iov == nil || msg.Hdr.Iov.Base == nil {
+		return hdr, iov, errors.New("rawMessage missing payload buffer")
+	}
+	payloadLen := int(msg.Hdr.Iov.Len)
+	if payloadLen < 0 {
+		return hdr, iov, fmt.Errorf("invalid payload length: %d", payloadLen)
+	}
+	iov.Base = msg.Hdr.Iov.Base
+	iov.Len = uint32(payloadLen)
+	hdr.Iov = &iov
+	hdr.Iovlen = 1
+	hdr.Name = msg.Hdr.Name
+	// CRITICAL: Always set to full buffer size for receive, not what kernel wrote last time
+	if hdr.Name != nil {
+		hdr.Namelen = uint32(unix.SizeofSockaddrInet6)
+	} else {
+		hdr.Namelen = 0
+	}
+	hdr.Control = msg.Hdr.Control
+	// CRITICAL: Use the allocated size, not what was previously returned
+	if hdr.Control != nil {
+		// Control buffer size is stored in Controllen from PrepareRawMessages
+		hdr.Controllen = msg.Hdr.Controllen
+	} else {
+		hdr.Controllen = 0
+	}
+	hdr.Flags = 0 // Reset flags for new receive
+	return hdr, iov, nil
+}
+
+func updateRawMessageFromUnixMsghdr(msg *rawMessage, hdr *unix.Msghdr, n int) {
+	if msg == nil || hdr == nil {
+		return
+	}
+	msg.Hdr.Namelen = hdr.Namelen
+	msg.Hdr.Controllen = hdr.Controllen
+	msg.Hdr.Flags = hdr.Flags
+	if n < 0 {
+		n = 0
+	}
+	msg.Len = uint32(n)
 }

udp/udp_linux_64.go

@@ -7,6 +7,9 @@
 package udp
 
 import (
+	"errors"
+	"fmt"
+
 	"golang.org/x/sys/unix"
 )
 
@@ -33,25 +36,99 @@ type rawMessage struct {
 	Pad0 [4]byte
 }
 
-func (u *StdConn) PrepareRawMessages(n int) ([]rawMessage, [][]byte, [][]byte) {
+func (u *StdConn) PrepareRawMessages(n int) ([]rawMessage, [][]byte, [][]byte, [][]byte) {
+	controlLen := int(u.controlLen.Load())
+
 	msgs := make([]rawMessage, n)
 	buffers := make([][]byte, n)
 	names := make([][]byte, n)
 
+	var controls [][]byte
+	if controlLen > 0 {
+		controls = make([][]byte, n)
+	}
+
 	for i := range msgs {
-		buffers[i] = make([]byte, MTU)
+		size := int(u.groBufSize.Load())
+		if size < MTU {
+			size = MTU
+		}
+		buf := u.borrowRxBuffer(size)
+		buffers[i] = buf
 		names[i] = make([]byte, unix.SizeofSockaddrInet6)
 
-		vs := []iovec{
-			{Base: &buffers[i][0], Len: uint64(len(buffers[i]))},
-		}
+		vs := []iovec{{Base: &buf[0], Len: uint64(len(buf))}}
 
 		msgs[i].Hdr.Iov = &vs[0]
 		msgs[i].Hdr.Iovlen = uint64(len(vs))
 
 		msgs[i].Hdr.Name = &names[i][0]
 		msgs[i].Hdr.Namelen = uint32(len(names[i]))
+
+		if controlLen > 0 {
+			controls[i] = make([]byte, controlLen)
+			msgs[i].Hdr.Control = &controls[i][0]
+			msgs[i].Hdr.Controllen = controllen(len(controls[i]))
+		} else {
+			msgs[i].Hdr.Control = nil
+			msgs[i].Hdr.Controllen = controllen(0)
+		}
 	}
 
-	return msgs, buffers, names
+	return msgs, buffers, names, controls
+}
+
+func setIovecBase(msg *rawMessage, buf []byte) {
+	iov := (*iovec)(msg.Hdr.Iov)
+	iov.Base = &buf[0]
+	iov.Len = uint64(len(buf))
+}
+
+func rawMessageToUnixMsghdr(msg *rawMessage) (unix.Msghdr, unix.Iovec, error) {
+	var hdr unix.Msghdr
+	var iov unix.Iovec
+	if msg == nil {
+		return hdr, iov, errors.New("nil rawMessage")
+	}
+	if msg.Hdr.Iov == nil || msg.Hdr.Iov.Base == nil {
+		return hdr, iov, errors.New("rawMessage missing payload buffer")
+	}
+	payloadLen := int(msg.Hdr.Iov.Len)
+	if payloadLen < 0 {
+		return hdr, iov, fmt.Errorf("invalid payload length: %d", payloadLen)
+	}
+	iov.Base = msg.Hdr.Iov.Base
+	iov.Len = uint64(payloadLen)
+	hdr.Iov = &iov
+	hdr.Iovlen = 1
+	hdr.Name = msg.Hdr.Name
+	// CRITICAL: Always set to full buffer size for receive, not what kernel wrote last time
+	if hdr.Name != nil {
+		hdr.Namelen = uint32(unix.SizeofSockaddrInet6)
+	} else {
+		hdr.Namelen = 0
+	}
+	hdr.Control = msg.Hdr.Control
+	// CRITICAL: Use the allocated size, not what was previously returned
+	if hdr.Control != nil {
+		// Control buffer size is stored in Controllen from PrepareRawMessages
+		hdr.Controllen = msg.Hdr.Controllen
+	} else {
+		hdr.Controllen = 0
+	}
+	hdr.Flags = 0 // Reset flags for new receive
+	return hdr, iov, nil
+}
+
+func updateRawMessageFromUnixMsghdr(msg *rawMessage, hdr *unix.Msghdr, n int) {
+	if msg == nil || hdr == nil {
+		return
+	}
+	msg.Hdr.Namelen = hdr.Namelen
+	msg.Hdr.Controllen = hdr.Controllen
+	msg.Hdr.Flags = hdr.Flags
+	if n < 0 {
+		n = 0
+	}
+	msg.Len = uint32(n)
 }

udp/udp_rio_windows.go

@@ -149,7 +149,7 @@ func (u *RIOConn) ListenOut(r EncReader) {
 			continue
 		}
 
-		r(netip.AddrPortFrom(netip.AddrFrom16(rua.Addr).Unmap(), (rua.Port>>8)|((rua.Port&0xff)<<8)), buffer[:n])
+		r(netip.AddrPortFrom(netip.AddrFrom16(rua.Addr).Unmap(), (rua.Port>>8)|((rua.Port&0xff)<<8)), buffer[:n], nil)
 	}
 }
 

udp/udp_tester.go

@@ -112,7 +112,7 @@ func (u *TesterConn) ListenOut(r EncReader) {
 		if !ok {
 			return
 		}
-		r(p.From, p.Data)
+		r(p.From, p.Data, func() {})
 	}
 }