RustPython/stdlib/src/socket.rs at main · mode9/RustPython

1935 lines (1794 loc) · 66.4 KB
use crate::vm::{PyObjectRef, VirtualMachine};
#[cfg(feature = "ssl")]
pub(super) use _socket::{sock_select, timeout_error_msg, PySocket, SelectKind};
pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
    #[cfg(windows)]
    crate::vm::stdlib::nt::init_winsock();
    _socket::make_module(vm)
#[pymodule]
mod _socket {
    use crate::common::lock::{PyMappedRwLockReadGuard, PyRwLock, PyRwLockReadGuard};
    use crate::vm::{
        builtins::{PyBaseExceptionRef, PyListRef, PyStrRef, PyTupleRef, PyTypeRef},
        convert::{IntoPyException, ToPyObject, TryFromBorrowedObject, TryFromObject},
        function::{ArgBytesLike, ArgMemoryBuffer, Either, OptionalArg, OptionalOption},
        types::{DefaultConstructor, Initializer},
        utils::ToCString,
        AsObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
    use crossbeam_utils::atomic::AtomicCell;
    use num_traits::ToPrimitive;
    use socket2::{Domain, Protocol, Socket, Type as SocketType};
    use std::{
        ffi,
        io::{self, Read, Write},
        mem::MaybeUninit,
        net::{self, Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr, ToSocketAddrs},
        time::{Duration, Instant},
    #[cfg(unix)]
    use libc as c;
    #[cfg(windows)]
    mod c {
        pub use winapi::shared::ifdef::IF_MAX_STRING_SIZE as IF_NAMESIZE;
        pub use winapi::shared::netioapi::{if_indextoname, if_nametoindex};
        pub use winapi::shared::ws2def::*;
        pub use winapi::um::winsock2::{
            SD_BOTH as SHUT_RDWR, SD_RECEIVE as SHUT_RD, SD_SEND as SHUT_WR, SOCK_DGRAM, SOCK_RAW,
            SOCK_RDM, SOCK_SEQPACKET, SOCK_STREAM, SOL_SOCKET, SO_BROADCAST, SO_ERROR, SO_LINGER,
            SO_OOBINLINE, SO_REUSEADDR, SO_TYPE, *,
    // constants
    #[pyattr(name = "has_ipv6")]
    const HAS_IPV6: bool = true;
    #[pyattr]
    // put IPPROTO_MAX later
    use c::{
        AF_DECnet, AF_APPLETALK, AF_INET, AF_INET6, AF_IPX, AF_UNSPEC, INADDR_ANY, INADDR_LOOPBACK,
        INADDR_NONE, IPPROTO_AH, IPPROTO_DSTOPTS, IPPROTO_EGP, IPPROTO_ESP, IPPROTO_FRAGMENT,
        IPPROTO_HOPOPTS, IPPROTO_ICMP, IPPROTO_ICMPV6, IPPROTO_IDP, IPPROTO_IGMP, IPPROTO_IP,
        IPPROTO_IP as IPPROTO_IPIP, IPPROTO_IPV6, IPPROTO_NONE, IPPROTO_PIM, IPPROTO_PUP,
        IPPROTO_RAW, IPPROTO_ROUTING, IPPROTO_TCP, IPPROTO_TCP as SOL_TCP, IPPROTO_UDP, MSG_CTRUNC,
        MSG_DONTROUTE, MSG_OOB, MSG_PEEK, MSG_TRUNC, MSG_WAITALL, NI_DGRAM, NI_MAXHOST,
        NI_NAMEREQD, NI_NOFQDN, NI_NUMERICHOST, NI_NUMERICSERV, SHUT_RD, SHUT_RDWR, SHUT_WR,
        SOCK_DGRAM, SOCK_STREAM, SOL_SOCKET, SO_BROADCAST, SO_ERROR, SO_LINGER, SO_OOBINLINE,
        SO_REUSEADDR, SO_TYPE, TCP_NODELAY,
    #[cfg(unix)]
    #[pyattr]
    use c::{AF_UNIX, SO_REUSEPORT};
    #[pyattr]
    use c::{AI_ADDRCONFIG, AI_NUMERICHOST, AI_NUMERICSERV, AI_PASSIVE};
    #[cfg(not(target_os = "redox"))]
    #[pyattr]
    use c::{SOCK_RAW, SOCK_RDM, SOCK_SEQPACKET};
    #[cfg(windows)]
    #[pyattr]
    use winapi::shared::ws2def::{
        IPPROTO_CBT, IPPROTO_ICLFXBM, IPPROTO_IGP, IPPROTO_L2TP, IPPROTO_PGM, IPPROTO_RDP,
        IPPROTO_SCTP, IPPROTO_ST,
    #[pyattr]
    fn error(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.exceptions.os_error.to_owned()
    #[pyattr(once)]
    fn timeout(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.new_exception_type(
            "socket",
            "timeout",
            Some(vec![vm.ctx.exceptions.os_error.to_owned()]),
    #[pyattr(once)]
    fn herror(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.new_exception_type(
            "socket",
            "herror",
            Some(vec![vm.ctx.exceptions.os_error.to_owned()]),
    #[pyattr(once)]
    fn gaierror(vm: &VirtualMachine) -> PyTypeRef {
        vm.ctx.new_exception_type(
            "socket",
            "gaierror",
            Some(vec![vm.ctx.exceptions.os_error.to_owned()]),
    #[pyfunction]
    fn htonl(x: u32) -> u32 {
        u32::to_be(x)
    #[pyfunction]
    fn htons(x: u16) -> u16 {
        u16::to_be(x)
    #[pyfunction]
    fn ntohl(x: u32) -> u32 {
        u32::from_be(x)
    #[pyfunction]
    fn ntohs(x: u16) -> u16 {
        u16::from_be(x)
    #[cfg(unix)]
    type RawSocket = std::os::unix::io::RawFd;
    #[cfg(windows)]
    type RawSocket = std::os::windows::raw::SOCKET;
    #[cfg(unix)]
    macro_rules! errcode {
        ($e:ident) => {
            c::$e
    #[cfg(windows)]
    macro_rules! errcode {
    ($e:ident) => {
        paste::paste!(c::[<WSA $e>])
    #[cfg(windows)]
    use winapi::shared::netioapi;
    fn get_raw_sock(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<RawSocket> {
        #[cfg(unix)]
        type CastFrom = libc::c_long;
        #[cfg(windows)]
        type CastFrom = libc::c_longlong;
        // should really just be to_index() but test_socket tests the error messages explicitly
        if obj.fast_isinstance(vm.ctx.types.float_type) {
            return Err(vm.new_type_error("integer argument expected, got float".to_owned()));
        let int = vm
            .to_index_opt(obj)
            .unwrap_or_else(|| Err(vm.new_type_error("an integer is required".to_owned())))?;
        int.try_to_primitive::<CastFrom>(vm)
            .map(|sock| sock as RawSocket)
    #[cfg(unix)]
    mod nullable_socket {
        use super::*;
        use std::os::unix::io::AsRawFd;
        #[derive(Debug)]
        #[repr(transparent)]
        pub struct NullableSocket(Option<socket2::Socket>);
        impl From<socket2::Socket> for NullableSocket {
            fn from(sock: socket2::Socket) -> Self {
                NullableSocket(Some(sock))
        impl NullableSocket {
            pub fn invalid() -> Self {
                Self(None)
            pub fn get(&self) -> Option<&socket2::Socket> {
                self.0.as_ref()
            pub fn fd(&self) -> RawSocket {
                self.get().map_or(INVALID_SOCKET, |sock| sock.as_raw_fd())
            pub fn insert(&mut self, sock: socket2::Socket) -> &mut socket2::Socket {
                self.0.insert(sock)
    #[cfg(windows)]
    mod nullable_socket {
        use super::*;
        use std::os::windows::io::{AsRawSocket, FromRawSocket};
        // TODO: may change if windows changes its TcpStream repr
        #[derive(Debug)]
        #[repr(transparent)]
        pub struct NullableSocket(socket2::Socket);
        impl From<socket2::Socket> for NullableSocket {
            fn from(sock: socket2::Socket) -> Self {
                NullableSocket(sock)
        impl NullableSocket {
            pub fn invalid() -> Self {
                // TODO: may become UB in the future; maybe see rust-lang/rust#74699
                Self(unsafe { socket2::Socket::from_raw_socket(INVALID_SOCKET) })
            pub fn get(&self) -> Option<&socket2::Socket> {
                (self.0.as_raw_socket() != INVALID_SOCKET).then(|| &self.0)
            pub fn fd(&self) -> RawSocket {
                self.0.as_raw_socket()
            pub fn insert(&mut self, sock: socket2::Socket) -> &mut socket2::Socket {
                self.0 = sock;
                &mut self.0
    use nullable_socket::NullableSocket;
    impl Default for NullableSocket {
        fn default() -> Self {
            Self::invalid()
    #[pyattr(name = "socket")]
    #[pyattr(name = "SocketType")]
    #[pyclass(name = "socket")]
    #[derive(Debug, PyPayload)]
    pub struct PySocket {
        kind: AtomicCell<i32>,
        family: AtomicCell<i32>,
        proto: AtomicCell<i32>,
        pub(crate) timeout: AtomicCell<f64>,
        sock: PyRwLock<NullableSocket>,
    impl Default for PySocket {
        fn default() -> Self {
            PySocket {
                kind: AtomicCell::default(),
                family: AtomicCell::default(),
                proto: AtomicCell::default(),
                timeout: AtomicCell::new(-1.0),
                sock: PyRwLock::new(NullableSocket::invalid()),
    #[cfg(windows)]
    const CLOSED_ERR: i32 = c::WSAENOTSOCK;
    #[cfg(unix)]
    const CLOSED_ERR: i32 = c::EBADF;
    impl Read for &PySocket {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            (&mut &*self.sock()?).read(buf)
    impl Write for &PySocket {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            (&mut &*self.sock()?).write(buf)
        fn flush(&mut self) -> std::io::Result<()> {
            (&mut &*self.sock()?).flush()
    impl PySocket {
        pub fn sock_opt(&self) -> Option<PyMappedRwLockReadGuard<'_, Socket>> {
            PyRwLockReadGuard::try_map(self.sock.read(), |sock| sock.get()).ok()
        pub fn sock(&self) -> io::Result<PyMappedRwLockReadGuard<'_, Socket>> {
            self.sock_opt()
                .ok_or_else(|| io::Error::from_raw_os_error(CLOSED_ERR))
        fn init_inner(
            &self,
            family: i32,
            socket_kind: i32,
            proto: i32,
            sock: Socket,
        ) -> io::Result<()> {
            self.family.store(family);
            self.kind.store(socket_kind);
            self.proto.store(proto);
            let mu
802E
t s = self.sock.write();
            let sock = s.insert(sock);
            let timeout = DEFAULT_TIMEOUT.load();
            self.timeout.store(timeout);
            if timeout >= 0.0 {
                sock.set_nonblocking(true)?;
            Ok(())
        /// returns Err(blocking)
        pub fn get_timeout(&self) -> Result<Duration, bool> {
            let timeout = self.timeout.load();
            if timeout > 0.0 {
                Ok(Duration::from_secs_f64(timeout))
            } else {
                Err(timeout != 0.0)
        fn sock_op<F, R>(
            &self,
            vm: &VirtualMachine,
            select: SelectKind,
            f: F,
        ) -> Result<R, IoOrPyException>
        where
            F: FnMut() -> io::Result<R>,
            self.sock_op_timeout_err(vm, select, self.get_timeout().ok(), f)
        fn sock_op_timeout_err<F, R>(
            &self,
            vm: &VirtualMachine,
            select: SelectKind,
            timeout: Option<Duration>,
            mut f: F,
        ) -> Result<R, IoOrPyException>
        where
            F: FnMut() -> io::Result<R>,
            let deadline = timeout.map(Deadline::new);
            loop {
                if deadline.is_some() || matches!(select, SelectKind::Connect) {
                    let interval = deadline.as_ref().map(|d| d.time_until()).transpose()?;
                    let res = sock_select(&*self.sock()?, select, interval);
                    match res {
                        Ok(true) => return Err(IoOrPyException::Timeout),
                        Err(e) if e.kind() == io::ErrorKind::Interrupted => {
                            vm.check_signals()?;
                            continue;
                        Err(e) => return Err(e.into()),
                        Ok(false) => {} // no timeout, continue as normal
                let err = loop {
                    // loop on interrupt
                    match f() {
                        Ok(x) => return Ok(x),
                        Err(e) if e.kind() == io::ErrorKind::Interrupted => vm.check_signals()?,
                        Err(e) => break e,
                if timeout.is_some() && err.kind() == io::ErrorKind::WouldBlock {
                    continue;
                return Err(err.into());
        fn extract_address(
            &self,
            addr: PyObjectRef,
            caller: &str,
            vm: &VirtualMachine,
        ) -> Result<socket2::SockAddr, IoOrPyException> {
            let family = self.family.load();
            match family {
                #[cfg(unix)]
                c::AF_UNIX => {
                    use std::os::unix::ffi::OsStrExt;
                    let buf = crate::vm::function::ArgStrOrBytesLike::try_from_object(vm, addr)?;
                    let path = &*buf.borrow_bytes();
                    if cfg!(any(target_os = "linux", target_os = "android"))
                        && path.first() == Some(&0)
                        use libc::{sa_family_t, socklen_t};
                        use {socket2::SockAddr, std::ptr};
                        unsafe {
                            // based on SockAddr::unix
                            // TODO: upstream or fix socklen check for SockAddr::unix()?
                            SockAddr::init(|storage, len| {
                                // Safety: `SockAddr::init` zeros the address, which is a valid
                                // representation.
                                let storage: &mut libc::sockaddr_un = &mut *storage.cast();
                                let len: &mut socklen_t = &mut *len;
                                let bytes = path;
                                if bytes.len() > storage.sun_path.len() {
                                    return Err(io::Error::new(
                                        io::ErrorKind::InvalidInput,
                                        "path must be shorter than SUN_LEN",
                                storage.sun_family = libc::AF_UNIX as sa_family_t;
                                // Safety: `bytes` and `addr.sun_path` are not overlapping and
                                // both point to valid memory.
                                // `SockAddr::init` zeroes the memory, so the path is already
                                // null terminated.
                                ptr::copy_nonoverlapping(
                                    bytes.as_ptr(),
                                    storage.sun_path.as_mut_ptr() as *mut u8,
                                    bytes.len(),
                                let base = storage as *const _ as usize;
                                let path = &storage.sun_path as *const _ as usize;
                                let sun_path_offset = path - base;
                                let length = sun_path_offset + bytes.len();
                                *len = length as socklen_t;
                        .map(|(_, addr)| addr)
                    } else {
                        socket2::SockAddr::unix(ffi::OsStr::from_bytes(path))
                    .map_err(|_| vm.new_os_error("AF_UNIX path too long".to_owned()).into())
                c::AF_INET => {
                    let tuple: PyTupleRef = addr.downcast().map_err(|obj| {
                        vm.new_type_error(format!(
                            "{}(): AF_INET address must be tuple, not {}",
                            obj.class().name()
                    if tuple.len() != 2 {
                        return Err(vm
                            .new_type_error(
                                "AF_INET address must be a pair (host, post)".to_owned(),
                            .into());
                    let addr = Address::from_tuple(&tuple, vm)?;
                    let mut addr4 = get_addr(vm, addr.host, c::AF_INET)?;
                    match &mut addr4 {
                        SocketAddr::V4(addr4) => {
                            addr4.set_port(addr.port);
                        SocketAddr::V6(_) => unreachable!(),
                    Ok(addr4.into())
                c::AF_INET6 => {
                    let tuple: PyTupleRef = addr.downcast().map_err(|obj| {
                        vm.new_type_error(format!(
                            "{}(): AF_INET6 address must be tuple, not {}",
                            obj.class().name()
                    match tuple.len() {
                        _ => return Err(vm.new_type_error(
                            "AF_INET6 address must be a tuple (host, port[, flowinfo[, scopeid]])"
                                .to_owned(),
                        ).into()),
                    let (addr, flowinfo, scopeid) = Address::from_tuple_ipv6(&tuple, vm)?;
                    let mut addr6 = get_addr(vm, addr.host, c::AF_INET6)?;
                    match &mut addr6 {
                        SocketAddr::V6(addr6) => {
                            addr6.set_port(addr.port);
                            addr6.set_flowinfo(flowinfo);
                            addr6.set_scope_id(scopeid);
                        SocketAddr::V4(_) => unreachable!(),
                    Ok(addr6.into())
                _ => Err(vm.new_os_error(format!("{}(): bad family", caller)).into()),
        fn connect_inner(
            &self,
            address: PyObjectRef,
            caller: &str,
            vm: &VirtualMachine,
        ) -> Result<(), IoOrPyException> {
            let sock_addr = self.extract_address(address, caller, vm)?;
            let err = match self.sock()?.connect(&sock_addr) {
                Ok(()) => return Ok(()),
                Err(e) => e,
            let wait_connect = if err.kind() == io::ErrorKind::Interrupted {
                vm.check_signals()?;
                self.timeout.load() != 0.0
            } else {
                #[cfg(unix)]
                use c::EINPROGRESS;
                #[cfg(windows)]
                use c::WSAEWOULDBLOCK as EINPROGRESS;
                self.timeout.load() > 0.0 && err.raw_os_error() == Some(EINPROGRESS)
            if wait_connect {
                // basically, connect() is async, and it registers an "error" on the socket when it's
                // done connecting. SelectKind::Connect fills the errorfds fd_set, so if we wake up
                // from poll and the error is EISCONN then we know that the connect is done
                self.sock_op(vm, SelectKind::Connect, || {
                    let sock = self.sock()?;
                    let err = sock.take_error()?;
                    match err {
                        Some(e) if e.raw_os_error() == Some(libc::EISCONN) => Ok(()),
                        Some(e) => Err(e),
                        // TODO: is this accurate?
                        None => Ok(()),
            } else {
                Err(err.into())
    impl DefaultConstructor for PySocket {}
    impl Initializer for PySocket {
        type Args = (
            OptionalArg<i32>,
            OptionalArg<i32>,
            OptionalArg<i32>,
            OptionalOption<PyObjectRef>,
        fn init(zelf: PyRef<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
            Self::_init(zelf, args, vm).map_err(|e| e.into_pyexception(vm))
    #[pyimpl(with(DefaultConstructor, Initializer), flags(BASETYPE))]
    impl PySocket {
        fn _init(
            zelf: PyRef<Self>,
            (family, socket_kind, proto, fileno): <Self as Initializer>::Args,
            vm: &VirtualMachine,
        ) -> Result<(), IoOrPyException> {
            let mut family = family.unwrap_or(-1);
            let mut socket_kind = socket_kind.unwrap_or(-1);
            let mut proto = proto.unwrap_or(-1);
            let fileno = fileno
                .flatten()
                .map(|obj| get_raw_sock(obj, vm))
                .transpose()?;
            let sock;
            if let Some(fileno) = fileno {
                sock = sock_from_raw(fileno, vm)?;
                match sock.local_addr() {
                    Ok(addr) if family == -1 => family = addr.family() as i32,
                    Err(e)
                        if family == -1
                            || matches!(
                                e.raw_os_error(),
                                Some(errcode!(ENOTSOCK)) | Some(errcode!(EBADF))
                        std::mem::forget(sock);
                        return Err(e.into());
                if socket_kind == -1 {
                    // TODO: when socket2 cuts a new release, type will be available on all os
                    // socket_kind = sock.r#type().map_err(|e| e.into_pyexception(vm))?.into();
                    let res = unsafe {
                        c::getsockopt(
                            sock_fileno(&sock) as _,
                            c::SOL_SOCKET,
                            c::SO_TYPE,
                            &mut socket_kind as *mut libc::c_int as *mut _,
                            &mut (std::mem::size_of::<i32>() as _),
                    if res < 0 {
                        return Err(crate::common::os::errno().into());
                cfg_if::cfg_if! {
                    if #[cfg(any(
                        target_os = "android",
                        target_os = "freebsd",
                        target_os = "fuchsia",
                        target_os = "linux",
                        if proto == -1 {
                            proto = sock.protocol()?.map_or(0, Into::into);
                    } else {
                        proto = 0;
            } else {
                if family == -1 {
                    family = c::AF_INET as i32
                if socket_kind == -1 {
                    socket_kind = c::SOCK_STREAM
                if proto == -1 {
                    proto = 0
                sock = Socket::new(
                    Domain::from(family),
                    SocketType::from(socket_kind),
                    Some(Protocol::from(proto)),
            Ok(zelf.init_inner(family, socket_kind, proto, sock)?)
        #[pymethod]
        fn connect(
            &self,
            address: PyObjectRef,
            vm: &VirtualMachine,
        ) -> Result<(), IoOrPyException> {
            self.connect_inner(address, "connect", vm)
        #[pymethod]
        fn connect_ex(&self, address: PyObjectRef, vm: &VirtualMachine) -> PyResult<i32> {
            match self.connect_inner(address, "connect_ex", vm) {
                Ok(()) => Ok(0),
                Err(err) => err.errno(),
        #[pymethod]
        fn bind(&self, address: PyObjectRef, vm: &VirtualMachine) -> Result<(), IoOrPyException> {
            let sock_addr = self.extract_address(address, "bind", vm)?;
            Ok(self.sock()?.bind(&sock_addr)?)
        #[pymethod]
        fn listen(&self, backlog: OptionalArg<i32>) -> io::Result<()> {
            let backlog = backlog.unwrap_or(128);
            let backlog = if backlog < 0 { 0 } else { backlog };
            self.sock()?.listen(backlog)
        #[pymethod]
        fn _accept(
            &self,
            vm: &VirtualMachine,
        ) -> Result<(RawSocket, PyObjectRef), IoOrPyException> {
            let (sock, addr) = self.sock_op(vm, SelectKind::Read, || self.sock()?.accept())?;
            let fd = into_sock_fileno(sock);
            Ok((fd, get_addr_tuple(&addr, vm)))
        #[pymethod]
        fn recv(
            &self,
            bufsize: usize,
            flags: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> Result<Vec<u8>, IoOrPyException> {
            let flags = flags.unwrap_or(0);
            let mut buffer = Vec::with_capacity(bufsize);
            let sock = self.sock()?;
            let n = self.sock_op(vm, SelectKind::Read, || {
                sock.recv_with_flags(buffer.spare_capacity_mut(), flags)
            })?;
           
8B92
 unsafe { buffer.set_len(n) };
            Ok(buffer)
        #[pymethod]
        fn recv_into(
            &self,
            buf: ArgMemoryBuffer,
            flags: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> Result<usize, IoOrPyException> {
            let flags = flags.unwrap_or(0);
            let sock = self.sock()?;
            let mut buf = buf.borrow_buf_mut();
            let buf = &mut *buf;
            self.sock_op(vm, SelectKind::Read, || {
                sock.recv_with_flags(slice_as_uninit(buf), flags)
        #[pymethod]
        fn recvfrom(
            &self,
            bufsize: isize,
            flags: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> Result<(Vec<u8>, PyObjectRef), IoOrPyException> {
            let flags = flags.unwrap_or(0);
            let bufsize = bufsize
                .to_usize()
                .ok_or_else(|| vm.new_value_error("negative buffersize in recvfrom".to_owned()))?;
            let mut buffer = Vec::with_capacity(bufsize);
            let (n, addr) = self.sock_op(vm, SelectKind::Read, || {
                self.sock()?
                    .recv_from_with_flags(buffer.spare_capacity_mut(), flags)
            })?;
            unsafe { buffer.set_len(n) };
            Ok((buffer, get_addr_tuple(&addr, vm)))
        #[pymethod]
        fn recvfrom_into(
            &self,
            buf: ArgMemoryBuffer,
            nbytes: OptionalArg<isize>,
            flags: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> Result<(usize, PyObjectRef), IoOrPyException> {
            let mut buf = buf.borrow_buf_mut();
            let buf = &mut *buf;
            let buf = match nbytes {
                OptionalArg::Present(i) => {
                    let i = i.to_usize().ok_or_else(|| {
                        vm.new_value_error("negative buffersize in recvfrom_into".to_owned())
                    buf.get_mut(..i).ok_or_else(|| {
                        vm.new_value_error(
                            "nbytes is greater than the length of the buffer".to_owned(),
                OptionalArg::Missing => buf,
            let flags = flags.unwrap_or(0);
            let sock = self.sock()?;
            let (n, addr) = self.sock_op(vm, SelectKind::Read, || {
                sock.recv_from_with_flags(slice_as_uninit(buf), flags)
            })?;
            Ok((n, get_addr_tuple(&addr, vm)))
        #[pymethod]
        fn send(
            &self,
            bytes: ArgBytesLike,
            flags: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> Result<usize, IoOrPyException> {
            let flags = flags.unwrap_or(0);
            let buf = bytes.borrow_buf();
            let buf = &*buf;
            self.sock_op(vm, SelectKind::Write, || {
                self.sock()?.send_with_flags(buf, flags)
        #[pymethod]
        fn sendall(
            &self,
            bytes: ArgBytesLike,
            flags: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> Result<(), IoOrPyException> {
            let flags = flags.unwrap_or(0);
            let timeout = self.get_timeout().ok();
            let deadline = timeout.map(Deadline::new);
            let buf = bytes.borrow_buf();
            let buf = &*buf;
            let mut buf_offset = 0;
            // now we have like 3 layers of interrupt loop :)
            while buf_offset < buf.len() {
                let interval = deadline.as_ref().map(|d| d.time_until()).transpose()?;
                self.sock_op_timeout_err(vm, SelectKind::Write, interval, || {
                    let subbuf = &buf[buf_offset..];
                    buf_offset += self.sock()?.send_with_flags(subbuf, flags)?;
                    Ok(())
                vm.check_signals()?;
            Ok(())
        #[pymethod]
        fn sendto(
            &self,
            bytes: ArgBytesLike,
            arg2: PyObjectRef,
            arg3: OptionalArg<PyObjectRef>,
            vm: &VirtualMachine,
        ) -> Result<usize, IoOrPyException> {
            // signature is bytes[, flags], address
            let (flags, address) = match arg3 {
                OptionalArg::Present(arg3) => {
                    // should just be i32::try_from_obj but tests check for error message
                    let int = vm.to_index_opt(arg2).unwrap_or_else(|| {
                        Err(vm.new_type_error("an integer is required".to_owned()))
                    let flags = int.try_to_primitive::<i32>(vm)?;
                    (flags, arg3)
                OptionalArg::Missing => (0, arg2),
            let addr = self.extract_address(address, "sendto", vm)?;
            let buf = bytes.borrow_buf();
            let buf = &*buf;
            self.sock_op(vm, SelectKind::Write, || {
                self.sock()?.send_to_with_flags(buf, &addr, flags)
        #[pymethod]
        fn close(&self) -> io::Result<()> {
            let sock = self.detach();
            if sock != INVALID_SOCKET {
                close_inner(sock)?;
            Ok(())
        #[pymethod]
        #[inline]
        fn detach(&self) -> RawSocket {
            let sock = std::mem::replace(&mut *self.sock.write(), NullableSocket::invalid());
            std::mem::ManuallyDrop::new(sock).fd()
        #[pymethod]
        fn fileno(&self) -> RawSocket {
            self.sock.read().fd()
        #[pymethod]
        fn getsockname(&self, vm: &VirtualMachine) -> std::io::Result<PyObjectRef> {
            let addr = self.sock()?.local_addr()?;
            Ok(get_addr_tuple(&addr, vm))
        #[pymethod]
        fn getpeername(&self, vm: &VirtualMachine) -> std::io::Result<PyObjectRef> {
            let addr = self.sock()?.peer_addr()?;
            Ok(get_addr_tuple(&addr, vm))
        #[pymethod]
        fn gettimeout(&self) -> Option<f64> {
            let timeout = self.timeout.load();
            if timeout >= 0.0 {
                Some(timeout)
            } else {
        #[pymethod]
        fn setblocking(&self, block: bool) -> io::Result<()> {
            self.timeout.store(if block { -1.0 } else { 0.0 });
            self.sock()?.set_nonblocking(!block)
        #[pymethod]
        fn getblocking(&self) -> bool {
            self.timeout.load() != 0.0
        #[pymethod]
        fn settimeout(&self, timeout: Option<Duration>) -> io::Result<()> {
            self.timeout
                .store(timeout.map_or(-1.0, |d| d.as_secs_f64()));
            // even if timeout is > 0 the socket needs to be nonblocking in order for us to select() on
            // it
            self.sock()?.set_nonblocking(timeout.is_some())
        #[pymethod]
        fn getsockopt(
            &self,
            level: i32,
            name: i32,
            buflen: OptionalArg<i32>,
            vm: &VirtualMachine,
        ) -> PyResult {
            let fd = self.sock.read().fd() as _;
            let buflen = buflen.unwrap_or(0);
            if buflen == 0 {
                let mut flag: libc::c_int = 0;
                let mut flagsize = std::mem::size_of::<libc::c_int>() as _;
                let ret = unsafe {
                    c::getsockopt(
                        &mut flag as *mut libc::c_int as *mut _,
                        &mut flagsize,
                if ret < 0 {
                    return Err(crate::vm::stdlib::os::errno_err(vm));
                Ok(vm.ctx.new_int(flag).into())
            } else {
                if buflen <= 0 || buflen > 1024 {
                    return Err(vm.new_os_error("getsockopt buflen out of range".to_owned()));
                let mut buf = vec![0u8; buflen as usize];
                let mut buflen = buflen as _;
                let ret = unsafe {
                    c::getsockopt(fd, level, name, buf.as_mut_ptr() as *mut _, &mut buflen)
                buf.truncate(buflen as usize);
                if ret < 0 {
                    return Err(crate::vm::stdlib::os::errno_err(vm));
                Ok(vm.ctx.new_bytes(buf).into())
        #[pymethod]
        fn setsockopt(
            &self,
            level: i32,
            name: i32,
            value: Option<Either<ArgBytesLike, i32>>,
            optlen: OptionalArg<u32>,
            vm: &VirtualMachine,
        ) -> PyResult<()> {
            let fd = self.sock.read().fd() as _;
            let ret = match (value, optlen) {
                (Some(Either::A(b)), OptionalArg::Missing) => b.with_ref(|b| unsafe {
                    c::setsockopt(fd, level, name, b.as_ptr() as *const _, b.len() as _)
                (Some(Either::B(ref val)), OptionalArg::Missing) => unsafe {
                    c::setsockopt(
                        val as *const i32 as *const _,
                        std::mem::size_of::<i32>() as _,
                (None, OptionalArg::Present(optlen)) => unsafe {
                    c::setsockopt(fd, level, name, std::ptr::null(), optlen as _)
                    return Err(
                        vm.new_type_error("expected the value arg xor the optlen arg".to_owned())
            if ret < 0 {
                Err(crate::vm::stdlib::os::errno_err(vm))
            } else {
                Ok(())
        #[pymethod]
        fn shutdown(&self, how: i32, vm: &VirtualMachine) -> Result<(), IoOrPyException> {
            let how = match how {
                c::SHUT_RD => Shutdown::Read,
                c::SHUT_WR => Shutdown::Write,
                c::SHUT_RDWR => Shutdown::Both,
                    return Err(vm
                        .new_value_error("`how` must be SHUT_RD, SHUT_WR, or SHUT_RDWR".to_owned())
                        .into())
            Ok(self.sock()?.shutdown(how)?)
        #[pyproperty(name = "type")]
        fn kind(&self) -> i32 {
            self.kind.load()
        #[pyproperty]
        fn family(&self) -> i32 {
            self.family.load()
        #[pyproperty]
        fn proto(&self) -> i32 {
            self.proto.load()
        #[pymethod(magic)]
        fn repr(&self) -> String {
            format!(
                "<socket object, fd={}, family={}, type={}, proto={}>",
                // cast because INVALID_SOCKET is unsigned, so would show usize::MAX instead of -1
View remainder of file in raw view
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

socket.rs

Latest commit

History

socket.rs

File metadata and controls
