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
1 change: 0 additions & 1 deletion Lib/test/test_codeccallbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1038,7 +1038,6 @@ def __getitem__(self, key):
self.assertRaises(ValueError, codecs.charmap_encode, "\xff", err, D())
self.assertRaises(TypeError, codecs.charmap_encode, "\xff", err, {0xff: 300})

@unittest.expectedFailure # TODO: RUSTPYTHON
def test_decodehelper_bug36819(self):
handler = RepeatedPosReturn("x")
codecs.register_error("test.bug36819", handler.handle)
Expand Down
48 changes: 38 additions & 10 deletions crates/vm/src/codecs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1185,6 +1185,16 @@ impl<'a> PyDecodeContext<'a> {
exception: OnceCell::new(),
}
}

pub(crate) fn new_owned(encoding: &'a str, data: PyBytesRef, vm: &'a VirtualMachine) -> Self {
Self {
vm,
encoding,
data: PyDecodeData::Modified(data),
pos: 0,
exception: OnceCell::new(),
}
}
}

impl CodecContext for PyDecodeContext<'_> {
Expand Down Expand Up @@ -1333,6 +1343,7 @@ impl<'a> DecodeErrorHandler<PyDecodeContext<'a>> for StandardError {
pub(crate) struct ErrorsHandler<'a> {
errors: &'a Py<PyUtf8Str>,
resolved: OnceCell<ResolvedError>,
use_registry: bool,
}

enum ResolvedError {
Expand All @@ -1347,22 +1358,34 @@ impl<'a> ErrorsHandler<'a> {
Self {
errors,
resolved: OnceCell::new(),
use_registry: false,
}
} else {
Self {
errors: identifier_utf8!(vm, strict),
resolved: OnceCell::from(ResolvedError::Standard(StandardError::Strict)),
use_registry: false,
}
}
}

pub(crate) fn new_registered(errors: Option<&'a Py<PyUtf8Str>>, vm: &VirtualMachine) -> Self {
Self {
errors: errors.unwrap_or_else(|| identifier_utf8!(vm, strict)),
resolved: OnceCell::new(),
use_registry: true,
}
}

#[inline]
fn resolve(&self, vm: &VirtualMachine) -> PyResult<&ResolvedError> {
if let Some(val) = self.resolved.get() {
return Ok(val);
}
let errors_str = self.errors.as_str();
let val = if let Ok(standard) = errors_str.parse() {
let val = if !self.use_registry
&& let Ok(standard) = errors_str.parse()
{
ResolvedError::Standard(standard)
} else {
vm.state
Expand Down Expand Up @@ -1448,14 +1471,19 @@ impl<'a> DecodeErrorHandler<PyDecodeContext<'a>> for ErrorsHandler<'_> {
ResolvedError::Handler(handler) => handler,
};
let decode_exc = ctx.error_decoding(byte_range, reason);
let data_bytes = decode_exc
.as_object()
.get_attr("object", vm)?
.downcast::<PyBytes>()
.map_err(|_| vm.new_type_error("object attribute must be bytes"))?;
// CPython resumes from the exception's bytes object, not the original
// mutable buffer. Switch before calling Python to release its data lock.
ctx.data = PyDecodeData::Modified(data_bytes);
if matches!(ctx.data, PyDecodeData::Original(_)) {
// Release the original buffer's read guard before calling Python.
// A retained exception may have been changed by a mapping callback;
// only validate its object after the error handler can repair it.
let data_bytes = decode_exc
.as_object()
.get_attr("object", vm)?
.downcast::<PyBytes>()
.map_err(|_| {
vm.new_type_error("UnicodeError 'object' attribute must be a bytes")
})?;
ctx.data = PyDecodeData::Modified(data_bytes);
}
let res = handler.call((decode_exc.clone(),), vm)?;
let tuple_err = || vm.new_type_error("decoding error handler must return (str, int) tuple");
let (replace, restart) = match res.downcast_ref::<PyTuple>().map(|tup| tup.as_slice()) {
Expand All @@ -1477,7 +1505,7 @@ impl<'a> DecodeErrorHandler<PyDecodeContext<'a>> for ErrorsHandler<'_> {
.as_object()
.get_attr("object", vm)?
.downcast::<PyBytes>()
.map_err(|_| vm.new_type_error("object attribute must be bytes"))?;
.map_err(|_| vm.new_type_error("UnicodeError 'object' attribute must be a bytes"))?;
ctx.data = PyDecodeData::Modified(new_data);
let restart = if restart < 0 {
ctx.data.len() as isize + restart
Expand Down
115 changes: 102 additions & 13 deletions crates/vm/src/stdlib/_codecs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,17 @@ mod _codecs {
use core::hint::cold_path;

use crate::codecs::{ErrorsHandler, PyDecodeContext, PyEncodeContext};
use crate::common::encodings;
use crate::common::wtf8::Wtf8Buf;
use crate::common::encodings::{self, DecodeContext};
use crate::common::wtf8::{CodePoint, Wtf8Buf};
use crate::{
AsObject, PyObjectRef, PyResult, VirtualMachine,
builtins::{PyBytesRef, PyStrRef, PyUtf8StrRef},
builtins::{PyBytesRef, PyInt, PyStr, PyStrRef, PyUtf8StrRef},
codecs,
convert::TryFromObject,
exceptions::nul_char_error,
function::ArgBytesLike,
function::{ArgBytesLike, OptionalOption},
};
use num_traits::ToPrimitive;

#[derive(FromArgs)]
struct NormalizeEncodingArgs {
Expand Down Expand Up @@ -671,9 +672,9 @@ mod _codecs {
#[derive(FromArgs)]
struct CharmapDecodeArgs {
#[pyarg(positional)]
data: PyObjectRef,
data: ArgBytesLike,
#[pyarg(positional, optional)]
errors: Option<PyObjectRef>,
errors: OptionalOption<PyUtf8StrRef>,
#[pyarg(positional, optional)]
mapping: Option<PyObjectRef>,
}
Expand Down Expand Up @@ -702,15 +703,103 @@ mod _codecs {
delegate_pycodecs!(charmap_encode, forwarded, vm)
}
#[pyfunction]
fn charmap_decode(args: CharmapDecodeArgs, vm: &VirtualMachine) -> PyResult {
let mut forwarded = vec![args.data];
if args.mapping.is_some() && args.errors.is_none() {
forwarded.push(vm.ctx.new_str("strict").into());
fn charmap_decode(args: CharmapDecodeArgs, vm: &VirtualMachine) -> DecodeResult {
let errors = args.errors.flatten();
if errors
.as_ref()
.is_some_and(|s| s.as_pystr().contains_nuls())
{
return Err(nul_char_error(vm));
}
// Charmap honors registry overrides even for standard error names.
let errors = ErrorsHandler::new_registered(errors.as_deref(), vm);
let Some(mapping) = args.mapping.filter(|mapping| !vm.is_none(mapping)) else {
let ctx = PyDecodeContext::new(encodings::latin_1::ENCODING_NAME, &args.data, vm);
return encodings::latin_1::decode(ctx, &errors);
};
let consumed = args.data.borrow_buf().len();
let mut ctx: Option<PyDecodeContext<'_>> = None;
let mut pos = 0;
let mut decoded = Wtf8Buf::with_capacity(consumed);
// Only exact strings bypass __getitem__; subclasses can override it.
let string_map = mapping
.downcast_ref_if_exact::<PyStr>(vm)
.map(|s| s.as_wtf8().code_points().take(256).collect::<Vec<_>>());
loop {
// Do not hold a buffer read guard across mapping callbacks: they
// may change the original bytearray without resizing it.
let byte = match &ctx {
Some(ctx) => ctx.full_data().get(pos).copied(),
None => args.data.borrow_buf().get(pos).copied(),
};
let Some(byte) = byte else { break };
let defined = if let Some(string_map) = &string_map {
if let Some(&cp) = string_map
.get(usize::from(byte))
.filter(|cp| cp.to_u32() != 0xfffe)
{
decoded.push(cp);
true
} else {
false
}
} else {
charmap_decode_mapping(&mapping, byte, &mut decoded, vm)?
};
if defined {
pos += 1;
} else {
// At the first error CPython switches to the exception's
// immutable bytes, including mutations made by this lookup.
let ctx = ctx.get_or_insert_with(|| {
let data = vm.ctx.new_bytes(args.data.borrow_buf().to_vec());
PyDecodeContext::new_owned("charmap", data, vm)
});
ctx.restart_from(pos)?;
let replacement =
ctx.handle_error(&errors, pos..pos + 1, Some("character maps to <undefined>"))?;
decoded.push_wtf8(replacement.as_wtf8());
pos = ctx.position();
}
}
// The callback can replace its exception's input, but the codec reports
// the size of the original buffer, as PyUnicode_DecodeCharmap does.
Ok((decoded, consumed))
}

fn charmap_decode_mapping(
mapping: &PyObjectRef,
byte: u8,
decoded: &mut Wtf8Buf,
vm: &VirtualMachine,
) -> PyResult<bool> {
let value = match mapping.get_item(vm.ctx.new_int(byte).as_object(), vm) {
Ok(value) => value,
Err(e) if e.fast_isinstance(vm.ctx.exceptions.lookup_error) => return Ok(false),
Err(e) => return Err(e),
};
if vm.is_none(&value) {
return Ok(false);
}
if let Some(integer) = value.downcast_ref::<PyInt>() {
let cp = integer
.as_bigint()
.to_u32()
.and_then(CodePoint::from_u32)
.ok_or_else(|| vm.new_type_error("character mapping must be in range(0x110000)"))?;
if cp.to_u32() == 0xfffe {
return Ok(false);
}
decoded.push(cp);
} else if let Some(text) = value.downcast_ref::<PyStr>() {
if text.as_wtf8() == "\u{fffe}" {
return Ok(false);
}
decoded.push_wtf8(text.as_wtf8());
} else {
push_optional(&mut forwarded, args.errors);
return Err(vm.new_type_error("character mapping must return integer, None or str"));
}
push_optional(&mut forwarded, args.mapping);
delegate_pycodecs!(charmap_decode, forwarded, vm)
Ok(true)
}
#[pyfunction]
fn charmap_build(args: CharmapBuildArgs, vm: &VirtualMachine) -> PyResult {
Expand Down
49 changes: 49 additions & 0 deletions extra_tests/snippets/stdlib_codecs_wide.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,3 +204,52 @@ def reentrant_search(name):


test_registry_reentrant_cleanup()


def test_charmap_invalid_callback_object_diagnostic():
def handler(exc):
exc.object = None
return ("", 1)

name = "test_charmap_invalid_callback_object"
codecs.register_error(name, handler)
try:
try:
codecs.charmap_decode(b"!", name, {})
except TypeError as exc:
assert str(exc) == "UnicodeError 'object' attribute must be a bytes"
else:
raise AssertionError("invalid callback input was accepted")
finally:
_codecs._unregister_error(name)


test_charmap_invalid_callback_object_diagnostic()


def test_charmap_handler_repairs_retained_exception():
state = {}

class Mapping:
def __getitem__(self, key):
if key == ord("a"):
state["exception"].object = None
return "A"
return None

def handler(exc):
if "exception" not in state:
state["exception"] = exc
return ("X", 1)
exc.object = b"!a?"
return ("Y", 3)

name = "test_charmap_repair_retained_exception"
codecs.register_error(name, handler)
try:
assert codecs.charmap_decode(b"!a?", name, Mapping()) == ("XAY", 3)
finally:
_codecs._unregister_error(name)


test_charmap_handler_repairs_retained_exception()
Loading