Coverage for zenmqtt/mqtt/protocol/stream.py: 100%

1from typing import Final, Optional, Tuple 

2 

3from zenmqtt.connection import MQTTConnection 

4from zenmqtt.exceptions import IncomingPacketTooLargeError, MalformedPacketError 

5from zenmqtt.mqtt.packet import ( 

6 FIXED_HEADER_SIZE, 

7 BytesReader, 

8 FixedHeader, 

9 PacketType, 

10 parse_variable_byte_integer, 

11) 

12 

13# bytes requested from the connection by one read 

14READ_SIZE: Final[int] = 64 * 1024 

15 

16# "Remaining Length" is encoded in at most 4 bytes 

17_MAX_LENGTH_SIZE: Final[int] = 4 

18 

19 

20class PacketReader: 

21 """ 

22 Reads whole packets from the connection. The connection is read in large 

23 chunks, and all complete packets in the buffer are returned before the 

24 connection is read again. 

25 """ 

26 

27 def __init__(self, connection: MQTTConnection) -> None: 

28 self._connection = connection 

29 self._buffer = bytearray() 

30 self._offset = 0 

31 

32 # "Maximum Packet Size" of the client, None means no limit 

33 self.maximum_packet_size: Optional[int] = None 

34 

35 async def read_packet(self) -> Optional[Tuple[FixedHeader, BytesReader, int]]: 

36 """ 

37 Returns the fixed header, a reader of the body and the size of the packet, or None 

38 when the connection is closed. 

39 

40 :raises IncomingPacketTooLargeError: the packet is bigger than 

41 maximum_packet_size; its body isn't read 

42 :raises MalformedPacketError: invalid "Remaining Length" 

43 """ 

44 # the first byte and "Remaining Length" of 1..4 bytes 

45 header_size = FIXED_HEADER_SIZE 

46 

47 while True: 

48 header_size += 1 

49 

50 if not await self._fill(header_size): 

51 return None 

52 

53 # "Remaining Length" is a variable byte integer: 7 bits of value per 

54 # byte, the high bit (0x80) means that one more byte follows. The 

55 # size of the length isn't known in advance, so the buffer is filled 

56 # byte by byte until a byte without this bit ends the length; only 

57 # then the length can be decoded and the body size is known. 

58 last_byte = self._buffer[self._offset + header_size - 1] 

59 

60 if not last_byte & 0x80: 

61 break 

62 

63 if header_size - FIXED_HEADER_SIZE == _MAX_LENGTH_SIZE: 

64 raise MalformedPacketError("Remaining length is longer than 4 bytes") 

65 

66 length, _ = parse_variable_byte_integer(self._buffer, self._offset + 1) 

67 size = header_size + length 

68 

69 if self.maximum_packet_size is not None and size > self.maximum_packet_size: 

70 raise IncomingPacketTooLargeError( 

71 f"Packet of {size} bytes, maximum is {self.maximum_packet_size}" 

72 ) 

73 

74 if not await self._fill(size): 

75 return None 

76 

77 # the packet starts at the offset: _fill() may have moved the data to 

78 # the beginning of the buffer 

79 start = self._offset 

80 first_byte = self._buffer[start] 

81 body_start, end = start + header_size, start + size 

82 

83 # one copy instead of two (slice of bytearray and bytes of it); the 

84 # view is released at once, the buffer can't be resized while it lives 

85 with memoryview(self._buffer) as view: 

86 body = bytes(view[body_start:end]) 

87 

88 self._offset = end 

89 

90 header = FixedHeader( 

91 packet_type=PacketType(first_byte >> 4), 

92 flags=first_byte & 0xF, 

93 length=length, 

94 ) 

95 

96 return header, BytesReader(body), size 

97 

98 async def _fill(self, size: int) -> bool: 

99 """ 

100 Reads the connection until the buffer has `size` bytes after the 

101 offset. Returns False if the connection is closed before. 

102 """ 

103 while len(self._buffer) - self._offset < size: 

104 # drop packets which were returned already 

105 if self._offset: 

106 del self._buffer[: self._offset] 

107 self._offset = 0 

108 

109 data = await self._connection.read(READ_SIZE) 

110 

111 if not data: 

112 return False 

113 

114 self._buffer += data 

115 

116 return True