use std::{
ffi::c_int,
io, mem,
net::{IpAddr, Ipv4Addr, SocketAddr},
os::windows::io::AsRawSocket,
ptr,
};
use libc::c_uchar;
use socket2::Socket;
use windows_sys::Win32::Networking::WinSock::{self, SOCKET};
use crate::{DEFAULT_TTL, Io, UdpSocketController};
const CMSG_LEN: usize = 128;
#[derive(Copy, Clone)]
#[repr(align(8))] pub(crate) struct Aligned<T>(pub(crate) T);
impl Io for UdpSocketController {
fn config(socket: &Socket, addr: SocketAddr) -> std::io::Result<()> {
const OPTION_ON: c_int = 1;
const OPTION_OFF: c_int = 0;
let io = socket.as_raw_socket().try_into().unwrap();
setsockopt(io, WinSock::SOL_SOCKET, WinSock::SO_RCVBUF, 2 * 1024 * 1024);
match addr {
SocketAddr::V4(_) => {
setsockopt(io, WinSock::IPPROTO_IP, WinSock::IP_RECVTOS, OPTION_ON);
setsockopt(io, WinSock::IPPROTO_IP, WinSock::IP_PKTINFO, OPTION_ON);
setsockopt(io, WinSock::IPPROTO_IP, WinSock::IP_RECVTTL, OPTION_ON);
setsockopt(io, WinSock::IPPROTO_IP, WinSock::IP_RECVDSTADDR, OPTION_ON);
setsockopt(io, WinSock::IPPROTO_IP, WinSock::IP_TTL, DEFAULT_TTL);
}
SocketAddr::V6(_) => {
setsockopt(io, WinSock::IPPROTO_IPV6, WinSock::IPV6_V6ONLY, OPTION_OFF);
setsockopt(io, WinSock::IPPROTO_IPV6, WinSock::IPV6_HOPLIMIT, OPTION_ON);
setsockopt(
io,
WinSock::IPPROTO_IPV6,
WinSock::IPV6_RECVTCLASS,
OPTION_ON,
);
setsockopt(io, WinSock::IPPROTO_IPV6, WinSock::IPV6_PKTINFO, OPTION_ON);
}
}
if let Err(e) = socket.bind(&addr.into()) {
tracing::error!("Failed to bind socket: {}", e);
return Err(io::Error::new(io::ErrorKind::AddrInUse, e));
}
Ok(())
}
fn sendmsg(
&self,
bufs: &[std::io::IoSlice<'_>],
hdr: &crate::DatagramHeader,
) -> std::io::Result<usize> {
let dst = socket2::SockAddr::from(hdr.dst);
let mut count = 0;
for buf in bufs {
let mut ctrl_buf = Aligned([0; CMSG_LEN]);
let mut data = WinSock::WSABUF {
buf: buf.as_ptr() as *mut _,
len: buf.len() as _,
};
let ctrl = WinSock::WSABUF {
buf: ctrl_buf.0.as_mut_ptr(),
len: ctrl_buf.0.len() as _,
};
let mut wsa_msg = WinSock::WSAMSG {
name: dst.as_ptr() as *mut _,
namelen: dst.len(),
lpBuffers: &mut data,
Control: ctrl,
dwBufferCount: 1,
dwFlags: 0,
};
let mut cmsg = unsafe { first_cmsg(&mut wsa_msg).as_mut() };
let mut cmsg_len = 0;
if !hdr.src.ip().is_unspecified() {
let src = socket2::SockAddr::from(hdr.src);
match src.family() {
WinSock::AF_INET => {
let src_ip =
unsafe { ptr::read(src.as_ptr() as *const WinSock::SOCKADDR_IN) };
let pktinfo = WinSock::IN_PKTINFO {
ipi_addr: src_ip.sin_addr,
ipi_ifindex: 0,
};
cmsg = append_cmsg(
&wsa_msg,
cmsg,
WinSock::IPPROTO_IP,
WinSock::IP_PKTINFO,
pktinfo,
&mut cmsg_len,
);
}
WinSock::AF_INET6 => {
let src_ip =
unsafe { ptr::read(src.as_ptr() as *const WinSock::SOCKADDR_IN6) };
let pktinfo = WinSock::IN6_PKTINFO {
ipi6_addr: src_ip.sin6_addr,
ipi6_ifindex: unsafe { src_ip.Anonymous.sin6_scope_id },
};
cmsg = append_cmsg(
&wsa_msg,
cmsg,
WinSock::IPPROTO_IPV6,
WinSock::IPV6_PKTINFO,
pktinfo,
&mut cmsg_len,
);
}
_ => {
return Err(io::Error::from(io::ErrorKind::InvalidInput));
}
}
}
if let Some(ecn) = hdr.ecn {
let is_ipv4 = hdr.dst.is_ipv4()
|| matches!(hdr.dst.ip(), IpAddr::V6(addr) if addr.to_ipv4_mapped().is_some());
if is_ipv4 {
_ = append_cmsg(
&wsa_msg,
cmsg,
WinSock::IPPROTO_IP,
WinSock::IP_ECN,
ecn,
&mut cmsg_len,
);
} else {
_ = append_cmsg(
&wsa_msg,
cmsg,
WinSock::IPPROTO_IPV6,
WinSock::IPV6_TCLASS,
ecn,
&mut cmsg_len,
);
}
}
wsa_msg.Control.len = cmsg_len as _;
if cmsg_len == 0 {
wsa_msg.Control = WinSock::WSABUF {
buf: ptr::null_mut(),
len: 0,
};
}
let mut len = 0;
let ret = unsafe {
WinSock::WSASendMsg(
self.io.as_raw_socket() as usize,
&wsa_msg,
0,
&mut len,
ptr::null_mut(),
None,
)
};
if ret != 0 {
let e = io::Error::last_os_error();
if e.kind() != io::ErrorKind::WouldBlock {
return Err(e);
}
}
count += 1;
}
Ok(count as usize)
}
fn recvmsg(
&self,
bufs: &mut [std::io::IoSliceMut<'_>],
hdr: &mut [crate::DatagramHeader],
) -> std::io::Result<usize> {
let wsa_recvmsg_ptr = WSARECVMSG_PTR.expect("valid function pointer for WSARecvMsg");
let mut ctrl_buf = Aligned([0; CMSG_LEN]);
let mut source: WinSock::SOCKADDR_INET = unsafe { mem::zeroed() };
let ctrl = WinSock::WSABUF {
buf: ctrl_buf.0.as_mut_ptr(),
len: ctrl_buf.0.len() as _,
};
let mut wsa_msg = WinSock::WSAMSG {
name: &mut source as *mut _ as *mut _,
namelen: mem::size_of_val(&source) as _,
lpBuffers: &mut WinSock::WSABUF {
buf: bufs[0].as_mut_ptr(),
len: bufs[0].len() as _,
},
Control: ctrl,
dwBufferCount: 1,
dwFlags: 0,
};
let mut len = 0;
unsafe {
let rc = (wsa_recvmsg_ptr)(
self.io.as_raw_socket() as usize,
&mut wsa_msg,
&mut len,
ptr::null_mut(),
None,
);
if rc == -1 {
return Err(io::Error::last_os_error());
}
}
let addr = unsafe {
let (_, addr) = socket2::SockAddr::try_init(|addr_storage, len| {
*len = mem::size_of_val(&source) as _;
ptr::copy_nonoverlapping(&source, addr_storage as _, 1);
Ok(())
})?;
addr.as_socket()
};
let mut ecn_bits = 0;
let mut dst_ip = None;
let mut cmsg: Option<&mut WinSock::CMSGHDR> = unsafe { first_cmsg(&mut wsa_msg).as_mut() };
while let Some(cur_cmsg) = cmsg {
match (cur_cmsg.cmsg_level, cur_cmsg.cmsg_type) {
(WinSock::IPPROTO_IP, WinSock::IP_PKTINFO) => {
let pktinfo = cmsg_decode::<WinSock::IN_PKTINFO>(cur_cmsg);
let ip4 = Ipv4Addr::from(u32::from_be(unsafe { pktinfo.ipi_addr.S_un.S_addr }));
dst_ip = Some(ip4.into());
}
(WinSock::IPPROTO_IPV6, WinSock::IPV6_PKTINFO) => {
let pktinfo = cmsg_decode::<WinSock::IN6_PKTINFO>(cur_cmsg);
dst_ip = Some(IpAddr::from(unsafe { pktinfo.ipi6_addr.u.Byte }));
}
(WinSock::IPPROTO_IP, WinSock::IP_ECN) => {
ecn_bits = cmsg_decode::<c_int>(cur_cmsg);
}
(WinSock::IPPROTO_IPV6, WinSock::IPV6_ECN) => {
ecn_bits = cmsg_decode::<c_int>(cur_cmsg);
}
_ => {}
}
cmsg = unsafe { next_cmsg(&wsa_msg, cur_cmsg).as_mut() };
}
let dst = if let Some(ip) = dst_ip {
crate::SocketAddr::new(ip, self.local_addr()?.port())
} else {
self.local_addr()?
};
hdr[0] = crate::DatagramHeader {
src: addr.unwrap(),
dst,
ttl: DEFAULT_TTL as u8,
ecn: Some(ecn_bits as u8),
seg_size: len as u16,
};
Ok(1)
}
}
fn append_cmsg<'a, V: Copy>(
msg: &'a WinSock::WSAMSG,
mut cmsg: Option<&'a mut WinSock::CMSGHDR>,
level: libc::c_int,
ty: libc::c_int,
data: V,
cmsg_len: &mut usize,
) -> Option<&'a mut WinSock::CMSGHDR> {
let space = cmsg_space(mem::size_of_val(&data));
let next = cmsg.take().expect("no available cmsghdr");
next.cmsg_level = level as _;
next.cmsg_type = ty as _;
next.cmsg_len = cmsg_data_len(mem::size_of_val(&data)) as _;
unsafe {
ptr::write(cmsg_data(next) as *const V as *mut V, data);
}
*cmsg_len += space;
unsafe { next_cmsg(msg, next).as_mut() }
}
fn cmsg_decode<T: Copy>(cmsg: &mut WinSock::CMSGHDR) -> T {
unsafe { ptr::read(cmsg_data(cmsg) as *const T) }
}
const fn cmsghdr_align(length: usize) -> usize {
(length + mem::align_of::<WinSock::CMSGHDR>() - 1) & !(mem::align_of::<WinSock::CMSGHDR>() - 1)
}
fn cmsgdata_align(length: usize) -> usize {
(length + mem::align_of::<usize>() - 1) & !(mem::align_of::<usize>() - 1)
}
fn cmsg_data_len(len: usize) -> usize {
mem::size_of::<WinSock::CMSGHDR>() + len
}
fn cmsg_space(len: usize) -> usize {
let total = mem::size_of::<WinSock::CMSGHDR>() + len;
let align = mem::align_of::<WinSock::CMSGHDR>();
(total + align - 1) & !(align - 1)
}
unsafe fn first_cmsg(msg: &mut WinSock::WSAMSG) -> *mut WinSock::CMSGHDR {
if msg.Control.len as usize >= mem::size_of::<WinSock::CMSGHDR>() {
msg.Control.buf as *mut WinSock::CMSGHDR
} else {
ptr::null_mut::<WinSock::CMSGHDR>()
}
}
fn next_cmsg(msg: &WinSock::WSAMSG, cmsg: &WinSock::CMSGHDR) -> *mut WinSock::CMSGHDR {
let next = (cmsg as *const _ as usize + cmsghdr_align(cmsg.cmsg_len)) as *mut WinSock::CMSGHDR;
let max = msg.Control.buf as usize + msg.Control.len as usize;
if unsafe { next.offset(1) } as usize > max {
ptr::null_mut()
} else {
next
}
}
fn cmsg_data(cmsg: &mut WinSock::CMSGHDR) -> *mut libc::c_uchar {
(cmsg as *const _ as usize + cmsgdata_align(mem::size_of::<WinSock::CMSGHDR>())) as *mut c_uchar
}
fn setsockopt(io: SOCKET, level: libc::c_int, name: libc::c_int, value: libc::c_int) {
unsafe {
WinSock::setsockopt(
io,
level,
name,
&value as *const _ as _,
mem::size_of_val(&value) as _,
)
};
}
static WSARECVMSG_PTR: std::sync::LazyLock<WinSock::LPFN_WSARECVMSG> =
std::sync::LazyLock::new(|| {
let s = unsafe { WinSock::socket(WinSock::AF_INET as _, WinSock::SOCK_DGRAM as _, 0) };
if s == WinSock::INVALID_SOCKET {
tracing::warn!(
"failed to create socket for WSARecvMsg function pointer: {}",
io::Error::last_os_error()
);
return None;
}
let guid = WinSock::WSAID_WSARECVMSG;
let mut wsa_recvmsg_ptr = None;
let mut len = 0;
let ret = unsafe {
WinSock::WSAIoctl(
s as _,
WinSock::SIO_GET_EXTENSION_FUNCTION_POINTER,
&guid as *const _ as *const _,
mem::size_of_val(&guid) as u32,
&mut wsa_recvmsg_ptr as *mut _ as *mut _,
mem::size_of_val(&wsa_recvmsg_ptr) as u32,
&mut len,
ptr::null_mut(),
None,
)
};
if ret == -1 {
tracing::warn!(
"failed to get WSARecvMsg function pointer: {}",
io::Error::last_os_error()
);
} else if len as usize != mem::size_of::<WinSock::LPFN_WSARECVMSG>() {
tracing::warn!(
"WSARecvMsg function pointer size mismatch: expected {}, got {}",
mem::size_of::<WinSock::LPFN_WSARECVMSG>(),
len
);
wsa_recvmsg_ptr = None;
}
unsafe {
WinSock::closesocket(s);
}
wsa_recvmsg_ptr
});