1# -*- coding: utf-8 -*- 

2 

3# -- stdlib -- 

4from enum import IntEnum 

5from typing import Any, Iterator, Sequence, Tuple, List 

6import logging 

7import zlib 

8 

9# -- third party -- 

10from gevent import socket 

11from gevent.lock import RLock 

12from gevent.timeout import Timeout 

13import msgpack 

14 

15# -- own -- 

16import wire 

17 

18 

19# -- code -- 

20log = logging.getLogger("Endpoint") 

21log.setLevel(logging.ERROR) 

22 

23 

24class EndpointDied(Exception): 

25 pass 

26 

27 

28class DecodeError(Exception): 

29 pass 

30 

31 

32class Format(IntEnum): 

33 Packed = 1 

34 BulkCompressed = 2 

35 

36 

37class Endpoint(object): 

38 

39 def __init__(self, sock, address): 

40 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) 

41 sock.read = sock.recv 

42 sock.write = sock.sendall 

43 

44 self.sock = sock 

45 self.unpacker = msgpack.Unpacker(sock, raw=False, strict_map_key=False) 

46 self.writelock = RLock() 

47 self.address = address 

48 self.link_state = 'connected' # or disconnected 

49 

50 self.active = False 

51 

52 def __repr__(self): 

53 return '%s:%s:%s' % ( 

54 self.__class__.__name__, 

55 self.address[0], 

56 self.address[1], 

57 ) 

58 

59 @staticmethod 

60 def encode(p: wire.Message) -> bytes: 

61 return msgpack.packb([Format.Packed, p.encode()], use_bin_type=True) 

62 

63 @staticmethod 

64 def encode_bulk(pl: Sequence[wire.Message]) -> bytes: 

65 data = msgpack.packb([i.encode() for i in pl], use_bin_type=True) 

66 return msgpack.packb([Format.BulkCompressed, zlib.compress(data)], use_bin_type=True) 

67 

68 def write(self, p) -> None: 

69 log.debug("%s SEND>> %s", repr(self), p) 

70 self.raw_write(self.encode(p)) 

71 

72 def write_bulk(self, pl: Sequence[wire.Message]) -> None: 

73 for p in pl: 

74 log.debug("%s SEND>> %s", repr(self), p) 

75 self.raw_write(self.encode_bulk(pl)) 

76 

77 def raw_write(self, s: bytes) -> None: 

78 if self.link_state == 'connected': 78 ↛ exitline 78 didn't return from function 'raw_write', because the condition on line 78 was never false

79 try: 

80 with self.writelock: 

81 self.sock.sendall(s) 

82 self.active = True 

83 except IOError: 

84 log.exception('Error raw_write, closing Endpoint') 

85 self.close() 

86 

87 def close(self): 

88 if not self.link_state == 'disconnected': 88 ↛ exitline 88 didn't return from function 'close', because the condition on line 88 was never false

89 self.link_state = 'disconnected' 

90 self.sock.close() 

91 self.active = False 

92 

93 @staticmethod 

94 def _decode_packet(p: Any) -> Tuple[Format, Any]: 

95 try: 

96 if not (isinstance(p, (list, tuple)) and len(p) == 2): 96 ↛ 97line 96 didn't jump to line 97, because the condition on line 96 was never true

97 raise DecodeError 

98 

99 fmt, data = p 

100 fmt = Format(fmt) 

101 if fmt == Format.Packed: 101 ↛ 103line 101 didn't jump to line 103, because the condition on line 101 was never false

102 return fmt, data 

103 elif fmt == Format.BulkCompressed: 

104 try: 

105 inflated = zlib.decompress(data) 

106 except Exception: 

107 raise DecodeError 

108 

109 return fmt, msgpack.unpackb(inflated, raw=False, ext_hook=lambda code, data: None) 

110 else: 

111 raise DecodeError 

112 except (ValueError, msgpack.UnpackValueError): 

113 raise DecodeError 

114 

115 @staticmethod 

116 def decode_bytes(s: bytes) -> List[wire.Message]: 

117 p = msgpack.unpackb(s, raw=False) 

118 fmt, data = Endpoint._decode_packet(p) 

119 if fmt == Format.Packed: 

120 msg = wire.Message.decode(data) 

121 if not msg: 

122 raise DecodeError 

123 return [msg] 

124 elif fmt == Format.BulkCompressed: 

125 return data 

126 

127 assert False, 'WTF' 

128 

129 def messages(self, timeout=90) -> Iterator[wire.Message]: 

130 if self.link_state != 'connected': 130 ↛ 131line 130 didn't jump to line 131, because the condition on line 130 was never true

131 raise EndpointDied 

132 

133 unpacker = self.unpacker 

134 _NONE = object() 

135 while True: 

136 try: 

137 v = _NONE 

138 with Timeout(timeout, False): 

139 try: 

140 v = next(unpacker) 

141 self.active = True 

142 except StopIteration: 

143 pass 

144 

145 if v is _NONE: 

146 break 

147 

148 fmt, data = self._decode_packet(v) 

149 if fmt == Format.Packed: 149 ↛ 154line 149 didn't jump to line 154, because the condition on line 149 was never false

150 msg = wire.Message.decode(data) 

151 if msg: 151 ↛ 136line 151 didn't jump to line 136, because the condition on line 151 was never false

152 log.debug("%s <<RECV %r", repr(self), msg) 

153 yield msg 

154 elif fmt == Format.BulkCompressed: 

155 for d in data: 

156 msg = wire.Message.decode(d) 

157 if msg: 

158 log.debug("%s <<RECV %r", repr(self), msg) 

159 yield msg 

160 

161 except DecodeError: 161 ↛ 162line 161 didn't jump to line 162, because the exception caught by line 161 didn't happen

162 log.info('DecodeError') 

163 break 

164 

165 except msgpack.UnpackValueError: 

166 log.exception('UnpackValueError') 

167 break 

168 

169 self.close() 

170 raise EndpointDied