Coverage for zenmqtt/connection/tcp.py: 100%

1import asyncio 

2import ssl 

3from typing import Final, Protocol 

4from urllib.parse import ParseResult 

5 

6from .base import ConnectionOptions, MQTTConnectionTransport 

7 

8DEFAULT_PORT: Final[int] = 1883 

9DEFAULT_TLS_PORT: Final[int] = 8883 

10 

11 

12class StreamReader(Protocol): 

13 async def read(self, size: int) -> bytes: 

14 """Up to size bytes, b"" at the end of the stream.""" 

15 

16 def feed_eof(self) -> None: 

17 """Ends the stream: read() returns b"" after the buffered data.""" 

18 

19 

20class StreamWriter(Protocol): 

21 def write(self, payload: bytes) -> None: 

22 """Puts the bytes into the send buffer.""" 

23 

24 async def drain(self) -> None: 

25 """Waits while the send buffer is above its high-water mark.""" 

26 

27 def write_eof(self) -> None: 

28 """Closes the write side of the connection (half-close).""" 

29 

30 def can_write_eof(self) -> bool: 

31 """False for TLS: it doesn't support half-close.""" 

32 

33 def is_closing(self) -> bool: 

34 """True when the stream is closed or is being closed.""" 

35 

36 def close(self) -> None: 

37 """Starts closing the stream.""" 

38 

39 async def wait_closed(self) -> None: 

40 """Waits until the stream is closed.""" 

41 

42 

43class TCPConnectionTransport(MQTTConnectionTransport): 

44 def __init__(self, reader: StreamReader, writer: StreamWriter): 

45 self._reader = reader 

46 self._writer = writer 

47 

48 async def write(self, payload: bytes) -> None: 

49 self._writer.write(payload) 

50 

51 # backpressure: waits while the send buffer is above its high-water 

52 # mark (64 KiB by default), e.g. the broker doesn't read 

53 try: 

54 await self._writer.drain() 

55 except ConnectionError: 

56 # the connection is lost, the read loop handles it; the data is 

57 # lost like with a write to a closed transport 

58 pass 

59 

60 async def close(self) -> None: 

61 self._reader.feed_eof() 

62 

63 # the connection may be broken already (e.g. reset by the peer): closing 

64 # it is still fine, the errors only say that the socket is gone 

65 try: 

66 # TLS transports don't support half-close 

67 if self._writer.can_write_eof(): 

68 self._writer.write_eof() 

69 except OSError: 

70 pass 

71 

72 self._writer.close() 

73 

74 try: 

75 await self._writer.wait_closed() 

76 except OSError: 

77 pass 

78 

79 async def read(self, size: int = -1) -> bytes: 

80 return await self._reader.read(size) 

81 

82 def is_closing(self) -> bool: 

83 return self._writer.is_closing() 

84 

85 

86async def build_tcp_connection_transport( 

87 url: ParseResult, options: ConnectionOptions 

88) -> TCPConnectionTransport: 

89 reader, writer = await asyncio.open_connection( 

90 host=url.hostname, port=url.port or DEFAULT_PORT 

91 ) 

92 return TCPConnectionTransport(reader, writer) 

93 

94 

95async def build_tls_connection_transport( 

96 url: ParseResult, options: ConnectionOptions 

97) -> TCPConnectionTransport: 

98 reader, writer = await asyncio.open_connection( 

99 host=url.hostname, 

100 port=url.port or DEFAULT_TLS_PORT, 

101 ssl=options.get("ssl") or ssl.create_default_context(), 

102 ) 

103 return TCPConnectionTransport(reader, writer)