Avoid storing memoryview in deserialized objects

This commit is contained in:
Lonami Exo 2023-08-31 12:25:00 +02:00
parent 352ccdd4fc
commit 9e582ca0fa
3 changed files with 7 additions and 11 deletions

View File

@ -31,17 +31,17 @@ def _bootstrap_get_ty(constructor_id: int) -> Optional[Type["Serializable"]]:
class Reader: class Reader:
__slots__ = ("_view", "_pos", "_len") __slots__ = ("_view", "_pos", "_len")
def __init__(self, buffer: bytes) -> None: def __init__(self, buffer: bytes | memoryview) -> None:
self._view = ( self._view = (
memoryview(buffer) if not isinstance(buffer, memoryview) else buffer memoryview(buffer) if not isinstance(buffer, memoryview) else buffer
) )
self._pos = 0 self._pos = 0
self._len = len(self._view) self._len = len(self._view)
def read_remaining(self) -> bytes: def read_remaining(self) -> memoryview:
return self.read(self._len - self._pos) return self.read(self._len - self._pos)
def read(self, n: int) -> bytes: def read(self, n: int) -> memoryview:
self._pos += n self._pos += n
assert self._pos <= self._len assert self._pos <= self._len
return self._view[self._pos - n : self._pos] return self._view[self._pos - n : self._pos]
@ -52,7 +52,7 @@ class Reader:
assert self._pos <= self._len assert self._pos <= self._len
return struct.unpack(fmt, self._view[self._pos - size : self._pos]) return struct.unpack(fmt, self._view[self._pos - size : self._pos])
def read_bytes(self) -> bytes: def read_bytes(self) -> memoryview:
if self._view[self._pos] == 254: if self._view[self._pos] == 254:
self._pos += 4 self._pos += 4
length = struct.unpack("<i", self._view[self._pos - 4 : self._pos])[0] >> 8 length = struct.unpack("<i", self._view[self._pos - 4 : self._pos])[0] >> 8

View File

@ -36,10 +36,6 @@ class Serializable(abc.ABC):
def __repr__(self) -> str: def __repr__(self) -> str:
fields = ((attr, getattr(self, attr)) for attr in self.__slots__) fields = ((attr, getattr(self, attr)) for attr in self.__slots__)
fields = (
(name, bytes(field) if isinstance(field, memoryview) else field)
for name, field in fields
)
attrs = ", ".join(f"{name}={field!r}" for name, field in fields) attrs = ", ".join(f"{name}={field!r}" for name, field in fields)
return f"{self.__class__.__name__}({attrs})" return f"{self.__class__.__name__}({attrs})"

View File

@ -9,8 +9,8 @@ from .common import inner_type_fmt, is_trivial, to_class_name, trivial_struct_fm
# Some implementations choose to create these types by hand. # Some implementations choose to create these types by hand.
# For consistency, we instead special-case the generator. # For consistency, we instead special-case the generator.
SPECIAL_CASED_OBJECT_READS = { SPECIAL_CASED_OBJECT_READS = {
0xF35C6D01: "reader.read_remaining()", # rpc_result 0xF35C6D01: "b'' + reader.read_remaining()", # rpc_result
0x5BB8E511: "reader.read(_bytes)", # message 0x5BB8E511: "b'' + reader.read(_bytes)", # message
} }
@ -22,7 +22,7 @@ def reader_read_fmt(ty: Type, constructor_id: int) -> Tuple[str, Optional[str]]:
elif ty.name == "string": elif ty.name == "string":
return f"str(reader.read_bytes(), 'utf-8', 'replace')", None return f"str(reader.read_bytes(), 'utf-8', 'replace')", None
elif ty.name == "bytes": elif ty.name == "bytes":
return f"reader.read_bytes()", None return f"b'' + reader.read_bytes()", None
elif ty.name == "int128": elif ty.name == "int128":
return f"int.from_bytes(reader.read(16))", None return f"int.from_bytes(reader.read(16))", None
elif ty.name == "int256": elif ty.name == "int256":