11130Szelenkov@nginx.comimport base64 21130Szelenkov@nginx.comimport hashlib 31130Szelenkov@nginx.comimport itertools 41596Szelenkov@nginx.comimport pytest 51477Szelenkov@nginx.comimport random 61477Szelenkov@nginx.comimport re 71477Szelenkov@nginx.comimport select 81477Szelenkov@nginx.comimport struct 91477Szelenkov@nginx.com 101130Szelenkov@nginx.comfrom unit.applications.proto import TestApplicationProto 111130Szelenkov@nginx.com 121130Szelenkov@nginx.comGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" 131130Szelenkov@nginx.com 141130Szelenkov@nginx.com 151130Szelenkov@nginx.comclass TestApplicationWebsocket(TestApplicationProto): 161130Szelenkov@nginx.com 171130Szelenkov@nginx.com OP_CONT = 0x00 181130Szelenkov@nginx.com OP_TEXT = 0x01 191130Szelenkov@nginx.com OP_BINARY = 0x02 201130Szelenkov@nginx.com OP_CLOSE = 0x08 211130Szelenkov@nginx.com OP_PING = 0x09 221130Szelenkov@nginx.com OP_PONG = 0x0A 231130Szelenkov@nginx.com CLOSE_CODES = [1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011] 241130Szelenkov@nginx.com 251130Szelenkov@nginx.com def key(self): 261130Szelenkov@nginx.com raw_key = bytes(random.getrandbits(8) for _ in range(16)) 271130Szelenkov@nginx.com return base64.b64encode(raw_key).decode() 281130Szelenkov@nginx.com 291130Szelenkov@nginx.com def accept(self, key): 301130Szelenkov@nginx.com sha1 = hashlib.sha1((key + GUID).encode()).digest() 311130Szelenkov@nginx.com return base64.b64encode(sha1).decode() 321130Szelenkov@nginx.com 331262Szelenkov@nginx.com def upgrade(self, headers=None): 341262Szelenkov@nginx.com key = None 351262Szelenkov@nginx.com 361262Szelenkov@nginx.com if headers is None: 371262Szelenkov@nginx.com key = self.key() 381262Szelenkov@nginx.com headers = { 391130Szelenkov@nginx.com 'Host': 'localhost', 401130Szelenkov@nginx.com 'Upgrade': 'websocket', 411130Szelenkov@nginx.com 'Connection': 'Upgrade', 421130Szelenkov@nginx.com 'Sec-WebSocket-Key': key, 43*1625Smax.romanov@nginx.com 'Sec-WebSocket-Protocol': 'chat, phone, video', 441130Szelenkov@nginx.com 'Sec-WebSocket-Version': 13, 451262Szelenkov@nginx.com } 461262Szelenkov@nginx.com 471262Szelenkov@nginx.com _, sock = self.get( 481262Szelenkov@nginx.com headers=headers, 491246Szelenkov@nginx.com no_recv=True, 501130Szelenkov@nginx.com start=True, 511130Szelenkov@nginx.com ) 521130Szelenkov@nginx.com 531246Szelenkov@nginx.com resp = '' 541444Szelenkov@nginx.com while True: 551444Szelenkov@nginx.com rlist = select.select([sock], [], [], 60)[0] 561444Szelenkov@nginx.com if not rlist: 571596Szelenkov@nginx.com pytest.fail('Can\'t read response from server.') 581444Szelenkov@nginx.com 591246Szelenkov@nginx.com resp += sock.recv(4096).decode() 601246Szelenkov@nginx.com 611609Smax.romanov@nginx.com if (resp.startswith('HTTP/') and '\r\n\r\n' in resp): 621246Szelenkov@nginx.com resp = self._resp_to_dict(resp) 631246Szelenkov@nginx.com break 641246Szelenkov@nginx.com 651130Szelenkov@nginx.com return (resp, sock, key) 661130Szelenkov@nginx.com 671130Szelenkov@nginx.com def apply_mask(self, data, mask): 681130Szelenkov@nginx.com return bytes(b ^ m for b, m in zip(data, itertools.cycle(mask))) 691130Szelenkov@nginx.com 701150Szelenkov@nginx.com def serialize_close(self, code=1000, reason=''): 711130Szelenkov@nginx.com return struct.pack('!H', code) + reason.encode('utf-8') 721130Szelenkov@nginx.com 731424Szelenkov@nginx.com def frame_read(self, sock, read_timeout=60): 741130Szelenkov@nginx.com def recv_bytes(sock, bytes): 751130Szelenkov@nginx.com data = b'' 761444Szelenkov@nginx.com while True: 771444Szelenkov@nginx.com rlist = select.select([sock], [], [], read_timeout)[0] 781444Szelenkov@nginx.com if not rlist: 791444Szelenkov@nginx.com # For all current cases if the "read_timeout" was changed 801444Szelenkov@nginx.com # than test do not expect to get a response from server. 811444Szelenkov@nginx.com if read_timeout == 60: 821596Szelenkov@nginx.com pytest.fail('Can\'t read response from server.') 831444Szelenkov@nginx.com break 841444Szelenkov@nginx.com 851154Szelenkov@nginx.com data += sock.recv(bytes - len(data)) 861154Szelenkov@nginx.com 871154Szelenkov@nginx.com if len(data) == bytes: 881130Szelenkov@nginx.com break 891130Szelenkov@nginx.com 901130Szelenkov@nginx.com return data 911130Szelenkov@nginx.com 921130Szelenkov@nginx.com frame = {} 931130Szelenkov@nginx.com 941130Szelenkov@nginx.com head1, = struct.unpack('!B', recv_bytes(sock, 1)) 951130Szelenkov@nginx.com head2, = struct.unpack('!B', recv_bytes(sock, 1)) 961130Szelenkov@nginx.com 971130Szelenkov@nginx.com frame['fin'] = bool(head1 & 0b10000000) 981130Szelenkov@nginx.com frame['rsv1'] = bool(head1 & 0b01000000) 991130Szelenkov@nginx.com frame['rsv2'] = bool(head1 & 0b00100000) 1001130Szelenkov@nginx.com frame['rsv3'] = bool(head1 & 0b00010000) 1011130Szelenkov@nginx.com frame['opcode'] = head1 & 0b00001111 1021130Szelenkov@nginx.com frame['mask'] = head2 & 0b10000000 1031130Szelenkov@nginx.com 1041130Szelenkov@nginx.com length = head2 & 0b01111111 1051130Szelenkov@nginx.com if length == 126: 1061130Szelenkov@nginx.com data = recv_bytes(sock, 2) 1071130Szelenkov@nginx.com length, = struct.unpack('!H', data) 1081130Szelenkov@nginx.com elif length == 127: 1091130Szelenkov@nginx.com data = recv_bytes(sock, 8) 1101130Szelenkov@nginx.com length, = struct.unpack('!Q', data) 1111130Szelenkov@nginx.com 1121130Szelenkov@nginx.com if frame['mask']: 1131130Szelenkov@nginx.com mask_bits = recv_bytes(sock, 4) 1141130Szelenkov@nginx.com 1151156Szelenkov@nginx.com data = b'' 1161156Szelenkov@nginx.com 1171156Szelenkov@nginx.com if length != 0: 1181156Szelenkov@nginx.com data = recv_bytes(sock, length) 1191156Szelenkov@nginx.com 1201130Szelenkov@nginx.com if frame['mask']: 1211130Szelenkov@nginx.com data = self.apply_mask(data, mask_bits) 1221130Szelenkov@nginx.com 1231130Szelenkov@nginx.com if frame['opcode'] == self.OP_CLOSE: 1241130Szelenkov@nginx.com if length >= 2: 1251130Szelenkov@nginx.com code, = struct.unpack('!H', data[:2]) 1261130Szelenkov@nginx.com reason = data[2:].decode('utf-8') 1271130Szelenkov@nginx.com if not (code in self.CLOSE_CODES or 3000 <= code < 5000): 1281596Szelenkov@nginx.com pytest.fail('Invalid status code') 1291130Szelenkov@nginx.com frame['code'] = code 1301130Szelenkov@nginx.com frame['reason'] = reason 1311130Szelenkov@nginx.com elif length == 0: 1321130Szelenkov@nginx.com frame['code'] = 1005 1331130Szelenkov@nginx.com frame['reason'] = '' 1341130Szelenkov@nginx.com else: 1351596Szelenkov@nginx.com pytest.fail('Close frame too short') 1361130Szelenkov@nginx.com 1371130Szelenkov@nginx.com frame['data'] = data 1381130Szelenkov@nginx.com 1391130Szelenkov@nginx.com if frame['mask']: 1401596Szelenkov@nginx.com pytest.fail('Received frame with mask') 1411130Szelenkov@nginx.com 1421130Szelenkov@nginx.com return frame 1431130Szelenkov@nginx.com 1441130Szelenkov@nginx.com def frame_to_send( 1451130Szelenkov@nginx.com self, 1461130Szelenkov@nginx.com opcode, 1471130Szelenkov@nginx.com data, 1481130Szelenkov@nginx.com fin=True, 1491130Szelenkov@nginx.com length=None, 1501130Szelenkov@nginx.com rsv1=False, 1511130Szelenkov@nginx.com rsv2=False, 1521130Szelenkov@nginx.com rsv3=False, 1531130Szelenkov@nginx.com mask=True, 1541130Szelenkov@nginx.com ): 1551130Szelenkov@nginx.com frame = b'' 1561130Szelenkov@nginx.com 1571130Szelenkov@nginx.com if isinstance(data, str): 1581130Szelenkov@nginx.com data = data.encode('utf-8') 1591130Szelenkov@nginx.com 1601130Szelenkov@nginx.com head1 = ( 1611130Szelenkov@nginx.com (0b10000000 if fin else 0) 1621130Szelenkov@nginx.com | (0b01000000 if rsv1 else 0) 1631130Szelenkov@nginx.com | (0b00100000 if rsv2 else 0) 1641130Szelenkov@nginx.com | (0b00010000 if rsv3 else 0) 1651130Szelenkov@nginx.com | opcode 1661130Szelenkov@nginx.com ) 1671130Szelenkov@nginx.com 1681130Szelenkov@nginx.com head2 = 0b10000000 if mask else 0 1691130Szelenkov@nginx.com 1701130Szelenkov@nginx.com data_length = len(data) if length is None else length 1711130Szelenkov@nginx.com if data_length < 126: 1721130Szelenkov@nginx.com frame += struct.pack('!BB', head1, head2 | data_length) 1731130Szelenkov@nginx.com elif data_length < 65536: 1741130Szelenkov@nginx.com frame += struct.pack('!BBH', head1, head2 | 126, data_length) 1751130Szelenkov@nginx.com else: 1761130Szelenkov@nginx.com frame += struct.pack('!BBQ', head1, head2 | 127, data_length) 1771130Szelenkov@nginx.com 1781130Szelenkov@nginx.com if mask: 1791130Szelenkov@nginx.com mask_bits = struct.pack('!I', random.getrandbits(32)) 1801130Szelenkov@nginx.com frame += mask_bits 1811130Szelenkov@nginx.com 1821130Szelenkov@nginx.com if mask: 1831130Szelenkov@nginx.com frame += self.apply_mask(data, mask_bits) 1841130Szelenkov@nginx.com else: 1851130Szelenkov@nginx.com frame += data 1861130Szelenkov@nginx.com 1871130Szelenkov@nginx.com return frame 1881130Szelenkov@nginx.com 1891130Szelenkov@nginx.com def frame_write(self, sock, *args, **kwargs): 1901130Szelenkov@nginx.com chopsize = kwargs.pop('chopsize') if 'chopsize' in kwargs else None 1911130Szelenkov@nginx.com 1921130Szelenkov@nginx.com frame = self.frame_to_send(*args, **kwargs) 1931130Szelenkov@nginx.com 1941130Szelenkov@nginx.com if chopsize is None: 1951152Szelenkov@nginx.com try: 1961152Szelenkov@nginx.com sock.sendall(frame) 1971152Szelenkov@nginx.com except BrokenPipeError: 1981152Szelenkov@nginx.com pass 1991130Szelenkov@nginx.com 2001130Szelenkov@nginx.com else: 2011130Szelenkov@nginx.com pos = 0 2021130Szelenkov@nginx.com frame_len = len(frame) 2031150Szelenkov@nginx.com while pos < frame_len: 2041130Szelenkov@nginx.com end = min(pos + chopsize, frame_len) 2051151Szelenkov@nginx.com try: 2061151Szelenkov@nginx.com sock.sendall(frame[pos:end]) 2071151Szelenkov@nginx.com except BrokenPipeError: 2081151Szelenkov@nginx.com end = frame_len 2091130Szelenkov@nginx.com pos = end 2101130Szelenkov@nginx.com 2111130Szelenkov@nginx.com def message(self, sock, type, message, fragmention_size=None, **kwargs): 2121130Szelenkov@nginx.com message_len = len(message) 2131130Szelenkov@nginx.com 2141130Szelenkov@nginx.com if fragmention_size is None: 2151130Szelenkov@nginx.com fragmention_size = message_len 2161130Szelenkov@nginx.com 2171130Szelenkov@nginx.com if message_len <= fragmention_size: 2181130Szelenkov@nginx.com self.frame_write(sock, type, message, **kwargs) 2191130Szelenkov@nginx.com return 2201130Szelenkov@nginx.com 2211130Szelenkov@nginx.com pos = 0 2221130Szelenkov@nginx.com op_code = type 2231150Szelenkov@nginx.com while pos < message_len: 2241130Szelenkov@nginx.com end = min(pos + fragmention_size, message_len) 2251150Szelenkov@nginx.com fin = end == message_len 2261150Szelenkov@nginx.com self.frame_write( 2271150Szelenkov@nginx.com sock, op_code, message[pos:end], fin=fin, **kwargs 2281150Szelenkov@nginx.com ) 2291130Szelenkov@nginx.com op_code = self.OP_CONT 2301130Szelenkov@nginx.com pos = end 2311130Szelenkov@nginx.com 2321433Szelenkov@nginx.com def message_read(self, sock, read_timeout=60): 2331130Szelenkov@nginx.com frame = self.frame_read(sock, read_timeout=read_timeout) 2341130Szelenkov@nginx.com 2351150Szelenkov@nginx.com while not frame['fin']: 2361130Szelenkov@nginx.com temp = self.frame_read(sock, read_timeout=read_timeout) 2371130Szelenkov@nginx.com frame['data'] += temp['data'] 2381130Szelenkov@nginx.com frame['fin'] = temp['fin'] 2391130Szelenkov@nginx.com 2401130Szelenkov@nginx.com return frame 241