xref: /unit/test/unit/applications/websockets.py (revision 2073:bc6ad31ce286)
11130Szelenkov@nginx.comimport base64
21130Szelenkov@nginx.comimport hashlib
31130Szelenkov@nginx.comimport itertools
41477Szelenkov@nginx.comimport random
51477Szelenkov@nginx.comimport select
61477Szelenkov@nginx.comimport struct
71477Szelenkov@nginx.com
81635Szelenkov@nginx.comimport pytest
91130Szelenkov@nginx.comfrom unit.applications.proto import TestApplicationProto
101130Szelenkov@nginx.com
111130Szelenkov@nginx.comGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
121130Szelenkov@nginx.com
131130Szelenkov@nginx.com
141130Szelenkov@nginx.comclass TestApplicationWebsocket(TestApplicationProto):
151130Szelenkov@nginx.com
161130Szelenkov@nginx.com    OP_CONT = 0x00
171130Szelenkov@nginx.com    OP_TEXT = 0x01
181130Szelenkov@nginx.com    OP_BINARY = 0x02
191130Szelenkov@nginx.com    OP_CLOSE = 0x08
201130Szelenkov@nginx.com    OP_PING = 0x09
211130Szelenkov@nginx.com    OP_PONG = 0x0A
221130Szelenkov@nginx.com    CLOSE_CODES = [1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011]
231130Szelenkov@nginx.com
241130Szelenkov@nginx.com    def key(self):
251130Szelenkov@nginx.com        raw_key = bytes(random.getrandbits(8) for _ in range(16))
261130Szelenkov@nginx.com        return base64.b64encode(raw_key).decode()
271130Szelenkov@nginx.com
281130Szelenkov@nginx.com    def accept(self, key):
291130Szelenkov@nginx.com        sha1 = hashlib.sha1((key + GUID).encode()).digest()
301130Szelenkov@nginx.com        return base64.b64encode(sha1).decode()
311130Szelenkov@nginx.com
321262Szelenkov@nginx.com    def upgrade(self, headers=None):
331262Szelenkov@nginx.com        key = None
341262Szelenkov@nginx.com
351262Szelenkov@nginx.com        if headers is None:
361262Szelenkov@nginx.com            key = self.key()
371262Szelenkov@nginx.com            headers = {
381130Szelenkov@nginx.com                'Host': 'localhost',
391130Szelenkov@nginx.com                'Upgrade': 'websocket',
401130Szelenkov@nginx.com                'Connection': 'Upgrade',
411130Szelenkov@nginx.com                'Sec-WebSocket-Key': key,
421625Smax.romanov@nginx.com                'Sec-WebSocket-Protocol': 'chat, phone, video',
431130Szelenkov@nginx.com                'Sec-WebSocket-Version': 13,
441262Szelenkov@nginx.com            }
451262Szelenkov@nginx.com
46*2073Szelenkov@nginx.com        _, sock = self.get(
47*2073Szelenkov@nginx.com            headers=headers,
48*2073Szelenkov@nginx.com            no_recv=True,
49*2073Szelenkov@nginx.com            start=True,
50*2073Szelenkov@nginx.com        )
511130Szelenkov@nginx.com
521246Szelenkov@nginx.com        resp = ''
531444Szelenkov@nginx.com        while True:
541444Szelenkov@nginx.com            rlist = select.select([sock], [], [], 60)[0]
551444Szelenkov@nginx.com            if not rlist:
561596Szelenkov@nginx.com                pytest.fail('Can\'t read response from server.')
571444Szelenkov@nginx.com
581246Szelenkov@nginx.com            resp += sock.recv(4096).decode()
591246Szelenkov@nginx.com
601848Szelenkov@nginx.com            if resp.startswith('HTTP/') and '\r\n\r\n' in resp:
611246Szelenkov@nginx.com                resp = self._resp_to_dict(resp)
621246Szelenkov@nginx.com                break
631246Szelenkov@nginx.com
641130Szelenkov@nginx.com        return (resp, sock, key)
651130Szelenkov@nginx.com
661130Szelenkov@nginx.com    def apply_mask(self, data, mask):
671130Szelenkov@nginx.com        return bytes(b ^ m for b, m in zip(data, itertools.cycle(mask)))
681130Szelenkov@nginx.com
691150Szelenkov@nginx.com    def serialize_close(self, code=1000, reason=''):
701130Szelenkov@nginx.com        return struct.pack('!H', code) + reason.encode('utf-8')
711130Szelenkov@nginx.com
721424Szelenkov@nginx.com    def frame_read(self, sock, read_timeout=60):
731130Szelenkov@nginx.com        def recv_bytes(sock, bytes):
741130Szelenkov@nginx.com            data = b''
751444Szelenkov@nginx.com            while True:
761444Szelenkov@nginx.com                rlist = select.select([sock], [], [], read_timeout)[0]
771444Szelenkov@nginx.com                if not rlist:
781444Szelenkov@nginx.com                    # For all current cases if the "read_timeout" was changed
791444Szelenkov@nginx.com                    # than test do not expect to get a response from server.
801444Szelenkov@nginx.com                    if read_timeout == 60:
811596Szelenkov@nginx.com                        pytest.fail('Can\'t read response from server.')
821444Szelenkov@nginx.com                    break
831444Szelenkov@nginx.com
841154Szelenkov@nginx.com                data += sock.recv(bytes - len(data))
851154Szelenkov@nginx.com
861154Szelenkov@nginx.com                if len(data) == bytes:
871130Szelenkov@nginx.com                    break
881130Szelenkov@nginx.com
891130Szelenkov@nginx.com            return data
901130Szelenkov@nginx.com
911130Szelenkov@nginx.com        frame = {}
921130Szelenkov@nginx.com
931848Szelenkov@nginx.com        (head1,) = struct.unpack('!B', recv_bytes(sock, 1))
941848Szelenkov@nginx.com        (head2,) = struct.unpack('!B', recv_bytes(sock, 1))
951130Szelenkov@nginx.com
961130Szelenkov@nginx.com        frame['fin'] = bool(head1 & 0b10000000)
971130Szelenkov@nginx.com        frame['rsv1'] = bool(head1 & 0b01000000)
981130Szelenkov@nginx.com        frame['rsv2'] = bool(head1 & 0b00100000)
991130Szelenkov@nginx.com        frame['rsv3'] = bool(head1 & 0b00010000)
1001130Szelenkov@nginx.com        frame['opcode'] = head1 & 0b00001111
1011130Szelenkov@nginx.com        frame['mask'] = head2 & 0b10000000
1021130Szelenkov@nginx.com
1031130Szelenkov@nginx.com        length = head2 & 0b01111111
1041130Szelenkov@nginx.com        if length == 126:
1051130Szelenkov@nginx.com            data = recv_bytes(sock, 2)
1061848Szelenkov@nginx.com            (length,) = struct.unpack('!H', data)
1071130Szelenkov@nginx.com        elif length == 127:
1081130Szelenkov@nginx.com            data = recv_bytes(sock, 8)
1091848Szelenkov@nginx.com            (length,) = struct.unpack('!Q', data)
1101130Szelenkov@nginx.com
1111130Szelenkov@nginx.com        if frame['mask']:
1121130Szelenkov@nginx.com            mask_bits = recv_bytes(sock, 4)
1131130Szelenkov@nginx.com
1141156Szelenkov@nginx.com        data = b''
1151156Szelenkov@nginx.com
1161156Szelenkov@nginx.com        if length != 0:
1171156Szelenkov@nginx.com            data = recv_bytes(sock, length)
1181156Szelenkov@nginx.com
1191130Szelenkov@nginx.com        if frame['mask']:
1201130Szelenkov@nginx.com            data = self.apply_mask(data, mask_bits)
1211130Szelenkov@nginx.com
1221130Szelenkov@nginx.com        if frame['opcode'] == self.OP_CLOSE:
1231130Szelenkov@nginx.com            if length >= 2:
1241848Szelenkov@nginx.com                (code,) = struct.unpack('!H', data[:2])
1251130Szelenkov@nginx.com                reason = data[2:].decode('utf-8')
1261130Szelenkov@nginx.com                if not (code in self.CLOSE_CODES or 3000 <= code < 5000):
1271596Szelenkov@nginx.com                    pytest.fail('Invalid status code')
1281130Szelenkov@nginx.com                frame['code'] = code
1291130Szelenkov@nginx.com                frame['reason'] = reason
1301130Szelenkov@nginx.com            elif length == 0:
1311130Szelenkov@nginx.com                frame['code'] = 1005
1321130Szelenkov@nginx.com                frame['reason'] = ''
1331130Szelenkov@nginx.com            else:
1341596Szelenkov@nginx.com                pytest.fail('Close frame too short')
1351130Szelenkov@nginx.com
1361130Szelenkov@nginx.com        frame['data'] = data
1371130Szelenkov@nginx.com
1381130Szelenkov@nginx.com        if frame['mask']:
1391596Szelenkov@nginx.com            pytest.fail('Received frame with mask')
1401130Szelenkov@nginx.com
1411130Szelenkov@nginx.com        return frame
1421130Szelenkov@nginx.com
1431130Szelenkov@nginx.com    def frame_to_send(
1441130Szelenkov@nginx.com        self,
1451130Szelenkov@nginx.com        opcode,
1461130Szelenkov@nginx.com        data,
1471130Szelenkov@nginx.com        fin=True,
1481130Szelenkov@nginx.com        length=None,
1491130Szelenkov@nginx.com        rsv1=False,
1501130Szelenkov@nginx.com        rsv2=False,
1511130Szelenkov@nginx.com        rsv3=False,
1521130Szelenkov@nginx.com        mask=True,
1531130Szelenkov@nginx.com    ):
1541130Szelenkov@nginx.com        frame = b''
1551130Szelenkov@nginx.com
1561130Szelenkov@nginx.com        if isinstance(data, str):
1571130Szelenkov@nginx.com            data = data.encode('utf-8')
1581130Szelenkov@nginx.com
1591130Szelenkov@nginx.com        head1 = (
1601130Szelenkov@nginx.com            (0b10000000 if fin else 0)
1611130Szelenkov@nginx.com            | (0b01000000 if rsv1 else 0)
1621130Szelenkov@nginx.com            | (0b00100000 if rsv2 else 0)
1631130Szelenkov@nginx.com            | (0b00010000 if rsv3 else 0)
1641130Szelenkov@nginx.com            | opcode
1651130Szelenkov@nginx.com        )
1661130Szelenkov@nginx.com
1671130Szelenkov@nginx.com        head2 = 0b10000000 if mask else 0
1681130Szelenkov@nginx.com
1691130Szelenkov@nginx.com        data_length = len(data) if length is None else length
1701130Szelenkov@nginx.com        if data_length < 126:
1711130Szelenkov@nginx.com            frame += struct.pack('!BB', head1, head2 | data_length)
1721130Szelenkov@nginx.com        elif data_length < 65536:
1731130Szelenkov@nginx.com            frame += struct.pack('!BBH', head1, head2 | 126, data_length)
1741130Szelenkov@nginx.com        else:
1751130Szelenkov@nginx.com            frame += struct.pack('!BBQ', head1, head2 | 127, data_length)
1761130Szelenkov@nginx.com
1771130Szelenkov@nginx.com        if mask:
1781130Szelenkov@nginx.com            mask_bits = struct.pack('!I', random.getrandbits(32))
1791130Szelenkov@nginx.com            frame += mask_bits
1801130Szelenkov@nginx.com
1811130Szelenkov@nginx.com        if mask:
1821130Szelenkov@nginx.com            frame += self.apply_mask(data, mask_bits)
1831130Szelenkov@nginx.com        else:
1841130Szelenkov@nginx.com            frame += data
1851130Szelenkov@nginx.com
1861130Szelenkov@nginx.com        return frame
1871130Szelenkov@nginx.com
1881130Szelenkov@nginx.com    def frame_write(self, sock, *args, **kwargs):
1891130Szelenkov@nginx.com        chopsize = kwargs.pop('chopsize') if 'chopsize' in kwargs else None
1901130Szelenkov@nginx.com
1911130Szelenkov@nginx.com        frame = self.frame_to_send(*args, **kwargs)
1921130Szelenkov@nginx.com
1931130Szelenkov@nginx.com        if chopsize is None:
1941152Szelenkov@nginx.com            try:
1951152Szelenkov@nginx.com                sock.sendall(frame)
1961152Szelenkov@nginx.com            except BrokenPipeError:
1971152Szelenkov@nginx.com                pass
1981130Szelenkov@nginx.com
1991130Szelenkov@nginx.com        else:
2001130Szelenkov@nginx.com            pos = 0
2011130Szelenkov@nginx.com            frame_len = len(frame)
2021150Szelenkov@nginx.com            while pos < frame_len:
2031130Szelenkov@nginx.com                end = min(pos + chopsize, frame_len)
2041151Szelenkov@nginx.com                try:
2051151Szelenkov@nginx.com                    sock.sendall(frame[pos:end])
2061151Szelenkov@nginx.com                except BrokenPipeError:
2071151Szelenkov@nginx.com                    end = frame_len
2081130Szelenkov@nginx.com                pos = end
2091130Szelenkov@nginx.com
2101130Szelenkov@nginx.com    def message(self, sock, type, message, fragmention_size=None, **kwargs):
2111130Szelenkov@nginx.com        message_len = len(message)
2121130Szelenkov@nginx.com
2131130Szelenkov@nginx.com        if fragmention_size is None:
2141130Szelenkov@nginx.com            fragmention_size = message_len
2151130Szelenkov@nginx.com
2161130Szelenkov@nginx.com        if message_len <= fragmention_size:
2171130Szelenkov@nginx.com            self.frame_write(sock, type, message, **kwargs)
2181130Szelenkov@nginx.com            return
2191130Szelenkov@nginx.com
2201130Szelenkov@nginx.com        pos = 0
2211130Szelenkov@nginx.com        op_code = type
2221150Szelenkov@nginx.com        while pos < message_len:
2231130Szelenkov@nginx.com            end = min(pos + fragmention_size, message_len)
2241150Szelenkov@nginx.com            fin = end == message_len
225*2073Szelenkov@nginx.com            self.frame_write(sock, op_code, message[pos:end], fin=fin, **kwargs)
2261130Szelenkov@nginx.com            op_code = self.OP_CONT
2271130Szelenkov@nginx.com            pos = end
2281130Szelenkov@nginx.com
2291433Szelenkov@nginx.com    def message_read(self, sock, read_timeout=60):
2301130Szelenkov@nginx.com        frame = self.frame_read(sock, read_timeout=read_timeout)
2311130Szelenkov@nginx.com
2321150Szelenkov@nginx.com        while not frame['fin']:
2331130Szelenkov@nginx.com            temp = self.frame_read(sock, read_timeout=read_timeout)
2341130Szelenkov@nginx.com            frame['data'] += temp['data']
2351130Szelenkov@nginx.com            frame['fin'] = temp['fin']
2361130Szelenkov@nginx.com
2371130Szelenkov@nginx.com        return frame
238