From 615d67252676c63939c3714826dcee48fecb6951 Mon Sep 17 00:00:00 2001 From: dot Date: Wed, 4 Nov 2026 07:59:24 +0900 Subject: [PATCH] Align socket receive and control argument binding Reuse keyword parsing for receive buffers, preserve conversion ordering and argument lifetimes, and avoid zero-length receive waits. Assisted-by: OpenAI Codex:model version not exposed --- crates/stdlib/src/socket.rs | 195 +++++++++++++++++++----------- crates/vm/src/function/getargs.rs | 12 +- crates/vm/src/function/mod.rs | 2 +- 3 files changed, 131 insertions(+), 78 deletions(-) diff --git a/crates/stdlib/src/socket.rs b/crates/stdlib/src/socket.rs index 1b79f6ded52..a567857304d 100644 --- a/crates/stdlib/src/socket.rs +++ b/crates/stdlib/src/socket.rs @@ -28,7 +28,7 @@ mod _socket { }, convert::{IntoPyException, ToPyObject, TryFromBorrowedObject, TryFromObject}, function::{ - ArgBytesLike, ArgIndex, ArgIntoFloat, ArgMemoryBuffer, ArgStrOrBytesLike, FsPath, + ArgBytesLike, ArgIntoFloat, ArgMemoryBuffer, ArgSpec, ArgStrOrBytesLike, FsPath, FuncArgs, OptionalArg, }, types::{Constructor, DefaultConstructor, Destructor, Initializer, Representable}, @@ -1895,6 +1895,61 @@ mod _socket { addr: Option, } + struct ReceiveIntoArgs { + buffer: ArgMemoryBuffer, + nbytes: usize, + flags: i32, + } + + impl ReceiveIntoArgs { + fn parse( + args: &FuncArgs, + name: &'static str, + too_small: &'static str, + vm: &VirtualMachine, + ) -> PyResult { + let buffer = core::cell::RefCell::new(None); + let nbytes = core::cell::Cell::new(0isize); + let flags = core::cell::Cell::new(0i32); + ArgSpec { + fname: name, + keywords: &["buffer", "nbytes", "flags"], + required: 1, + max_positional: 3, + } + .parse_with( + args, + |position, value, vm| { + match position { + 0 => { + let export = ArgMemoryBuffer::try_from_borrowed_object(vm, value)?; + *buffer.borrow_mut() = Some(export); + } + 1 => nbytes.set(value.try_index(vm)?.try_to_primitive::(vm)?), + 2 => flags.set(value.try_index(vm)?.try_to_primitive::(vm)?), + _ => unreachable!(), + } + Ok(()) + }, + vm, + )?; + let buffer = buffer.into_inner().expect("required buffer was parsed"); + let nbytes = usize::try_from(nbytes.get()) + .map_err(|_| vm.new_value_error(format!("negative buffersize in {name}")))?; + let nbytes = if nbytes == 0 { buffer.len() } else { nbytes }; + if nbytes > buffer.len() { + return Err(vm.new_value_error(too_small)); + } + // Keep only the export alive during conversions and I/O, not a + // payload lock: index callbacks may write to this same buffer. + Ok(Self { + buffer, + nbytes, + flags: flags.get(), + }) + } + } + impl Initializer for PySocket { type Args = SocketInitArgs; @@ -2140,6 +2195,9 @@ mod _socket { let bufsize = bufsize .to_usize() .ok_or_else(|| vm.new_value_error("negative buffersize in recv"))?; + if bufsize == 0 { + return Ok(Vec::new()); + } let mut buffer = alloc_recv_scratch(bufsize, flags, vm)?; let n = zelf.sock_op(vm, SockWaitKind::Read, || { zelf.sock_snapshot()? @@ -2154,27 +2212,23 @@ mod _socket { #[pymethod] fn recv_into( zelf: &Py, - buf: ArgMemoryBuffer, - nbytes: OptionalArg, - flags: OptionalArg, + args: FuncArgs, vm: &VirtualMachine, ) -> Result { - let flags = flags.unwrap_or(0); - - // Handle nbytes parameter - let read_len = if let OptionalArg::Present(nbytes) = nbytes { - let nbytes = nbytes - .to_usize() - .ok_or_else(|| vm.new_value_error("negative buffersize in recv_into"))?; - if nbytes > buf.len() { - return Err(vm - .new_value_error("buffer too small for requested bytes") - .into()); - } - if nbytes == 0 { buf.len() } else { nbytes } - } else { - buf.len() - }; + let ReceiveIntoArgs { + buffer: buf, + nbytes: read_len, + flags, + } = ReceiveIntoArgs::parse( + &args, + "recv_into", + "buffer too small for requested bytes", + vm, + )?; + // Like recv(), an empty receive does not enter the socket wait. + if read_len == 0 { + return Ok(0); + } let mut scratch = alloc_recv_scratch(read_len, flags, vm)?; let n = zelf.sock_op(vm, SockWaitKind::Read, || { @@ -2216,26 +2270,19 @@ mod _socket { #[pymethod] fn recvfrom_into( zelf: &Py, - buf: ArgMemoryBuffer, - nbytes: OptionalArg, - flags: OptionalArg, + args: FuncArgs, vm: &VirtualMachine, ) -> Result<(usize, PyObjectRef), IoOrPyException> { - let read_len = match nbytes { - OptionalArg::Present(i) => { - let i = i.to_usize().ok_or_else(|| { - vm.new_value_error("negative buffersize in recvfrom_into") - })?; - if i > buf.len() { - return Err(vm - .new_value_error("nbytes is greater than the length of the buffer") - .into()); - } - if i == 0 { buf.len() } else { i } - } - OptionalArg::Missing => buf.len(), - }; - let flags = flags.unwrap_or(0); + let ReceiveIntoArgs { + buffer: buf, + nbytes: read_len, + flags, + } = ReceiveIntoArgs::parse( + &args, + "recvfrom_into", + "nbytes is greater than the length of the buffer", + vm, + )?; let mut scratch = alloc_recv_scratch(read_len, flags, vm)?; let (n, addr) = zelf.sock_op(vm, SockWaitKind::Read, || { zelf.sock_snapshot()? @@ -2300,20 +2347,35 @@ mod _socket { .new_type_error("sendto() takes no keyword arguments") .into()); } - let n = args.args.len(); - let (bytes, flags, address) = match n { - 2 => (args.args[0].clone(), 0, args.args[1].clone()), - 3 => { - let flags = args.args[1].try_index(vm)?.try_to_primitive::(vm)?; - (args.args[0].clone(), flags, args.args[2].clone()) - } - _ => { - return Err(vm - .new_type_error(format!("sendto() takes 2 or 3 arguments ({n} given)")) - .into()); - } + // This overload's diagnostic cannot be expressed by the binder's + // ordinary minimum/maximum arity messages. + let nargs = args.args.len(); + if !(2..=3).contains(&nargs) { + return Err(vm + .new_type_error(format!("sendto() takes 2 or 3 arguments ({nargs} given)")) + .into()); + } + #[expect( + clippy::redundant_clone, + reason = "Keep original arguments alive through callbacks and the socket operation" + )] + let (data, flags_or_address, address): ( + PyObjectRef, + PyObjectRef, + OptionalArg, + ) = args.clone().bind_for(vm, "sendto")?; + // Export before converting flags, but defer borrowing its data + // until all Python callbacks, including address conversion, finish. + let bytes = ArgBytesLike::try_from_object(vm, data)?; + let (flags, address) = match address { + OptionalArg::Missing => (0, flags_or_address), + OptionalArg::Present(address) => ( + flags_or_address + .try_index(vm)? + .try_to_primitive::(vm)?, + address, + ), }; - let bytes = ArgBytesLike::try_from_object(vm, bytes)?; let addr = zelf.extract_address(address, "sendto", vm)?; let buf = bytes.borrow_buf_unlocked(vm)?; let buf = &*buf; @@ -2739,30 +2801,21 @@ mod _socket { .new_type_error("setsockopt() takes no keyword arguments") .into()); } - let nargs = args.args.len(); - if nargs < 3 { - return Err(vm - .new_type_error(format!( - "setsockopt() takes at least 3 arguments ({nargs} given)" - )) - .into()); - } - if nargs > 4 { - return Err(vm - .new_type_error(format!( - "setsockopt() takes at most 4 arguments ({nargs} given)" - )) - .into()); - } + #[expect( + clippy::redundant_clone, + reason = "Keep original arguments alive through callbacks and the socket operation" + )] let (level, name, value, optlen): ( - i32, - i32, + PyObjectRef, + PyObjectRef, Option, - OptionalArg, - ) = args.bind(vm)?; + OptionalArg, + ) = args.clone().bind_for(vm, "setsockopt")?; + let level = level.try_index(vm)?.try_to_primitive::(vm)?; + let name = name.try_index(vm)?.try_to_primitive::(vm)?; let optlen = match optlen { OptionalArg::Present(index) => { - let index = index.as_ref(); + let index = index.try_index(vm)?; // CPython's unsigned-int bitfield conversion masks the value, // warning only when it fits neither a signed nor unsigned int. if !index.as_bigint().to_i64().is_some_and(|value| { diff --git a/crates/vm/src/function/getargs.rs b/crates/vm/src/function/getargs.rs index 371b99137b6..7e9193fe7c4 100644 --- a/crates/vm/src/function/getargs.rs +++ b/crates/vm/src/function/getargs.rs @@ -6,17 +6,17 @@ use crate::{PyObject, PyObjectRef, PyResult, VirtualMachine, common::wtf8::Wtf8} /// The keyword list of a `PyArg_ParseTupleAndKeywords` call together with the /// `|`, `$` and `:` markers of its format string. Names are matched to the /// supplied arguments; converting them is left to the caller. -pub(crate) struct ArgSpec<'a> { +pub struct ArgSpec<'a> { /// What the `:name` suffix names the function in error messages. - pub(crate) fname: &'a str, + pub fname: &'a str, /// Every parameter name, in `kwlist` order. None may be positional-only. - pub(crate) keywords: &'a [&'a str], + pub keywords: &'a [&'a str], /// Where `|` sits: how many leading parameters are required. Equal to the /// parameter count when the format string has no `|`. - pub(crate) required: usize, + pub required: usize, /// Where `$` sits: how many leading parameters may be given positionally. /// Equal to the parameter count when the format string has no `$`. - pub(crate) max_positional: usize, + pub max_positional: usize, } impl ArgSpec<'_> { @@ -30,7 +30,7 @@ impl ArgSpec<'_> { /// `check` runs for each supplied argument in `kwlist` order, so a /// conversion failure is reported ahead of the checks that follow it. - pub(crate) fn parse_with( + pub fn parse_with( &self, args: &FuncArgs, check: impl Fn(usize, &PyObject, &VirtualMachine) -> PyResult<()>, diff --git a/crates/vm/src/function/mod.rs b/crates/vm/src/function/mod.rs index b434eb3bccc..24f7dbd7a95 100644 --- a/crates/vm/src/function/mod.rs +++ b/crates/vm/src/function/mod.rs @@ -30,7 +30,7 @@ pub use builtin::{IntoPyNativeFn, PyNativeFn, static_func, static_raw_func}; pub use doctext::{ItemDoc, db_doc, plain_doc}; pub use either::Either; pub use fspath::FsPath; -pub(crate) use getargs::ArgSpec; +pub use getargs::ArgSpec; pub use getset::PySetterValue; pub(super) use getset::{IntoPyGetterFunc, IntoPySetterFunc, PyGetterFunc, PySetterFunc}; pub use method::{HeapMethodDef, PyMethodDef, PyMethodFlags};