def encode(self, compression: Optional[cmp.Compression] = None) -> bytes:
"""
byte array representation of log entry.
append CRC32 checksum of header and k/v
-----------------------------------------------------------------------
| block size | meta | key length | value length | key | value | crc32 |
-----------------------------------------------------------------------
"""
key, value, meta = self.key, self.value, self.meta
if compression and compression.isenabled:
key = compression.compress(key)
value = compression.compress(value)
header = self._encode_header(key=key, value=value, meta=meta)
checksum = crc32(header)
encoded = bytearray(header)
checksum = crc32(key, checksum)
encoded += key
checksum = crc32(value, checksum)
encoded += value
encoded += uvarint.encode(checksum)
block_size = uvarint.encode(len(encoded))
return bytes([*block_size, *encoded])
async def test_write_message(fragmented_message: Tuple[StreamID, MplexFlag,
StreamData]):
reader_mock, writer_mock = get_connection_mock("127.0.0.1", 7777)
stream_id, flag, data = fragmented_message
mplex_protocol = MplexProtocol(reader_mock, writer_mock)
await mplex_protocol.write_message(
MplexMessage(stream_id=stream_id, flag=flag, data=data))
encoded_message = (uvarint.encode(stream_id << 3 | flag) +
uvarint.encode(len(data)) + data)
writer_mock.write.assert_called_with(encoded_message)
writer_mock.drain.assert_awaited()
async def test_read_message(fragmented_message: Tuple[StreamID, MplexFlag,
StreamData]):
reader_mock, writer_mock = get_connection_mock("127.0.0.1", 7777)
stream_id, flag, data = fragmented_message
mplex_protocol = MplexProtocol(reader_mock, writer_mock)
encoded_message = (uvarint.encode(stream_id << 3 | flag) +
uvarint.encode(len(data)) + data)
reader_mock.feed_data(encoded_message)
message = await mplex_protocol.read_message()
assert isinstance(message, MplexMessage)
assert message.stream_id == stream_id
assert message.flag == flag
assert message.data == data
async def test_read_uvarint_overflow():
uvarint_overflow = bytearray(
[0b10000000 for _ in range(UVARINT_MAX_BYTES + 1)])
reader_mock, writer_mock = get_connection_mock("127.0.0.1", 7777)
stream_id, flag = 12, MplexFlag.NEW_STREAM
encoded_message = uvarint.encode(stream_id << 3 | flag) + uvarint_overflow
reader_mock.feed_data(encoded_message)
mplex_protocol = MplexProtocol(reader_mock, writer_mock)
with pytest.raises(OverflowError):
message = await mplex_protocol.read_message()
def _encode_header(cls, key: types.Key, value: types.Value,
meta: int) -> bytes:
"""
byte array representation of header fields/metadata
------------------------------------
| meta | key length | value length |
------------------------------------
"""
header_fields = [len(key), len(value)]
header = bytearray([meta])
for val in header_fields:
header += uvarint.encode(val)
return header
def to_bytes(self):
"""
"""
ba = []
if bool(self._inferiors):
# fixed width binary representation for all lower halves
inferiors_bits = int(reduce(lambda a, b: a + b,
map(lambda inf: ("{0:0%db}" % self._lower_bits).format(inf),
self._inferiors)),
2)
inferiors_byte_count = max(1, math.ceil(inferiors_bits.bit_length() / 8.0))
else:
inferiors_byte_count = 0
if bool(self._superiors):
# negated unary representation for upper halves
superiors_bits = int(reduce(lambda a, b: a + b,
map(lambda sup: ("1" * sup) + "0",
self._superiors)),
2)
superiors_byte_count = max(1, math.ceil(superiors_bits.bit_length() / 8.0))
else:
superiors_byte_count = 0
header = uvarint.encode(0) # EliasFano: 0, MultiLevelEliasFano: 1
header += uvarint.encode(self._n)
header += uvarint.encode(self._lower_bits)
header += uvarint.encode(self._upper_bits)
header += uvarint.encode(inferiors_byte_count)
header += uvarint.encode(superiors_byte_count)
ba.append(header)
if inferiors_byte_count:
ba.append(inferiors_bits.to_bytes(inferiors_byte_count, 'little', signed=False))
if superiors_byte_count:
ba.append(superiors_bits.to_bytes(superiors_byte_count, 'little', signed=False))
return len(b''.join(ba)), b''.join(ba)
def get_encoded_message(stream_name: StreamName, flag: MplexFlag,
data: StreamData) -> bytes:
encoded_message = (
uvarint.encode(get_stream_id_from_name(stream_name) << 3 | flag) +
uvarint.encode(len(data)) + data)
return encoded_message
def _encode_message(self, message: MplexMessage) -> bytes:
return (
uvarint.encode(message.stream_id << 3 | message.flag)
+ uvarint.encode(len(message.data))
+ message.data
)
def test_encode(self) -> None:
decoded: int
encoded: bytes
for (decoded, encoded) in TestUvarint.all:
self.assertEqual(uvarint.encode(decoded), encoded)