From b2146c52681387578f9a057936a03b3a51978ff3 Mon Sep 17 00:00:00 2001 From: Brian Cloutier Date: Mon, 9 Sep 2019 14:12:38 -0700 Subject: [PATCH] Don't crash on large messages --- libp2p/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libp2p/utils.py b/libp2p/utils.py index 7374993..c69f61b 100644 --- a/libp2p/utils.py +++ b/libp2p/utils.py @@ -4,6 +4,8 @@ import math from libp2p.exceptions import ParseError from libp2p.io.abc import Reader +from .io.utils import read_exactly + # Unsigned LEB128(varint codec) # Reference: https://github.com/ethereum/py-wasm/blob/master/wasm/parsers/leb128.py @@ -62,7 +64,7 @@ def encode_varint_prefixed(msg_bytes: bytes) -> bytes: async def read_varint_prefixed_bytes(reader: Reader) -> bytes: len_msg = await decode_uvarint_from_stream(reader) - data = await reader.read(len_msg) + data = await read_exactly(reader, len_msg) if len(data) != len_msg: raise ValueError( f"failed to read enough bytes: len_msg={len_msg}, data={data!r}"