chia-blockchain/chia/util/struct_stream.py
Arvid Norberg ca3c8ce407
validate integers in StructStream constructors (#4298)
* validate integers in StructStream constructors

* don't step back sub slots beyond zero

* fix off-by-one in blocks-to-remove

* handle find_fork_point_in_chain() returning -1 in wallet

* fix handling of not finding a shared ancestor in wallet reorg

* terminate clean_block_record() when we reach height 0

* harden int/uint validation in constructor. add comment about int512
2021-06-08 10:45:56 -07:00

47 lines
1.4 KiB
Python

import io
import struct
from typing import Any, BinaryIO
class StructStream(int):
PACK = ""
"""
Create a class that can parse and stream itself based on a struct.pack template string.
"""
def __new__(cls: Any, value: int):
value = int(value)
try:
v1 = struct.unpack(cls.PACK, struct.pack(cls.PACK, value))[0]
if value != v1:
raise ValueError(f"Value {value} does not fit into {cls.__name__}")
except Exception:
bits = struct.calcsize(cls.PACK) * 8
raise ValueError(
f"Value {value} of size {value.bit_length()} does not fit into " f"{cls.__name__} of size {bits}"
)
return int.__new__(cls, value) # type: ignore
@classmethod
def parse(cls: Any, f: BinaryIO) -> Any:
bytes_to_read = struct.calcsize(cls.PACK)
read_bytes = f.read(bytes_to_read)
assert read_bytes is not None and len(read_bytes) == bytes_to_read
return cls(*struct.unpack(cls.PACK, read_bytes))
def stream(self, f):
f.write(struct.pack(self.PACK, self))
@classmethod
def from_bytes(cls: Any, blob: bytes) -> Any: # type: ignore
f = io.BytesIO(blob)
result = cls.parse(f)
assert f.read() == b""
return result
def __bytes__(self: Any) -> bytes:
f = io.BytesIO()
self.stream(f)
return bytes(f.getvalue())