From 38a5ef9e5979687c9b3e79288471f4a99f85c270 Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Thu, 28 Nov 2024 11:54:19 +0100
Subject: [PATCH] fix(udp): do not enable URO on Windows on ARM

On Windows on ARM with Windows Subsystem for Linux (WSL) `WSA_RECVMSG` does not
return the segment size of coalesced UDP datagrams. See
https://github.com/quinn-rs/quinn/issues/2041 for details.

While lacking a fix for the root cause, don't enable URO, i.e. don't coalesce
UDP datagrams on Windows on ARM.
---
 Cargo.toml               |  2 +-
 quinn-udp/src/windows.rs | 71 +++++++++++++++++++++++++++++++++-------
 2 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index b78d81eaa..71f5668c9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -53,7 +53,7 @@ tracing-subscriber = { version = "0.3.0", default-features = false, features = [
 url = "2"
 wasm-bindgen-test = { version = "0.3.45" }
 web-time = "1"
-windows-sys = { version = ">=0.52, <=0.59", features = ["Win32_Foundation", "Win32_System_IO", "Win32_Networking_WinSock"] }
+windows-sys = { version = ">=0.52, <=0.59", features = ["Win32_Foundation", "Win32_System_IO", "Win32_Networking_WinSock", "Win32_System_SystemInformation", "Win32_System_Threading"] }
 
 [profile.bench]
 debug = true
diff --git a/quinn-udp/src/windows.rs b/quinn-udp/src/windows.rs
index 223ab0fdf..28fe2b51c 100644
--- a/quinn-udp/src/windows.rs
+++ b/quinn-udp/src/windows.rs
@@ -10,7 +10,13 @@ use std::{
 
 use libc::{c_int, c_uint};
 use once_cell::sync::Lazy;
-use windows_sys::Win32::Networking::WinSock;
+use windows_sys::Win32::{
+    Networking::WinSock,
+    System::{
+        SystemInformation::IMAGE_FILE_MACHINE_ARM64,
+        Threading::{GetCurrentProcess, IsWow64Process2},
+    },
+};
 
 use crate::{
     cmsg::{self, CMsgHdr},
@@ -107,16 +113,33 @@ impl UdpSocketState {
             )?;
         }
 
-        // Opportunistically try to enable GRO
-        _ = set_socket_option(
-            &*socket.0,
-            WinSock::IPPROTO_UDP,
-            WinSock::UDP_RECV_MAX_COALESCED_SIZE,
-            // u32 per
-            // https://learn.microsoft.com/en-us/windows/win32/winsock/ipproto-udp-socket-options.
-            // Choice of 2^16 - 1 inspired by msquic.
-            u16::MAX as u32,
-        );
+        match &*IS_WINDOWS_ON_ARM {
+            Ok(true) => {
+                // Bug on Windows on ARM, not receiving `UDP_COALESCED_INFO` `CMSG`
+                // when _Virtual Machine Platform_ feature enabled. See
+                // <https://github.com/quinn-rs/quinn/issues/2041> for details.
+                debug!("detected Windows on ARM host thus not enabling URO")
+            }
+            Ok(false) => {
+                // Opportunistically try to enable URO
+                let result = set_socket_option(
+                    &*socket.0,
+                    WinSock::IPPROTO_UDP,
+                    WinSock::UDP_RECV_MAX_COALESCED_SIZE,
+                    // u32 per
+                    // https://learn.microsoft.com/en-us/windows/win32/winsock/ipproto-udp-socket-options.
+                    // Choice of 2^16 - 1 inspired by msquic.
+                    u16::MAX as u32,
+                );
+
+                if let Err(e) = result {
+                    debug!("failed to enable URO: {}", e);
+                }
+            }
+            Err(e) => {
+                debug!("failed to detect host system thus not enabling URO: {}", e);
+            }
+        }
 
         let now = Instant::now();
         Ok(Self {
@@ -470,3 +493,29 @@ static MAX_GSO_SEGMENTS: Lazy<usize> = Lazy::new(|| {
         Err(_) => 1,
     }
 });
+
+/// Evaluates to [`Ok(true)`] if executed either directly on Windows on ARM, or
+/// on an emulator which itself executes on Windows on ARM.
+///
+/// See
+/// <https://learn.microsoft.com/en-us/windows/arm/apps-on-arm-x86-emulation#detecting-emulation>
+/// for details.
+static IS_WINDOWS_ON_ARM: Lazy<io::Result<bool>> = Lazy::new(|| {
+    let mut process_machine: u16 = 0;
+    let mut native_machine: u16 = 0;
+
+    let result = unsafe {
+        IsWow64Process2(
+            GetCurrentProcess(),
+            &mut process_machine as *mut u16,
+            &mut native_machine as *mut u16,
+        )
+    };
+
+    match result {
+        // See
+        // <https://learn.microsoft.com/en-us/windows/win32/api/wow64apiset/nf-wow64apiset-iswow64process2#return-value>.
+        0 => Err(io::Error::last_os_error()),
+        _ => Ok(native_machine == IMAGE_FILE_MACHINE_ARM64),
+    }
+});