Fix byte preservation for Python 2 interned strings

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("<i", len(value))
+    return tag + size + value
+
+
+def reader(data, version="2.7"):
+    return VersionIndependentUnmarshaller(
+        io.BytesIO(data), magic2int(by_version[version]), False, code_objects={}
+    )
+
+
+@pytest.mark.parametrize("version", ("2.4", "2.7"))
+@pytest.mark.parametrize(
+    "value",
+    (b"", b"ascii_value", b"a\0b", b"\xffA\x80", b"\xc3\xa9", bytes(range(256))),
+    ids=("empty", "ascii", "nul", "non-utf8", "utf8", "all-bytes"),
+)
+def test_python2_interned_bytes_match_ordinary_strings(version, value):
+    ordinary = reader(string_record(b"s", value), version).r_object(bytes_for_s=True)
+    interned = reader(string_record(b"t", value), version).r_object(bytes_for_s=True)
+    assert isinstance(interned, bytes)
+    assert ordinary == interned == value
+
+
+@pytest.mark.parametrize("version", ("2.4", "2.7"))
+@pytest.mark.parametrize("first_bytes", (False, True))
+@pytest.mark.parametrize(
+    "raw, text", ((b"shared_value", "shared_value"), (b"\xffA\x80", "A"))
+)
+def test_python2_references_preserve_both_representations(
+    version, first_bytes, raw, text
+):
+    reference = b"R" + pack("<i", 0)
+    unmarshaler = reader(string_record(b"t", raw) + reference * 2, version)
+    first = unmarshaler.r_object(bytes_for_s=first_bytes)
+    other = unmarshaler.r_object(bytes_for_s=not first_bytes)
+    repeated = unmarshaler.r_object(bytes_for_s=first_bytes)
+    assert first == (raw if first_bytes else text)
+    assert other == (text if first_bytes else raw)
+    assert repeated is first
+
+
+def test_python2_interned_tuple_from_cpython():
+    # CPython 2.7.18: value = intern('\xffA\x80'); marshal.dumps((value, value), 2)
+    data = bytes.fromhex("28020000007403000000ff41805200000000")
+    result = reader(data).r_object(bytes_for_s=True)
+    assert result == (b"\xffA\x80", b"\xffA\x80")
+    assert result[0] is result[1]
+
+
+@pytest.mark.parametrize("bytes_for_s", (False, True))
+@pytest.mark.parametrize(
+    "version, tag, text",
+    (
+        ("3.3", b"t", "\u00e9"),
+        ("3.4", b"t", "\u00e9"),
+        ("3.4", b"A", "shared_value"),
+        ("3.4", b"Z", "shared_value"),
+    ),
+)
+def test_python3_interned_strings_remain_text(version, tag, text, bytes_for_s):
+    result = reader(string_record(tag, text.encode("utf8")), version).r_object(
+        bytes_for_s=bytes_for_s
+    )
+    assert isinstance(result, str)
+    assert result == text
+
+
+@pytest.mark.parametrize("bytes_for_s", (False, True))
+def test_python3_object_references_remain_text(bytes_for_s):
+    data = (
+        string_record(bytes([ord("t") | 0x80]), b"shared_value") + b"r" + pack("<i", 0)
+    )
+    unmarshaler = reader(data, "3.4")
+    first = unmarshaler.r_object(bytes_for_s=bytes_for_s)
+    second = unmarshaler.r_object(bytes_for_s=bytes_for_s)
+    assert first == "shared_value"
+    assert isinstance(first, str)
+    assert second is first
+
+
+@pytest.mark.parametrize("code_tag", (b"s", b"t"))
+def test_python27_pyc_keeps_interned_bytecode_and_text_metadata(code_tag):
+    # Generated by stock CPython 2.7.18: compile('fn(*args)', '<fixture>', '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("<I", 0) + data
+    version, _, _, code, *_ = load_module_from_file_object(io.BytesIO(pyc))
+    expected = bytes.fromhex("6500006501008c00000164000053")
+    assert version == (2, 7)
+    assert code.co_code == expected
+    assert isinstance(code.co_code, bytes)
+    assert code.co_consts == (None,)
+    assert code.co_names == ("fn", "args")
+    assert code.co_filename == "<fixture>"
+    assert code.co_name == "<module>"
+
+    # 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

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,