diff --git a/vm/src/stdlib/select.rs b/vm/src/stdlib/select.rs index 9e175870f..5a9718e0f 100644 --- a/vm/src/stdlib/select.rs +++ b/vm/src/stdlib/select.rs @@ -1,6 +1,6 @@ use crate::pyobject::{PyObjectRef, PyResult, TryFromObject}; use crate::vm::VirtualMachine; -use std::io; +use std::{io, mem}; pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef { super::socket::init_winsock(); @@ -17,28 +17,34 @@ pub(crate) fn make_module(vm: &VirtualMachine) -> PyObjectRef { mod platform { pub use libc::{fd_set, select, timeval, FD_ISSET, FD_SET, FD_SETSIZE, FD_ZERO}; pub use std::os::unix::io::RawFd; + + pub fn check_err(x: i32) -> bool { + x < 0 + } } #[allow(non_snake_case)] #[cfg(windows)] mod platform { - pub use winapi::um::winsock2::{fd_set, select, timeval, FD_SETSIZE, SOCKET as RawFd}; - // from winsock2.h: https://gist.github.com/piscisaureus/906386#file-winsock2-h-L128-L141 + use winapi::um::winsock2; + pub use winsock2::{fd_set, select, timeval, FD_SETSIZE, SOCKET as RawFd}; + + // based off winsock2.h: https://gist.github.com/piscisaureus/906386#file-winsock2-h-L128-L141 pub unsafe fn FD_SET(fd: RawFd, set: *mut fd_set) { - let mut i = 0; - for idx in 0..(*set).fd_count as usize { - i = idx; - if (*set).fd_array[i] == fd { - break; + let mut slot = std::ptr::addr_of_mut!((*set).fd_array).cast::(); + let fd_count = (*set).fd_count; + for _ in 0..fd_count { + if *slot == fd { + return; } + slot = slot.add(1); } - if i == (*set).fd_count as usize { - if (*set).fd_count < FD_SETSIZE as u32 { - (*set).fd_array[i] = fd as _; - (*set).fd_count += 1; - } + // slot == &fd_array[fd_count] at this point + if fd_count < FD_SETSIZE as u32 { + *slot = fd as RawFd; + (*set).fd_count += 1; } } @@ -50,6 +56,10 @@ mod platform { use winapi::um::winsock2::__WSAFDIsSet; __WSAFDIsSet(fd as _, set) != 0 } + + pub fn check_err(x: i32) -> bool { + x == winsock2::SOCKET_ERROR + } } pub use platform::timeval; @@ -72,8 +82,9 @@ impl TryFromObject for Selectable { } } +// Keep it in a MaybeUninit, since on windows FD_ZERO doesn't actually zero the whole thing #[repr(transparent)] -pub struct FdSet(platform::fd_set); +pub struct FdSet(mem::MaybeUninit); impl FdSet { pub fn new() -> FdSet { @@ -81,29 +92,25 @@ impl FdSet { // interacting with it is in C, so it's safe to zero let mut fdset = std::mem::MaybeUninit::zeroed(); unsafe { platform::FD_ZERO(fdset.as_mut_ptr()) }; - FdSet(unsafe { fdset.assume_init() }) + FdSet(fdset) } pub fn insert(&mut self, fd: RawFd) { - unsafe { platform::FD_SET(fd, &mut self.0) }; + unsafe { platform::FD_SET(fd, self.0.as_mut_ptr()) }; } pub fn contains(&mut self, fd: RawFd) -> bool { - unsafe { platform::FD_ISSET(fd, &mut self.0) } + unsafe { platform::FD_ISSET(fd, self.0.as_mut_ptr()) } } pub fn clear(&mut self) { - unsafe { platform::FD_ZERO(&mut self.0) }; + unsafe { platform::FD_ZERO(self.0.as_mut_ptr()) }; } pub fn highest(&mut self) -> Option { - for i in (0..platform::FD_SETSIZE as RawFd).rev() { - if self.contains(i) { - return Some(i); - } - } - - None + (0..platform::FD_SETSIZE as RawFd) + .rev() + .find(|&i| self.contains(i)) } } @@ -121,13 +128,13 @@ pub fn select( let ret = unsafe { platform::select( nfds, - &mut readfds.0, - &mut writefds.0, - &mut errfds.0, + readfds.0.as_mut_ptr(), + writefds.0.as_mut_ptr(), + errfds.0.as_mut_ptr(), timeout, ) }; - if ret < 0 { + if platform::check_err(ret) { Err(io::Error::last_os_error()) } else { Ok(ret)