diff --git a/pytest/test_unmarshal.py b/pytest/test_unmarshal.py new file mode 100644 index 00000000..785bb813 --- /dev/null +++ b/pytest/test_unmarshal.py @@ -0,0 +1,125 @@ +"""Lossless Python 2 interned strings and context-sensitive references.""" + +import io +from struct import pack + +import pytest +from xdis.load import load_module_from_file_object +from xdis.magics import by_version, magic2int +from xdis.marsh import dumps +from xdis.unmarshal import VersionIndependentUnmarshaller, load_code + + +def string_record(tag, value): + size = bytes([len(value)]) if tag == b"Z" else pack("', 'exec'), + # rebuild using types.CodeType with intern(co.co_code), then marshal.dumps(co, 2). + # CALL_FUNCTION_VAR (0x8c) is not valid standalone UTF-8. The other variant + # uses the ordinary string tag for the same bytecode. + data = bytes.fromhex( + "6300000000000000000200000040000000740e0000006500006501008c00000164000053" + "28010000004e28020000007402000000666e74040000006172677328000000002800000000" + "280000000073090000003c666978747572653e74080000003c6d6f64756c653e01000000" + "7400000000" + ) + data = data[:17] + code_tag + data[18:] + pyc = by_version["2.7"] + pack("" + assert code.co_name == "" + + # Keep the reference metadata used by the marshal writer, for both bytecode + # and textual names. No interpreter or filesystem fixture is required. + dumped = dumps(code, python_version=(2, 7)) + assert dumped[17:18] == code_tag + assert string_record(b"t", b"fn") in dumped + restored = load_code(dumped, magic2int(by_version["2.7"]), code_objects={}) + assert restored.co_code == expected + assert restored.co_names == code.co_names diff --git a/xdis/unmarshal.py b/xdis/unmarshal.py index 036bd717..589b90b0 100644 --- a/xdis/unmarshal.py +++ b/xdis/unmarshal.py @@ -191,6 +191,7 @@ def __init__(self, fp, magic_int, bytes_for_s, code_objects={}) -> None: self.marshal_version = 0 self.intern_strings = [] + self.intern_string_bytes = {} self.intern_objects = [] self.is_pypy = magic_int in PYPY3_MAGICS self.is_rust = magic_int in RUSTPYTHON_MAGICS @@ -227,6 +228,7 @@ def load(self): if self.marshal_version == 0: self.intern_strings = [] + self.intern_string_bytes = {} if self.marshal_version < 3: assert self.intern_objects == [] @@ -428,8 +430,15 @@ def t_short_ASCII_interned(self, save_ref, bytes_for_s: bool = False): def t_interned(self, save_ref, bytes_for_s: bool = False): strsize = self.read_uint32() - interned = compat_str(self.fp.read(strsize)) + raw = self.fp.read(strsize) + interned = compat_str(raw) + if self.version_triple < (3, 0): + # A string first read as text can later be referenced as bytecode. + # Keep the original bytes alongside the existing text representation. + self.intern_string_bytes[len(self.intern_strings)] = raw self.intern_strings.append(interned) + if bytes_for_s and self.version_triple < (3, 0): + interned = raw return self.r_ref(interned, save_ref) def t_unicode(self, save_ref, bytes_for_s: bool = False): @@ -526,8 +535,17 @@ def t_dict(self, save_ref, bytes_for_s: bool = False) -> dict: def t_python2_string_reference(self, save_ref, bytes_for_s: bool = False): refnum = self.read_uint32() + if bytes_for_s and self.version_triple < (3, 0): + return self.intern_string_bytes[refnum] return self.intern_strings[refnum] + def _reference_objects(self): + references = set(self.intern_objects + self.intern_strings) + if self.version_triple < (3, 0): + # The marshal writer can encounter either representation. + references.update(self.intern_string_bytes.values()) + return references + def t_slice(self, save_ref, bytes_for_s: bool = False): """TYPE_SLICE introducted in Marshal version 5""" retval, idx = self.r_ref_reserve(slice(None, None, None), save_ref) @@ -687,7 +705,7 @@ def t_code(self, save_ref, bytes_for_s: bool = False): co_firstlineno = -1 # Bogus sentinel value co_lnotab = b"" - reference_objects = set(self.intern_objects + self.intern_strings) + reference_objects = self._reference_objects() code = to_portable( co_argcount=co_argcount, @@ -781,7 +799,7 @@ def t_code_old(self, _, bytes_for_s: bool = False): co_firstlineno = -1 # Bogus sentinel value co_lnotab = b"" - reference_objects = set(self.intern_objects + self.intern_strings) + reference_objects = self._reference_objects() code = to_portable( co_argcount=co_argcount,