Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions Lib/test/test_asyncio/test_sock_lowlevel.py
Original file line number Diff line number Diff line change
Expand Up @@ -663,10 +663,6 @@ async def recvfrom_into(socket):
self._basetest_datagram_send_to_non_listening_address(
recvfrom_into))

@unittest.expectedFailureIfWindows("TODO: RUSTPYTHON; AssertionError: ValueError not raised")
def test_recvfrom_into_wrong_size(self):
return super().test_recvfrom_into_wrong_size()

else:
import selectors

Expand Down
1 change: 0 additions & 1 deletion Lib/test/test_os/test_os.py
Original file line number Diff line number Diff line change
Expand Up @@ -3520,7 +3520,6 @@ def test_spawnve_noargs(self):
self.assertRaises(ValueError, os.spawnve, os.P_NOWAIT, program, ('',), {})
self.assertRaises(ValueError, os.spawnve, os.P_NOWAIT, program, [''], {})

@unittest.expectedFailureIfWindows("TODO: RUSTPYTHON; AssertionError: \"must contain only strings\" does not match \"expected str, bytes or os.PathLike object, not int\"")
@requires_native_spawnv
def test_spawnv_arg_conversion_errors(self):
# A non-path argv item gets a TypeError naming the argument...
Expand Down
20 changes: 13 additions & 7 deletions crates/common/src/str.rs
Original file line number Diff line number Diff line change
Expand Up @@ -743,24 +743,31 @@ pub mod levenshtein {
/// Replace all tabs in a string with spaces, using the given tab size.
#[must_use]
pub fn expandtabs(input: &Wtf8, tab_size: usize) -> Wtf8Buf {
let capacity = if tab_size == 0 { 0 } else { input.len() };
let mut expanded_str = Wtf8Buf::with_capacity(capacity);
expandtabs_into(input, tab_size, &mut expanded_str);
expanded_str
}

/// Append tab-expanded input to a buffer whose capacity the caller may reserve.
pub fn expandtabs_into(input: &Wtf8, tab_stop: usize, expanded_str: &mut Wtf8Buf) {
// A tab size of zero, which is also where a negative one lands, leaves no
// column for a tab to advance to: the tabs come out and nothing else moves.
// Going through the arithmetic anyway subtracts the current column from a
// tab stop of zero and underflows on the first tab, so the width asked for
// next is `usize::MAX`. The bytes version of this already returns here.
if tab_size == 0 {
return input.code_points().filter(|ch| *ch != '\t').collect();
// next is `usize::MAX`.
if tab_stop == 0 {
expanded_str.extend(input.code_points().filter(|ch| *ch != '\t'));
return;
}

let tab_stop = tab_size;
let mut expanded_str = Wtf8Buf::with_capacity(input.len());
let mut tab_size = tab_stop;
let mut col_count = 0usize;
for ch in input.code_points() {
if ch == '\t' {
let num_spaces = tab_size - col_count;
col_count += num_spaces;
expanded_str.push_str(&" ".repeat(num_spaces));
expanded_str.extend(core::iter::repeat_n(CodePoint::from(' '), num_spaces));
} else {
expanded_str.push(ch);
if ch == '\r' || ch == '\n' {
Expand All @@ -774,7 +781,6 @@ pub fn expandtabs(input: &Wtf8, tab_size: usize) -> Wtf8Buf {
tab_size += tab_stop;
}
}
expanded_str
}

/// Creates an [`AsciiStr`][ascii::AsciiStr] from a string literal, throwing a compile error if the
Expand Down
40 changes: 40 additions & 0 deletions crates/stdlib/src/overlapped.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1048,6 +1048,10 @@ mod _overlapped {
return Err(vm.new_value_error("buffer too large"));
}

if size as usize > buf_len {
return Err(vm.new_value_error("nbytes is greater than the length of the buffer"));
}

let address: host_overlapped::SocketAddrV6 = unsafe { core::mem::zeroed() };
let address_length = core::mem::size_of::<host_overlapped::SocketAddrV6>() as i32;

Expand Down Expand Up @@ -1293,4 +1297,40 @@ mod _overlapped {
host_winapi::reset_event(handle as host_winapi::Handle)
.map_err(|err| set_from_windows_err(err.raw_os_error().unwrap_or(0) as u32, vm))
}

#[cfg(test)]
mod tests {
use super::*;
use crate::vm::{TryFromObject, builtins::PyByteArray};

#[test]
fn recvfrom_into_rejects_oversized_request_before_starting_io() {
let builder = crate::vm::Interpreter::builder(Default::default());
let defs = crate::stdlib_module_defs(&builder.ctx);
let interp = builder.add_native_modules(&defs).build();
interp.enter(|vm| {
let module = vm.import("_overlapped", 0).unwrap();
let overlapped = module
.get_attr("Overlapped", vm)
.unwrap()
.call((0,), vm)
.unwrap()
.downcast::<Overlapped>()
.unwrap();
let buffer = PyByteArray::from(vec![0; 8]).into_ref(&vm.ctx);
let buffer = ArgMemoryBuffer::try_from_object(vm, buffer.into()).unwrap();
let err = Overlapped::WSARecvFromInto(
&overlapped,
0,
buffer,
9,
OptionalArg::Missing,
vm,
)
.unwrap_err();
assert!(err.fast_isinstance(vm.ctx.exceptions.value_error));
assert!(matches!(overlapped.inner.lock().data, OverlappedData::None));
});
}
}
}
71 changes: 71 additions & 0 deletions crates/vm/src/anystr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,45 @@ use crate::{
function::{ArgIndex, PySsize},
};

pub(crate) fn checked_replace_len(
original: usize,
old: usize,
new: usize,
count: usize,
) -> Option<usize> {
let result = if new >= old {
original.checked_add(count.checked_mul(new - old)?)?
} else {
original.checked_sub(count.checked_mul(old - new)?)?
};
isize::try_from(result).ok().map(|_| result)
}

pub(crate) fn checked_expandtabs_len(
original: usize,
chars: impl Iterator<Item = u32>,
tabsize: usize,
) -> Option<usize> {
let mut result = original;
let mut column = 0;
for ch in chars {
if ch == u32::from(b'\t') {
let spaces = if tabsize == 0 {
0
} else {
tabsize - column % tabsize
};
result = checked_replace_len(result, 1, spaces, 1)?;
column += spaces;
} else if ch == u32::from(b'\r') || ch == u32::from(b'\n') {
column = 0;
} else {
column += 1;
}
}
Some(result)
}

#[derive(FromArgs)]
pub struct SplitArgs<T: TryFromObject> {
#[pyarg(any, optional)]
Expand Down Expand Up @@ -515,3 +554,35 @@ where
})?;
(predicate)(single)
}

#[cfg(test)]
mod tests {
use super::{checked_expandtabs_len, checked_replace_len};

#[test]
fn replacement_length_checks_signed_and_unsigned_overflow() {
let max = isize::MAX as usize;
assert_eq!(checked_replace_len(4, 0, 1, 5), Some(9));
assert_eq!(checked_replace_len(6, 1, 0, 4), Some(2));
assert_eq!(checked_replace_len(max, 0, 1, 0), Some(max));
assert_eq!(checked_replace_len(max, 0, 1, 1), None);
assert_eq!(checked_replace_len(1, 0, 2, usize::MAX), None);
}

#[test]
fn expanded_tabs_check_the_complete_output_before_allocating() {
let max = isize::MAX as usize;
assert_eq!(
checked_expandtabs_len(1, "\t".chars().map(u32::from), max),
Some(max)
);
assert_eq!(
checked_expandtabs_len(3, "\t\n\t".chars().map(u32::from), max),
None
);
assert_eq!(
checked_expandtabs_len(3, "a\tb".chars().map(u32::from), 0),
Some(2)
);
}
}
8 changes: 6 additions & 2 deletions crates/vm/src/builtins/bytearray.rs
Original file line number Diff line number Diff line change
Expand Up @@ -999,8 +999,12 @@ impl Py<PyByteArray> {
}

#[pymethod]
fn expandtabs(&self, options: anystr::ExpandTabsArgs) -> PyByteArray {
self.inner().expandtabs(options).into()
fn expandtabs(
&self,
options: anystr::ExpandTabsArgs,
vm: &VirtualMachine,
) -> PyResult<PyByteArray> {
self.inner().expandtabs(options, vm).map(Into::into)
}

#[pymethod]
Expand Down
8 changes: 6 additions & 2 deletions crates/vm/src/builtins/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -541,8 +541,12 @@ impl Py<PyBytes> {
}

#[pymethod]
fn expandtabs(&self, options: anystr::ExpandTabsArgs) -> PyBytes {
self.inner.expandtabs(options).into()
fn expandtabs(
&self,
options: anystr::ExpandTabsArgs,
vm: &VirtualMachine,
) -> PyResult<PyBytes> {
self.inner.expandtabs(options, vm).map(Into::into)
}

#[pymethod]
Expand Down
58 changes: 48 additions & 10 deletions crates/vm/src/builtins/str.rs
Original file line number Diff line number Diff line change
Expand Up @@ -825,17 +825,29 @@ impl PyStr {
})
}
let sep = zelf.as_wtf8();
let mut len = sep.len().saturating_mul(items.len().saturating_sub(1));
let mut len = 0usize;
let overflow = || vm.new_overflow_error("join() result is too long for a Python string");
for (i, obj) in items.iter().enumerate() {
len = len.saturating_add(item_str(i, obj, vm)?.as_wtf8().len());
len = len
.checked_add(item_str(i, obj, vm)?.as_wtf8().len())
.ok_or_else(overflow)?;
if i > 0 {
len = len.checked_add(sep.len()).ok_or_else(overflow)?;
}
if isize::try_from(len).is_err() {
return Err(overflow());
}
}
if let [only] = items {
let only = item_str(0, only, vm)?;
if only.class().is(vm.ctx.types.str_type) {
return Ok(only.to_owned());
}
}
let mut joined = Wtf8Buf::with_capacity(len);
let mut joined = Wtf8Buf::new();
joined
.try_reserve_exact(len)
.map_err(|_| vm.new_memory_error(""))?;
for (i, obj) in items.iter().enumerate() {
if i > 0 {
joined.push_wtf8(sep);
Expand Down Expand Up @@ -1250,10 +1262,14 @@ impl Py<PyStr> {
}

#[pymethod]
fn replace(zelf: PyRef<PyStr>, args: ReplaceArgs, vm: &VirtualMachine) -> PyRef<PyStr> {
fn replace(
zelf: PyRef<PyStr>,
args: ReplaceArgs,
vm: &VirtualMachine,
) -> PyResult<PyRef<PyStr>> {
let ReplaceArgs { old, new, count } = args;
if count == 0 || old.byte_len() > zelf.byte_len() || old.as_wtf8() == new.as_wtf8() {
return PyStr::result_unchanged(zelf, vm);
return Ok(PyStr::result_unchanged(zelf, vm));
}

let s = zelf.as_wtf8();
Expand All @@ -1262,15 +1278,26 @@ impl Py<PyStr> {
} else {
count as usize
};
let capacity = if new.byte_len() > old.byte_len() {
let matches = if old.is_empty() {
(zelf.char_len() + 1).min(limit)
} else {
s.find_iter(old.as_wtf8()).take(limit).count()
};
anystr::checked_replace_len(s.len(), old.byte_len(), new.byte_len(), matches)
.ok_or_else(|| vm.new_overflow_error("replace string is too long"))?
} else {
s.len()
};
let replaced = if old.is_empty() {
s.replacen(old.as_wtf8(), new.as_wtf8(), limit)
} else {
let mut matches = s.find_iter(old.as_wtf8()).take(limit);
let Some(first) = matches.next() else {
drop(matches);
return PyStr::result_unchanged(zelf, vm);
return Ok(PyStr::result_unchanged(zelf, vm));
};
let mut result = Vec::with_capacity(s.len());
let mut result = Vec::with_capacity(capacity);
let mut last = 0;
for start in core::iter::once(first).chain(matches) {
result.extend_from_slice(&s.as_bytes()[last..start]);
Expand All @@ -1281,7 +1308,7 @@ impl Py<PyStr> {
// SAFETY: A nonempty WTF-8 needle matches only at code point boundaries.
unsafe { Wtf8Buf::from_bytes_unchecked(result) }
};
vm.ctx.new_str(replaced)
Ok(vm.ctx.new_str(replaced))
}

#[pymethod]
Expand Down Expand Up @@ -1500,8 +1527,19 @@ impl Py<PyStr> {
}

#[pymethod]
fn expandtabs(&self, args: anystr::ExpandTabsArgs) -> Wtf8Buf {
rustpython_common::str::expandtabs(self.as_wtf8(), args.tabsize())
fn expandtabs(&self, args: anystr::ExpandTabsArgs, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
let capacity = anystr::checked_expandtabs_len(
self.byte_len(),
self.as_wtf8().code_points().map(|ch| ch.to_u32()),
args.tabsize(),
)
.ok_or_else(|| vm.new_overflow_error("new string is too long"))?;
let mut result = Wtf8Buf::new();
result
.try_reserve_exact(capacity)
.map_err(|_| vm.new_memory_error(""))?;
rustpython_common::str::expandtabs_into(self.as_wtf8(), args.tabsize(), &mut result);
Ok(result)
}

#[pymethod]
Expand Down
Loading
Loading