diff --git a/crates/stdlib/src/pickle.rs b/crates/stdlib/src/pickle.rs index 031c700fdf9..51427c0c9b0 100644 --- a/crates/stdlib/src/pickle.rs +++ b/crates/stdlib/src/pickle.rs @@ -818,13 +818,18 @@ mod _pickle { .ok() .filter(|n| isize::try_from(*n).is_ok()) .ok_or_else(|| { - new_unpickling_error( - vm, + let message = if what == "FRAME" { + format!( + "FRAME length exceeds system's maximum of {} bytes", + isize::MAX + ) + } else { format!( "{what} exceeds system's maximum size of {} bytes", isize::MAX - ), - ) + ) + }; + vm.new_overflow_error(message) }) } @@ -1090,7 +1095,7 @@ mod _pickle { } Some(BINBYTES8) => { let arr: [u8; 8] = st.read_n(8, vm)?.try_into().unwrap(); - let n = read_size(&arr, vm, "BINBYTES8")?; + let n = read_size(&arr, vm, "BINBYTES")?; let data = st.read_n(n, vm)?.to_vec(); stack.push(vm.ctx.new_bytes(data).into()); } @@ -1129,7 +1134,7 @@ mod _pickle { } Some(BINUNICODE8) => { let arr: [u8; 8] = st.read_n(8, vm)?.try_into().unwrap(); - let n = read_size(&arr, vm, "BINUNICODE8")?; + let n = read_size(&arr, vm, "BINUNICODE")?; let data = st.read_n(n, vm)?.to_vec(); stack.push(utf8_surrogatepass(&data, vm)?); } diff --git a/extra_tests/snippets/stdlib_pickle.py b/extra_tests/snippets/stdlib_pickle.py index 855bd79691e..e1872a24bfa 100644 --- a/extra_tests/snippets/stdlib_pickle.py +++ b/extra_tests/snippets/stdlib_pickle.py @@ -1,5 +1,8 @@ import _pickle import io +import sys + +from testutils import assert_raises # BINPUT leaves holes; MEMOIZE uses the number of present entries. unpickler = _pickle.Unpickler(io.BytesIO(b"\x80\x04]q\x07\x94h\x01\x86.")) @@ -13,3 +16,22 @@ # Slot wrapper descriptors reduce to lookup on their defining class. assert object.__str__.__reduce__() == (getattr, (object, "__str__")) assert _pickle.loads(_pickle.dumps(object.__str__)) is object.__str__ + +# Oversized advertised lengths are rejected before trying to read or allocate. +for opcode, name in ( + (b"\x8e", "BINBYTES"), + (b"\x8d", "BINUNICODE"), + (b"\x95", "FRAME"), +): + data = b"\x80\x04" + opcode + (sys.maxsize + 1).to_bytes(8, "little") + for load in ( + _pickle.loads, + lambda data: _pickle.Unpickler(io.BytesIO(data)).load(), + ): + with assert_raises(OverflowError) as caught: + load(data) + if name == "FRAME": + expected = f"FRAME length exceeds system's maximum of {sys.maxsize} bytes" + else: + expected = f"{name} exceeds system's maximum size of {sys.maxsize} bytes" + assert str(caught.exception) == expected