From 9647fce6d5c944fefd82583b60de982206211775 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Mon, 24 Aug 2020 19:22:49 -0500 Subject: [PATCH] Add some more socket methods + other misc fixes --- Cargo.toml | 2 -- vm/Cargo.toml | 2 +- vm/src/py_io.rs | 19 ++++++++++ vm/src/stdlib/socket.rs | 80 ++++++++++++++++++++++++++++++----------- 4 files changed, 80 insertions(+), 23 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 431053833..94a530eb1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -68,7 +68,5 @@ opt-level = 3 [patch.crates-io] # REDOX START, Uncommment when you want to compile/check with redoxer -# # following patches are just waiting on a new version to be released to crates.io -# socket2 = { git = "https://github.com/alexcrichton/socket2-rs", branch = "v0.3.x" } # REDOX END diff --git a/vm/Cargo.toml b/vm/Cargo.toml index d5cb9632a..2088b1f03 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -106,7 +106,7 @@ uname = "0.1.1" crc32fast = "1.2.0" adler32 = "1.0.3" gethostname = "0.2.0" -socket2 = "0.3" +socket2 = "0.3.19" rustyline = "6.0" openssl = { version = "0.10", features = ["vendored"], optional = true } openssl-sys = { version = "0.9", optional = true } diff --git a/vm/src/py_io.rs b/vm/src/py_io.rs index 17c007772..3974dac5b 100644 --- a/vm/src/py_io.rs +++ b/vm/src/py_io.rs @@ -30,6 +30,25 @@ impl Write for PyWriter<'_> { } } +pub fn write_all( + mut buf: &[u8], + mut write: impl FnMut(&[u8]) -> io::Result, +) -> io::Result<()> { + while !buf.is_empty() { + match write(buf) { + Ok(0) => { + return Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write whole buffer", + )) + } + Ok(n) => buf = &buf[n..], + Err(e) => return Err(e), + } + } + Ok(()) +} + pub fn file_readline(obj: &PyObjectRef, size: Option, vm: &VirtualMachine) -> PyResult { let args = size.map_or_else(Vec::new, |size| vec![vm.ctx.new_int(size)]); let ret = vm.call_method(obj, "readline", args)?; diff --git a/vm/src/stdlib/socket.rs b/vm/src/stdlib/socket.rs index ca7a6962c..23e876afe 100644 --- a/vm/src/stdlib/socket.rs +++ b/vm/src/stdlib/socket.rs @@ -4,7 +4,7 @@ use gethostname::gethostname; use nix::unistd::sethostname; use socket2::{Domain, Protocol, Socket, Type as SocketType}; use std::convert::TryFrom; -use std::io::{self, prelude::*}; +use std::io; use std::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr, ToSocketAddrs}; use std::time::Duration; @@ -20,7 +20,7 @@ use crate::pyobject::{ BorrowValue, Either, IntoPyObject, PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, StaticType, TryFromObject, }; -use crate::vm::VirtualMachine; +use crate::{py_io, VirtualMachine}; #[cfg(unix)] type RawSocket = std::os::unix::io::RawFd; @@ -48,7 +48,8 @@ mod c { 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_STREAM, SOL_SOCKET, SO_BROADCAST, SO_REUSEADDR, SO_TYPE, *, + SOCK_RDM, SOCK_STREAM, SOL_SOCKET, SO_BROADCAST, SO_ERROR, SO_OOBINLINE, SO_REUSEADDR, + SO_TYPE, *, }; } @@ -71,7 +72,7 @@ pub type PySocketRef = PyRef; #[pyimpl(flags(BASETYPE))] impl PySocket { - fn sock(&self) -> PyRwLockReadGuard<'_, Socket> { + pub fn sock(&self) -> PyRwLockReadGuard<'_, Socket> { self.sock.read() } @@ -167,52 +168,90 @@ impl PySocket { } #[pymethod] - fn recv(&self, bufsize: usize, vm: &VirtualMachine) -> PyResult> { + fn recv( + &self, + bufsize: usize, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult> { + let flags = flags.unwrap_or(0); let mut buffer = vec![0u8; bufsize]; let n = self .sock() - .recv(&mut buffer) + .recv_with_flags(&mut buffer, flags) .map_err(|err| convert_sock_error(vm, err))?; buffer.truncate(n); Ok(buffer) } #[pymethod] - fn recv_into(&self, buf: PyRwBytesLike, vm: &VirtualMachine) -> PyResult { - buf.with_ref(|buf| self.sock().recv(buf)) + fn recv_into( + &self, + buf: PyRwBytesLike, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + let flags = flags.unwrap_or(0); + buf.with_ref(|buf| self.sock().recv_with_flags(buf, flags)) .map_err(|err| convert_sock_error(vm, err)) } #[pymethod] - fn recvfrom(&self, bufsize: usize, vm: &VirtualMachine) -> PyResult<(Vec, AddrTuple)> { + fn recvfrom( + &self, + bufsize: usize, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult<(Vec, AddrTuple)> { + let flags = flags.unwrap_or(0); let mut buffer = vec![0u8; bufsize]; let (n, addr) = self .sock() - .recv_from(&mut buffer) + .recv_from_with_flags(&mut buffer, flags) .map_err(|err| convert_sock_error(vm, err))?; buffer.truncate(n); Ok((buffer, get_addr_tuple(addr))) } #[pymethod] - fn send(&self, bytes: PyBytesLike, vm: &VirtualMachine) -> PyResult { + fn send( + &self, + bytes: PyBytesLike, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult { + let flags = flags.unwrap_or(0); bytes - .with_ref(|b| self.sock().send(b)) + .with_ref(|b| self.sock().send_with_flags(b, flags)) .map_err(|err| convert_sock_error(vm, err)) } #[pymethod] - fn sendall(&self, bytes: PyBytesLike, vm: &VirtualMachine) -> PyResult<()> { + fn sendall( + &self, + bytes: PyBytesLike, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult<()> { + let flags = flags.unwrap_or(0); + let sock = self.sock(); bytes - .with_ref(|b| self.sock_mut().write_all(b)) + .with_ref(|buf| py_io::write_all(buf, |b| sock.send_with_flags(b, flags))) .map_err(|err| convert_sock_error(vm, err)) } #[pymethod] - fn sendto(&self, bytes: PyBytesLike, address: Address, vm: &VirtualMachine) -> PyResult<()> { + fn sendto( + &self, + bytes: PyBytesLike, + address: Address, + flags: OptionalArg, + vm: &VirtualMachine, + ) -> PyResult<()> { + let flags = flags.unwrap_or(0); let addr = get_addr(vm, address, Some(self.family.load()))?; bytes - .with_ref(|b| self.sock().send_to(b, &addr)) + .with_ref(|b| self.sock().send_to_with_flags(b, &addr, flags)) .map_err(|err| convert_sock_error(vm, err))?; Ok(()) } @@ -411,15 +450,15 @@ impl PySocket { impl io::Read for PySocketRef { fn read(&mut self, buf: &mut [u8]) -> io::Result { - ::read(&mut self.sock_mut(), buf) + <&Socket as io::Read>::read(&mut &*self.sock(), buf) } } impl io::Write for PySocketRef { fn write(&mut self, buf: &[u8]) -> io::Result { - ::write(&mut self.sock_mut(), buf) + <&Socket as io::Write>::write(&mut &*self.sock(), buf) } fn flush(&mut self) -> io::Result<()> { - ::flush(&mut self.sock_mut()) + <&Socket as io::Write>::flush(&mut &*self.sock()) } } @@ -850,7 +889,8 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { "SO_REUSEADDR" => ctx.new_int(c::SO_REUSEADDR), "SO_TYPE" => ctx.new_int(c::SO_TYPE), "SO_BROADCAST" => ctx.new_int(c::SO_BROADCAST), - // "SO_EXCLUSIVEADDRUSE" => ctx.new_int(c::SO_EXCLUSIVEADDRUSE), + "SO_OOBINLINE" => ctx.new_int(c::SO_OOBINLINE), + "SO_ERROR" => ctx.new_int(c::SO_ERROR), "TCP_NODELAY" => ctx.new_int(c::TCP_NODELAY), "AI_ALL" => ctx.new_int(c::AI_ALL), "AI_PASSIVE" => ctx.new_int(c::AI_PASSIVE),