diff --git a/Cargo.lock b/Cargo.lock index 4dc06d285..63615acd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1330,6 +1330,7 @@ dependencies = [ "sha-1 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)", "sha2 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)", "sha3 0.8.2 (registry+https://github.com/rust-lang/crates.io-index)", + "socket2 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)", "statrs 0.12.0 (registry+https://github.com/rust-lang/crates.io-index)", "subprocess 0.1.18 (registry+https://github.com/rust-lang/crates.io-index)", "unic 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)", @@ -1478,6 +1479,17 @@ name = "smallvec" version = "0.6.10" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "socket2" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "cfg-if 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.62 (registry+https://github.com/rust-lang/crates.io-index)", + "redox_syscall 0.1.56 (registry+https://github.com/rust-lang/crates.io-index)", + "winapi 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "sourcefile" version = "0.1.4" @@ -2281,6 +2293,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" "checksum sha3 0.8.2 (registry+https://github.com/rust-lang/crates.io-index)" = "dd26bc0e7a2e3a7c959bc494caf58b72ee0c71d67704e9520f736ca7e4853ecf" "checksum siphasher 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)" = "0b8de496cf83d4ed58b6be86c3a275b8602f6ffe98d3024a869e124147a9a3ac" "checksum smallvec 0.6.10 (registry+https://github.com/rust-lang/crates.io-index)" = "ab606a9c5e214920bb66c458cd7be8ef094f813f20fe77a54cc7dbfff220d4b7" +"checksum socket2 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)" = "e8b74de517221a2cb01a53349cf54182acdc31a074727d3079068448c0676d85" "checksum sourcefile 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)" = "4bf77cb82ba8453b42b6ae1d692e4cdc92f9a47beaf89a847c8be83f4e328ad3" "checksum static_assertions 0.3.4 (registry+https://github.com/rust-lang/crates.io-index)" = "7f3eb36b47e512f8f1c9e3d10c2c1965bc992bd9cdb024fa581e2194501c83d3" "checksum statrs 0.12.0 (registry+https://github.com/rust-lang/crates.io-index)" = "cce16f6de653e88beca7bd13780d08e09d4489dbca1f9210e041bc4852481382" diff --git a/vm/Cargo.toml b/vm/Cargo.toml index 222277095..013f05021 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -83,7 +83,8 @@ flate2 = { version = "1.0", features = ["zlib"], default-features = false } libz-sys = "1.0.25" gethostname = "0.2.0" subprocess = "0.1.18" -num_cpus = "1.0" +num_cpus = "1" +socket2 = { version = "0.3", features = ["unix"] } [target."cfg(windows)".dependencies] winapi = { version = "0.3", features = ["winsock2", "handleapi", "std"] } diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs index 91eeb3624..9361feaf7 100644 --- a/vm/src/obj/objbytearray.rs +++ b/vm/src/obj/objbytearray.rs @@ -37,7 +37,7 @@ use std::mem::size_of; pub struct PyByteArray { pub inner: RefCell, } -type PyByteArrayRef = PyRef; +pub type PyByteArrayRef = PyRef; impl PyByteArray { pub fn new(data: Vec) -> Self { diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs index 164fe3740..e481cf77b 100644 --- a/vm/src/obj/objbyteinner.rs +++ b/vm/src/obj/objbyteinner.rs @@ -5,8 +5,8 @@ use num_bigint::BigInt; use num_integer::Integer; use num_traits::ToPrimitive; -use super::objbytearray::PyByteArray; -use super::objbytes::PyBytes; +use super::objbytearray::{PyByteArray, PyByteArrayRef}; +use super::objbytes::{PyBytes, PyBytesRef}; use super::objint::{self, PyInt, PyIntRef}; use super::objlist::PyList; use super::objmemory::PyMemoryView; @@ -1303,3 +1303,34 @@ fn split_slice_reverse<'a>(slice: &'a [u8], sep: &[u8], maxsplit: i32) -> Vec<&' splitted.reverse(); splitted } + +pub enum PyBytesLike { + Bytes(PyBytesRef), + Bytearray(PyByteArrayRef), + Vec(Vec), +} + +impl TryFromObject for PyBytesLike { + fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + match_class!(match obj { + b @ PyBytes => Ok(PyBytesLike::Bytes(b)), + b @ PyByteArray => Ok(PyBytesLike::Bytearray(b)), + m @ PyMemoryView => Ok(PyBytesLike::Vec(m.get_obj_value().unwrap())), + l @ PyList => Ok(PyBytesLike::Vec(l.get_byte_inner(vm)?.elements)), + obj => Err(vm.new_type_error(format!( + "a bytes-like object is required, not {}", + obj.class() + ))), + }) + } +} + +impl PyBytesLike { + pub fn to_cow(&self) -> std::borrow::Cow<[u8]> { + match self { + PyBytesLike::Bytes(b) => b.get_value().into(), + PyBytesLike::Bytearray(b) => b.inner.borrow().elements.clone().into(), + PyBytesLike::Vec(b) => b.as_slice().into(), + } + } +} diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index 82ca42f53..b1d5f66c7 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -1,493 +1,258 @@ -use std::cell::RefCell; +use std::cell::Cell; use std::io; -use std::io::Read; -use std::io::Write; -use std::net::{Ipv4Addr, SocketAddr, TcpListener, TcpStream, ToSocketAddrs, UdpSocket}; +use std::net::{Ipv4Addr, Shutdown, SocketAddr, ToSocketAddrs}; use std::time::Duration; use byteorder::{BigEndian, ByteOrder}; use gethostname::gethostname; #[cfg(all(unix, not(target_os = "redox")))] use nix::unistd::sethostname; -use num_traits::ToPrimitive; +use socket2::{Domain, Socket, Type as SocketType}; -use crate::function::PyFuncArgs; +use super::os::convert_io_error; +#[cfg(unix)] +use super::os::convert_nix_error; +use super::time_module::duration_to_f64; +use crate::function::OptionalArg; +use crate::obj::objbyteinner::PyBytesLike; use crate::obj::objbytes::PyBytesRef; -use crate::obj::objint::PyIntRef; use crate::obj::objstr::PyStringRef; use crate::obj::objtuple::PyTupleRef; use crate::obj::objtype::PyClassRef; use crate::pyobject::{PyObjectRef, PyRef, PyResult, PyValue, TryFromObject}; -#[cfg(unix)] -use crate::stdlib::os::convert_nix_error; use crate::vm::VirtualMachine; -#[derive(Debug, Copy, Clone)] -enum AddressFamily { - Unix = 1, - Inet = 2, - Inet6 = 3, -} - -impl TryFromObject for AddressFamily { - fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { - match i32::try_from_object(vm, obj)? { - 1 => Ok(AddressFamily::Unix), - 2 => Ok(AddressFamily::Inet), - 3 => Ok(AddressFamily::Inet6), - value => Err(vm.new_os_error(format!("Unknown address family value: {}", value))), - } - } -} - -#[derive(Debug, Copy, Clone)] -enum SocketKind { - Stream = 1, - Dgram = 2, -} - -impl TryFromObject for SocketKind { - fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { - match i32::try_from_object(vm, obj)? { - 1 => Ok(SocketKind::Stream), - 2 => Ok(SocketKind::Dgram), - value => Err(vm.new_os_error(format!("Unknown socket kind value: {}", value))), - } - } -} - #[derive(Debug)] -enum Connection { - TcpListener(TcpListener), - TcpStream(TcpStream), - UdpSocket(UdpSocket), +pub struct PySocket { + timeout: Cell>, + sock: Socket, } -impl Connection { - fn accept(&mut self) -> io::Result<(TcpStream, SocketAddr)> { - match self { - Connection::TcpListener(con) => con.accept(), - _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), - } - } - - fn local_addr(&self) -> io::Result { - match self { - Connection::TcpListener(con) => con.local_addr(), - Connection::UdpSocket(con) => con.local_addr(), - Connection::TcpStream(con) => con.local_addr(), - } - } - - fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { - match self { - Connection::UdpSocket(con) => con.recv_from(buf), - _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), - } - } - - fn send_to(&self, buf: &[u8], addr: A) -> io::Result { - match self { - Connection::UdpSocket(con) => con.send_to(buf, addr), - _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), - } - } - - #[cfg(unix)] - fn fileno(&self) -> i64 { - use std::os::unix::io::AsRawFd; - let raw_fd = match self { - Connection::TcpListener(con) => con.as_raw_fd(), - Connection::UdpSocket(con) => con.as_raw_fd(), - Connection::TcpStream(con) => con.as_raw_fd(), - }; - i64::from(raw_fd) - } - - #[cfg(windows)] - fn fileno(&self) -> i64 { - use std::os::windows::io::AsRawSocket; - let raw_fd = match self { - Connection::TcpListener(con) => con.as_raw_socket(), - Connection::UdpSocket(con) => con.as_raw_socket(), - Connection::TcpStream(con) => con.as_raw_socket(), - }; - raw_fd as i64 - } - - #[cfg(all(not(unix), not(windows)))] - fn fileno(&self) -> i64 { - unimplemented!(); - } - - fn setblocking(&mut self, value: bool) -> io::Result<()> { - match self { - Connection::TcpListener(con) => con.set_nonblocking(!value), - Connection::UdpSocket(con) => con.set_nonblocking(!value), - Connection::TcpStream(con) => con.set_nonblocking(!value), - } - } - - fn settimeout(&mut self, duration: Duration) -> io::Result<()> { - match self { - // net - Connection::TcpListener(_con) => Ok(()), - Connection::UdpSocket(con) => con - .set_read_timeout(Some(duration)) - .and_then(|_| con.set_write_timeout(Some(duration))), - Connection::TcpStream(con) => con - .set_read_timeout(Some(duration)) - .and_then(|_| con.set_write_timeout(Some(duration))), - } - } -} - -impl Read for Connection { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - match self { - Connection::TcpStream(con) => con.read(buf), - Connection::UdpSocket(con) => con.recv(buf), - _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), - } - } -} - -impl Write for Connection { - fn write(&mut self, buf: &[u8]) -> io::Result { - match self { - Connection::TcpStream(con) => con.write(buf), - Connection::UdpSocket(con) => con.send(buf), - _ => Err(io::Error::new(io::ErrorKind::Other, "oh no!")), - } - } - fn flush(&mut self) -> io::Result<()> { - Ok(()) - } -} - -#[derive(Debug)] -pub struct Socket { - address_family: AddressFamily, - socket_kind: SocketKind, - con: RefCell>, - timeout: RefCell>, -} - -impl PyValue for Socket { +impl PyValue for PySocket { fn class(vm: &VirtualMachine) -> PyClassRef { vm.class("socket", "socket") } } -impl Socket { - fn new(address_family: AddressFamily, socket_kind: SocketKind) -> Socket { - Socket { - address_family, - socket_kind, - con: RefCell::new(None), - timeout: RefCell::new(None), - } - } -} +type PySocketRef = PyRef; -type SocketRef = PyRef; - -impl SocketRef { +impl PySocketRef { fn new( cls: PyClassRef, - family: AddressFamily, - kind: SocketKind, + domain: OptionalArg, + socket_type: OptionalArg, vm: &VirtualMachine, - ) -> PyResult { - Socket::new(family, kind).into_ref_with_type(vm, cls) - } - - fn enter(self, _vm: &VirtualMachine) -> SocketRef { - self - } - - fn exit(self, _args: PyFuncArgs, _vm: &VirtualMachine) { - self.close(_vm) + ) -> PyResult { + let domain = domain.unwrap_or(libc::AF_INET); + let socket_type = socket_type.unwrap_or(libc::SOCK_STREAM); + let domain = match domain { + libc::AF_INET => Domain::ipv4(), + libc::AF_INET6 => Domain::ipv6(), + #[cfg(unix)] + libc::AF_UNIX => Domain::unix(), + _ => return Err(vm.new_os_error(format!("Unknown address family value: {}", domain))), + }; + let socket_type = match socket_type { + libc::SOCK_STREAM => SocketType::stream(), + libc::SOCK_DGRAM => SocketType::dgram(), + _ => return Err(vm.new_os_error(format!("Unknown socket kind value: {}", socket_type))), + }; + let sock = + Socket::new(domain, socket_type, None).map_err(|err| convert_io_error(vm, err))?; + PySocket { + sock, + timeout: Cell::default(), + } + .into_ref_with_type(vm, cls) } fn connect(self, address: Address, vm: &VirtualMachine) -> PyResult<()> { - let address_string = address.get_address_string(); - - match self.socket_kind { - SocketKind::Stream => { - let con = if let Some(duration) = self.timeout.borrow().as_ref() { - let sock_addr = match address_string.to_socket_addrs() { - Ok(mut sock_addrs) => { - if sock_addrs.len() == 0 { - let error_type = vm.class("socket", "gaierror"); - return Err(vm.new_exception( - error_type, - "nodename nor servname provided, or not known".to_string(), - )); - } else { - sock_addrs.next().unwrap() - } - } - Err(e) => { - let error_type = vm.class("socket", "gaierror"); - return Err(vm.new_exception(error_type, e.to_string())); - } - }; - TcpStream::connect_timeout(&sock_addr, *duration) - } else { - TcpStream::connect(address_string) - }; - match con { - Ok(stream) => { - self.con.borrow_mut().replace(Connection::TcpStream(stream)); - Ok(()) - } - Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { - let socket_timeout = vm.class("socket", "timeout"); - Err(vm.new_exception(socket_timeout, "Timed out".to_string())) - } - Err(s) => Err(vm.new_os_error(s.to_string())), - } - } - SocketKind::Dgram => { - if let Some(Connection::UdpSocket(con)) = self.con.borrow().as_ref() { - match con.connect(address_string) { - Ok(_) => Ok(()), - Err(s) => Err(vm.new_os_error(s.to_string())), - } - } else { - Err(vm.new_type_error("".to_string())) - } - } - } + let sock_addr = get_addr(vm, address)?; + let res = if let Some(duration) = self.timeout.get() { + self.sock.connect_timeout(&sock_addr, duration) + } else { + self.sock.connect(&sock_addr) + }; + res.map_err(|err| convert_io_error(vm, err)) } fn bind(self, address: Address, vm: &VirtualMachine) -> PyResult<()> { - let address_string = address.get_address_string(); - - match self.socket_kind { - SocketKind::Stream => match TcpListener::bind(address_string) { - Ok(stream) => { - self.con - .borrow_mut() - .replace(Connection::TcpListener(stream)); - Ok(()) - } - Err(s) => Err(vm.new_os_error(s.to_string())), - }, - SocketKind::Dgram => match UdpSocket::bind(address_string) { - Ok(dgram) => { - self.con.borrow_mut().replace(Connection::UdpSocket(dgram)); - Ok(()) - } - Err(s) => Err(vm.new_os_error(s.to_string())), - }, - } + let sock_addr = get_addr(vm, address)?; + self.sock + .bind(&sock_addr) + .map_err(|err| convert_io_error(vm, err)) } - fn listen(self, _num: PyIntRef, _vm: &VirtualMachine) {} + fn listen(self, backlog: OptionalArg, vm: &VirtualMachine) -> PyResult<()> { + let backlog = backlog.unwrap_or(128); + let backlog = if backlog < 0 { 0 } else { backlog }; + self.sock + .listen(backlog) + .map_err(|err| convert_io_error(vm, err)) + } fn accept(self, vm: &VirtualMachine) -> PyResult { - let ret = match self.con.borrow_mut().as_mut() { - Some(v) => v.accept(), - None => return Err(vm.new_type_error("".to_string())), - }; + let (sock, addr) = self + .sock + .accept() + .map_err(|err| convert_io_error(vm, err))?; - let (tcp_stream, addr) = match ret { - Ok((socket, addr)) => (socket, addr), - Err(s) => return Err(vm.new_os_error(s.to_string())), - }; - - let socket = Socket { - address_family: self.address_family, - socket_kind: self.socket_kind, - con: RefCell::new(Some(Connection::TcpStream(tcp_stream))), - timeout: RefCell::new(None), + let socket = PySocket { + sock, + timeout: self.timeout.clone(), } .into_ref(vm); - let addr_tuple = get_addr_tuple(vm, addr)?; + let addr_tuple = get_addr_tuple(vm, addr); Ok(vm.ctx.new_tuple(vec![socket.into_object(), addr_tuple])) } fn recv(self, bufsize: usize, vm: &VirtualMachine) -> PyResult { let mut buffer = vec![0u8; bufsize]; - match self.con.borrow_mut().as_mut() { - Some(v) => match v.read_exact(&mut buffer) { - Ok(_) => (), - Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { - let socket_timeout = vm.class("socket", "timeout"); - return Err(vm.new_exception(socket_timeout, "Timed out".to_string())); - } - Err(s) => return Err(vm.new_os_error(s.to_string())), - }, - None => return Err(vm.new_type_error("".to_string())), - }; - Ok(vm.ctx.new_bytes(buffer)) + match self.sock.recv(&mut buffer) { + Ok(_) => Ok(vm.ctx.new_bytes(buffer)), + Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { + let socket_timeout = vm.class("socket", "timeout"); + Err(vm.new_exception(socket_timeout, "Timed out".to_string())) + } + Err(err) => Err(convert_io_error(vm, err)), + } } fn recvfrom(self, bufsize: usize, vm: &VirtualMachine) -> PyResult { let mut buffer = vec![0u8; bufsize]; - let ret = match self.con.borrow().as_ref() { - Some(v) => v.recv_from(&mut buffer), - None => return Err(vm.new_type_error("".to_string())), + let addr = match self.sock.recv_from(&mut buffer) { + Ok((_, addr)) => addr, + Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { + let socket_timeout = vm.class("socket", "timeout"); + return Err(vm.new_exception(socket_timeout, "Timed out".to_string())); + } + Err(err) => return Err(convert_io_error(vm, err)), }; - let addr = match ret { - Ok((_size, addr)) => addr, - Err(s) => return Err(vm.new_os_error(s.to_string())), - }; - - let addr_tuple = get_addr_tuple(vm, addr)?; + let addr_tuple = get_addr_tuple(vm, addr); Ok(vm.ctx.new_tuple(vec![vm.ctx.new_bytes(buffer), addr_tuple])) } - fn send(self, bytes: PyBytesRef, vm: &VirtualMachine) -> PyResult<()> { - match self.con.borrow_mut().as_mut() { - Some(v) => match v.write(&bytes) { - Ok(_) => (), - Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { - let socket_timeout = vm.class("socket", "timeout"); - return Err(vm.new_exception(socket_timeout, "Timed out".to_string())); - } - Err(s) => return Err(vm.new_os_error(s.to_string())), - }, - None => return Err(vm.new_type_error("Socket is not connected".to_string())), - }; + fn send(self, bytes: PyBytesLike, vm: &VirtualMachine) -> PyResult { + match self.sock.send(bytes.to_cow().as_ref()) { + Ok(i) => Ok(i), + Err(ref e) if e.kind() == io::ErrorKind::TimedOut => { + let socket_timeout = vm.class("socket", "timeout"); + Err(vm.new_exception(socket_timeout, "Timed out".to_string())) + } + Err(err) => Err(convert_io_error(vm, err)), + } + } + + fn sendto(self, bytes: PyBytesLike, address: Address, vm: &VirtualMachine) -> PyResult<()> { + let addr = get_addr(vm, address)?; + self.sock + .send_to(bytes.to_cow().as_ref(), &addr) + .map_err(|err| convert_io_error(vm, err))?; Ok(()) } - fn sendto(self, bytes: PyBytesRef, address: Address, vm: &VirtualMachine) -> PyResult<()> { - let address_string = address.get_address_string(); - - match self.socket_kind { - SocketKind::Dgram => { - if let Some(v) = self.con.borrow().as_ref() { - return match v.send_to(&bytes, address_string) { - Ok(_) => Ok(()), - Err(s) => Err(vm.new_os_error(s.to_string())), - }; - } - // Doing implicit bind - match UdpSocket::bind("0.0.0.0:0") { - Ok(dgram) => match dgram.send_to(&bytes, address_string) { - Ok(_) => { - self.con.borrow_mut().replace(Connection::UdpSocket(dgram)); - Ok(()) - } - Err(s) => Err(vm.new_os_error(s.to_string())), - }, - Err(s) => Err(vm.new_os_error(s.to_string())), - } - } - _ => Err(vm.new_not_implemented_error("".to_string())), + fn close(self, vm: &VirtualMachine) -> PyResult<()> { + let fd = self.clone().fileno(vm); + #[cfg(unix)] + let ret = unsafe { libc::close(fd) }; + #[cfg(windows)] + let ret = { + extern "system" { + fn closesocket(s: std::os::windows::raw::SOCKET) -> c_int; + }; + unsafe { closesocket(fd) } + }; + if ret != 0 { + Err(convert_io_error(vm, io::Error::last_os_error())) + } else { + Ok(()) } } - fn close(self, _vm: &VirtualMachine) { - self.con.borrow_mut().take(); + #[cfg(unix)] + fn fileno(self, _vm: &VirtualMachine) -> std::os::unix::io::RawFd { + use std::os::unix::io::AsRawFd; + self.sock.as_raw_fd() } - - fn fileno(self, vm: &VirtualMachine) -> PyResult { - let fileno = match self.con.borrow_mut().as_mut() { - Some(v) => v.fileno(), - None => return Err(vm.new_type_error("".to_string())), - }; - Ok(vm.ctx.new_int(fileno)) + #[cfg(windows)] + fn fileno(self, _vm: &VirtualMachine) -> std::os::windows::raw::SOCKET { + use std::os::windows::io::AsRawSocket; + self.sock.as_raw_socket() } fn getsockname(self, vm: &VirtualMachine) -> PyResult { - let addr = match self.con.borrow().as_ref() { - Some(v) => v.local_addr(), - None => return Err(vm.new_type_error("".to_string())), + let addr = self + .sock + .local_addr() + .map_err(|err| convert_io_error(vm, err))?; + + Ok(get_addr_tuple(vm, addr)) + } + + fn gettimeout(self, _vm: &VirtualMachine) -> Option { + self.timeout.get().map(duration_to_f64) + } + + fn setblocking(self, block: bool, vm: &VirtualMachine) -> PyResult<()> { + self.sock + .set_nonblocking(!block) + .map_err(|err| convert_io_error(vm, err))?; + if block { + self.timeout.set(None); + } else { + self.timeout.set(Some(Duration::from_secs(0))); + } + Ok(()) + } + + fn getblocking(self, _vm: &VirtualMachine) -> bool { + self.timeout.get().map_or(false, |d| d.as_secs() != 0) + } + + fn settimeout(self, timeout: f64, vm: &VirtualMachine) -> PyResult<()> { + let secs: u64 = timeout.trunc() as u64; + let nanos: u32 = (timeout.fract() * 1e9) as u32; + let duration = Duration::new(secs, nanos); + + self.timeout.set(Some(duration)); + + self.sock + .set_nonblocking(duration.as_secs() == 0) + .map_err(|err| convert_io_error(vm, err))?; + + Ok(()) + } + + fn shutdown(self, how: i32, vm: &VirtualMachine) -> PyResult<()> { + let how = match how { + libc::SHUT_RD => Shutdown::Read, + libc::SHUT_WR => Shutdown::Write, + libc::SHUT_RDWR => Shutdown::Both, + _ => { + return Err( + vm.new_value_error("`how` must be SHUT_RD, SHUT_WR, or SHUT_RDWR".to_string()) + ) + } }; - - match addr { - Ok(addr) => get_addr_tuple(vm, addr), - Err(s) => Err(vm.new_os_error(s.to_string())), - } - } - - fn gettimeout(self, _vm: &VirtualMachine) -> PyResult> { - match self.timeout.borrow().as_ref() { - Some(duration) => Ok(Some(duration.as_secs() as f64)), - None => Ok(None), - } - } - - fn setblocking(self, block: Option, vm: &VirtualMachine) -> PyResult<()> { - match block { - Some(value) => { - if value { - self.timeout.replace(None); - } else { - self.timeout.borrow_mut().replace(Duration::from_secs(0)); - } - if let Some(conn) = self.con.borrow_mut().as_mut() { - match conn.setblocking(value) { - Ok(_) => Ok(()), - Err(err) => Err(vm.new_os_error(err.to_string())), - } - } else { - Ok(()) - } - } - None => { - // Avoid converting None to bool - Err(vm.new_type_error("an bool is required".to_string())) - } - } - } - - fn getblocking(self, _vm: &VirtualMachine) -> PyResult> { - match self.timeout.borrow().as_ref() { - Some(duration) => { - if duration.as_secs() != 0 { - Ok(Some(true)) - } else { - Ok(Some(false)) - } - } - None => Ok(Some(true)), - } - } - - fn settimeout(self, timeout: Option, vm: &VirtualMachine) -> PyResult<()> { - match timeout { - Some(timeout) => { - self.timeout - .borrow_mut() - .replace(Duration::from_secs(timeout as u64)); - - let block = timeout > 0.0; - - if let Some(conn) = self.con.borrow_mut().as_mut() { - conn.setblocking(block) - .and_then(|_| conn.settimeout(Duration::from_secs(timeout as u64))) - .map_err(|err| vm.new_os_error(err.to_string())) - .map(|_| ()) - } else { - Ok(()) - } - } - None => { - self.timeout.replace(None); - Ok(()) - } - } + self.sock + .shutdown(how) + .map_err(|err| convert_io_error(vm, err)) } } struct Address { - host: String, - port: usize, + host: PyStringRef, + port: u16, } -impl Address { - fn get_address_string(self) -> String { - format!("{}:{}", self.host, self.port.to_string()) +impl ToSocketAddrs for Address { + type Iter = std::vec::IntoIter; + fn to_socket_addrs(&self) -> io::Result { + (self.host.as_str(), self.port).to_socket_addrs() } } @@ -498,23 +263,26 @@ impl TryFromObject for Address { Err(vm.new_type_error("Address tuple should have only 2 values".to_string())) } else { Ok(Address { - host: PyStringRef::try_from_object(vm, tuple.elements[0].clone())? - .as_str() - .to_owned(), - port: PyIntRef::try_from_object(vm, tuple.elements[1].clone())? - .as_bigint() - .to_usize() - .unwrap(), + host: PyStringRef::try_from_object(vm, tuple.elements[0].clone())?, + port: u16::try_from_object(vm, tuple.elements[1].clone())?, }) } } } -fn get_addr_tuple(vm: &VirtualMachine, addr: SocketAddr) -> PyResult { - let port = vm.ctx.new_int(addr.port()); - let ip = vm.ctx.new_str(addr.ip().to_string()); +fn get_addr_tuple>(vm: &VirtualMachine, addr: A) -> PyObjectRef { + let addr = addr.into(); + let (port, ip) = if let Some(addr) = addr.as_inet() { + (addr.port(), addr.ip().to_string()) + } else if let Some(addr) = addr.as_inet6() { + (addr.port(), addr.ip().to_string()) + } else { + (0, String::new()) + }; + let port = vm.ctx.new_int(port); + let ip = vm.ctx.new_str(ip); - Ok(vm.ctx.new_tuple(vec![ip, port])) + vm.ctx.new_tuple(vec![ip, port]) } fn socket_gethostname(vm: &VirtualMachine) -> PyResult { @@ -541,7 +309,7 @@ fn socket_inet_ntoa(packed_ip: PyBytesRef, vm: &VirtualMachine) -> PyResult { if packed_ip.len() != 4 { return Err(vm.new_os_error("packed IP wrong length for inet_ntoa".to_string())); } - let ip_num = BigEndian::read_u32(packed_ip.get_value()); + let ip_num = BigEndian::read_u32(&packed_ip); Ok(vm.new_str(Ipv4Addr::from(ip_num).to_string())) } @@ -549,39 +317,66 @@ fn socket_htonl(host: u32, vm: &VirtualMachine) -> PyResult { Ok(vm.new_int(host.to_be())) } +fn get_addr(vm: &VirtualMachine, addr: T) -> PyResult +where + T: ToSocketAddrs, + I: ExactSizeIterator, +{ + match addr.to_socket_addrs() { + Ok(mut sock_addrs) => { + if sock_addrs.len() == 0 { + let error_type = vm.class("socket", "gaierror"); + Err(vm.new_exception( + error_type, + "nodename nor servname provided, or not known".to_string(), + )) + } else { + Ok(sock_addrs.next().unwrap().into()) + } + } + Err(e) => { + let error_type = vm.class("socket", "gaierror"); + Err(vm.new_exception(error_type, e.to_string())) + } + } +} + pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { let ctx = &vm.ctx; let socket_timeout = ctx.new_class("socket.timeout", vm.ctx.exceptions.os_error.clone()); let socket_gaierror = ctx.new_class("socket.gaierror", vm.ctx.exceptions.os_error.clone()); let socket = py_class!(ctx, "socket", ctx.object(), { - (slot new) => SocketRef::new, - "__enter__" => ctx.new_rustfunc(SocketRef::enter), - "__exit__" => ctx.new_rustfunc(SocketRef::exit), - "connect" => ctx.new_rustfunc(SocketRef::connect), - "recv" => ctx.new_rustfunc(SocketRef::recv), - "send" => ctx.new_rustfunc(SocketRef::send), - "bind" => ctx.new_rustfunc(SocketRef::bind), - "accept" => ctx.new_rustfunc(SocketRef::accept), - "listen" => ctx.new_rustfunc(SocketRef::listen), - "close" => ctx.new_rustfunc(SocketRef::close), - "getsockname" => ctx.new_rustfunc(SocketRef::getsockname), - "sendto" => ctx.new_rustfunc(SocketRef::sendto), - "recvfrom" => ctx.new_rustfunc(SocketRef::recvfrom), - "fileno" => ctx.new_rustfunc(SocketRef::fileno), - "getblocking" => ctx.new_rustfunc(SocketRef::getblocking), - "setblocking" => ctx.new_rustfunc(SocketRef::setblocking), - "gettimeout" => ctx.new_rustfunc(SocketRef::gettimeout), - "settimeout" => ctx.new_rustfunc(SocketRef::settimeout), + (slot new) => PySocketRef::new, + "connect" => ctx.new_rustfunc(PySocketRef::connect), + "recv" => ctx.new_rustfunc(PySocketRef::recv), + "send" => ctx.new_rustfunc(PySocketRef::send), + "bind" => ctx.new_rustfunc(PySocketRef::bind), + "accept" => ctx.new_rustfunc(PySocketRef::accept), + "listen" => ctx.new_rustfunc(PySocketRef::listen), + "close" => ctx.new_rustfunc(PySocketRef::close), + "getsockname" => ctx.new_rustfunc(PySocketRef::getsockname), + "sendto" => ctx.new_rustfunc(PySocketRef::sendto), + "recvfrom" => ctx.new_rustfunc(PySocketRef::recvfrom), + "fileno" => ctx.new_rustfunc(PySocketRef::fileno), + "getblocking" => ctx.new_rustfunc(PySocketRef::getblocking), + "setblocking" => ctx.new_rustfunc(PySocketRef::setblocking), + "gettimeout" => ctx.new_rustfunc(PySocketRef::gettimeout), + "settimeout" => ctx.new_rustfunc(PySocketRef::settimeout), + "shutdown" => ctx.new_rustfunc(PySocketRef::shutdown), }); let module = py_module!(vm, "socket", { "error" => ctx.exceptions.os_error.clone(), "timeout" => socket_timeout, "gaierror" => socket_gaierror, - "AF_INET" => ctx.new_int(AddressFamily::Inet as i32), - "SOCK_STREAM" => ctx.new_int(SocketKind::Stream as i32), - "SOCK_DGRAM" => ctx.new_int(SocketKind::Dgram as i32), + "AF_INET" => ctx.new_int(libc::AF_INET), + "AF_INET6" => ctx.new_int(libc::AF_INET6), + "SOCK_STREAM" => ctx.new_int(libc::SOCK_STREAM), + "SOCK_DGRAM" => ctx.new_int(libc::SOCK_DGRAM), + "SHUT_RD" => ctx.new_int(libc::SHUT_RD), + "SHUT_WR" => ctx.new_int(libc::SHUT_WR), + "SHUT_RDWR" => ctx.new_int(libc::SHUT_RDWR), "socket" => socket, "inet_aton" => ctx.new_rustfunc(socket_inet_aton), "inet_ntoa" => ctx.new_rustfunc(socket_inet_ntoa), @@ -589,22 +384,20 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "htonl" => ctx.new_rustfunc(socket_htonl), }); - extend_module_platform_specific(vm, module) -} + extend_module_platform_specific(vm, &module); -#[cfg(not(unix))] -fn extend_module_platform_specific(_vm: &VirtualMachine, module: PyObjectRef) -> PyObjectRef { module } +#[cfg(not(unix))] +fn extend_module_platform_specific(_vm: &VirtualMachine, _module: &PyObjectRef) {} + #[cfg(unix)] -fn extend_module_platform_specific(vm: &VirtualMachine, module: PyObjectRef) -> PyObjectRef { +fn extend_module_platform_specific(vm: &VirtualMachine, module: &PyObjectRef) { let ctx = &vm.ctx; #[cfg(not(target_os = "redox"))] extend_module!(vm, module, { "sethostname" => ctx.new_rustfunc(socket_sethostname), }); - - module }