diff --git a/fickling/fickle.py b/fickling/fickle.py index c71b7f1..bfb7d2a 100644 --- a/fickling/fickle.py +++ b/fickling/fickle.py @@ -2066,6 +2066,24 @@ def run(self, interpreter: Interpreter): interpreter.stack.append(make_constant(False)) +class NextBuffer(Opcode): + name = "NEXT_BUFFER" + + def run(self, interpreter: Interpreter): + # The buffer contents are out-of-band: they are not part of the pickle + # stream, so the analysis models them as an opaque value. + interpreter.stack.append(ast.Name("out_of_band_buffer", ast.Load())) + + +class ReadOnlyBuffer(Opcode): + name = "READONLY_BUFFER" + + def run(self, interpreter: Interpreter): + # Marks the buffer on the top of the stack read-only; the value itself + # is unchanged, so this is a no-op for the AST. + pass + + class Tuple(StackSliceOpcode): name = "TUPLE" @@ -2386,6 +2404,22 @@ class BinBytes8(BinBytes): length_bytes = 8 +class ByteArray8(BinBytes8): + name = "BYTEARRAY8" + priority = BinBytes8.priority + 1 + + def encode_body(self) -> bytes: + return bytes(self.arg) + + @classmethod + def validate(cls, obj): + if not isinstance(obj, bytearray): + raise ValueError( + f"{cls.__name__} must be instantiated with an object of type bytearray, not {obj!r}" + ) + return super().validate(obj) + + class Long1(ConstantInt): name = "LONG1" num_bytes = 1 @@ -2418,6 +2452,20 @@ class Long(Int): priority = Int.priority + 1 +class Float(ConstantOpcode): + name = "FLOAT" + priority = Long.priority + 1 + + def encode_body(self) -> bytes: + return f"{self.arg}\n".encode() + + @classmethod + def validate(cls, obj): + if not isinstance(obj, float): + raise ValueError(f"{cls.__name__} expects a float, but received {obj!r}") + return obj + + class Dict(Opcode): name = "DICT" diff --git a/test/test_pickle.py b/test/test_pickle.py index 9318244..1e429aa 100644 --- a/test/test_pickle.py +++ b/test/test_pickle.py @@ -1,3 +1,5 @@ +import ast +import io import pickle from ast import unparse from contextlib import redirect_stdout @@ -491,3 +493,37 @@ def test_ext_with_function(self): self.assertEqual(result, real_result) finally: copyreg.remove_extension("os.path", "join", 200) + + +class TestUnimplementedOpcodeRegression(TestCase): + """Regression tests: these four standard opcodes previously raised + NotImplementedError at load time, crashing the security analysis on benign + pickles (a float in a legacy protocol-0 file; a bytearray or out-of-band + buffer in a protocol-5 file).""" + + def test_float_proto0(self): + pickled = Pickled.load(io.BytesIO(dumps(1.5, protocol=0))) + self.assertEqual(1.5, get_result(pickled)) + + def test_bytearray8_proto5(self): + pickled = Pickled.load(io.BytesIO(dumps(bytearray(b"hi"), protocol=5))) + constants = [n.value for n in ast.walk(pickled.ast) if isinstance(n, ast.Constant)] + self.assertIn(bytearray(b"hi"), constants) + + def test_next_buffer_and_readonly_buffer_do_not_crash(self): + # PROTO 5, NEXT_BUFFER, READONLY_BUFFER, STOP: an out-of-band buffer's + # opcode sequence (the buffer contents live outside the stream). + stream = b"\x80\x05\x97\x98." + pickled = Pickled.load(io.BytesIO(stream)) + names = [n.id for n in ast.walk(pickled.ast) if isinstance(n, ast.Name)] + self.assertIn("out_of_band_buffer", names) + check_safety(pickled) + + def test_load_never_raises_notimplementederror(self): + for stream in ( + dumps(1.5, protocol=0), + dumps(bytearray(b"hi"), protocol=5), + b"\x80\x05\x97\x98.", + ): + with self.subTest(stream=stream): + Pickled.load(io.BytesIO(stream))