"""WebSocket protocol versions 13 and 8.""" import base64 import binascii import collections import hashlib import json import os import random import sys from enum import IntEnum from struct import Struct from aiohttp import errors, hdrs from aiohttp.log import ws_logger __all__ = ('WebSocketParser', 'WebSocketWriter', 'do_handshake', 'WSMessage', 'WebSocketError', 'WSMsgType', 'WSCloseCode') class WSCloseCode(IntEnum): OK = 1000 GOING_AWAY = 1001 PROTOCOL_ERROR = 1002 UNSUPPORTED_DATA = 1003 INVALID_TEXT = 1007 POLICY_VIOLATION = 1008 MESSAGE_TOO_BIG = 1009 MANDATORY_EXTENSION = 1010 INTERNAL_ERROR = 1011 SERVICE_RESTART = 1012 TRY_AGAIN_LATER = 1013 ALLOWED_CLOSE_CODES = {int(i) for i in WSCloseCode} class WSMsgType(IntEnum): CONTINUATION = 0x0 TEXT = 0x1 BINARY = 0x2 PING = 0x9 PONG = 0xa CLOSE = 0x8 CLOSED = 0x101 ERROR = 0x102 text = TEXT binary = BINARY ping = PING pong = PONG close = CLOSE closed = CLOSED error = ERROR WS_KEY = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11' UNPACK_LEN2 = Struct('!H').unpack_from UNPACK_LEN3 = Struct('!Q').unpack_from UNPACK_CLOSE_CODE = Struct('!H').unpack PACK_LEN1 = Struct('!BB').pack PACK_LEN2 = Struct('!BBH').pack PACK_LEN3 = Struct('!BBQ').pack PACK_CLOSE_CODE = Struct('!H').pack MSG_SIZE = 2 ** 14 _WSMessageBase = collections.namedtuple('_WSMessageBase', ['type', 'data', 'extra']) class WSMessage(_WSMessageBase): def json(self, *, loads=json.loads): """Return parsed JSON data. .. versionadded:: 0.22 """ return loads(self.data) @property def tp(self): return self.type CLOSED_MESSAGE = WSMessage(WSMsgType.CLOSED, None, None) class WebSocketError(Exception): """WebSocket protocol parser error.""" def __init__(self, code, message): self.code = code super().__init__(message) def WebSocketParser(out, buf): while True: fin, opcode, payload = yield from parse_frame(buf) if opcode == WSMsgType.CLOSE: if len(payload) >= 2: close_code = UNPACK_CLOSE_CODE(payload[:2])[0] if close_code < 3000 and close_code not in ALLOWED_CLOSE_CODES: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Invalid close code: {}'.format(close_code)) try: close_message = payload[2:].decode('utf-8') except UnicodeDecodeError as exc: raise WebSocketError( WSCloseCode.INVALID_TEXT, 'Invalid UTF-8 text message') from exc msg = WSMessage(WSMsgType.CLOSE, close_code, close_message) elif payload: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Invalid close frame: {} {} {!r}'.format( fin, opcode, payload)) else: msg = WSMessage(WSMsgType.CLOSE, 0, '') out.feed_data(msg, 0) elif opcode == WSMsgType.PING: out.feed_data(WSMessage(WSMsgType.PING, payload, ''), len(payload)) elif opcode == WSMsgType.PONG: out.feed_data(WSMessage(WSMsgType.PONG, payload, ''), len(payload)) elif opcode not in (WSMsgType.TEXT, WSMsgType.BINARY): raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Unexpected opcode={!r}".format(opcode)) else: # load text/binary data = [payload] while not fin: fin, _opcode, payload = yield from parse_frame(buf, True) # We can receive ping/close in the middle of # text message, Case 5.* if _opcode == WSMsgType.PING: out.feed_data( WSMessage(WSMsgType.PING, payload, ''), len(payload)) fin, _opcode, payload = yield from parse_frame(buf, True) elif _opcode == WSMsgType.CLOSE: if len(payload) >= 2: close_code = UNPACK_CLOSE_CODE(payload[:2])[0] if (close_code not in ALLOWED_CLOSE_CODES and close_code < 3000): raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Invalid close code: {}'.format(close_code)) try: close_message = payload[2:].decode('utf-8') except UnicodeDecodeError as exc: raise WebSocketError( WSCloseCode.INVALID_TEXT, 'Invalid UTF-8 text message') from exc msg = WSMessage(WSMsgType.CLOSE, close_code, close_message) elif payload: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Invalid close frame: {} {} {!r}'.format( fin, opcode, payload)) else: msg = WSMessage(WSMsgType.CLOSE, 0, '') out.feed_data(msg, 0) fin, _opcode, payload = yield from parse_frame(buf, True) if _opcode != WSMsgType.CONTINUATION: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'The opcode in non-fin frame is expected ' 'to be zero, got {!r}'.format(_opcode)) else: data.append(payload) if opcode == WSMsgType.TEXT: try: text = b''.join(data).decode('utf-8') out.feed_data(WSMessage(WSMsgType.TEXT, text, ''), len(text)) except UnicodeDecodeError as exc: raise WebSocketError( WSCloseCode.INVALID_TEXT, 'Invalid UTF-8 text message') from exc else: data = b''.join(data) out.feed_data( WSMessage(WSMsgType.BINARY, data, ''), len(data)) native_byteorder = sys.byteorder def _websocket_mask_python(mask, data): """Websocket masking function. `mask` is a `bytes` object of length 4; `data` is a `bytes` object of any length. Returns a `bytes` object of the same length as `data` with the mask applied as specified in section 5.3 of RFC 6455. This pure-python implementation may be replaced by an optimized version when available. """ assert isinstance(data, bytearray), data assert len(mask) == 4, mask datalen = len(data) if datalen == 0: # everything work without this, but may be changed later in Python. return bytearray() data = int.from_bytes(data, native_byteorder) mask = int.from_bytes(mask * (datalen // 4) + mask[: datalen % 4], native_byteorder) return (data ^ mask).to_bytes(datalen, native_byteorder) if bool(os.environ.get('AIOHTTP_NO_EXTENSIONS')): _websocket_mask = _websocket_mask_python else: try: from ._websocket import _websocket_mask_cython _websocket_mask = _websocket_mask_cython except ImportError: # pragma: no cover _websocket_mask = _websocket_mask_python def parse_frame(buf, continuation=False): """Return the next frame from the socket.""" # read header data = yield from buf.read(2) first_byte, second_byte = data fin = (first_byte >> 7) & 1 rsv1 = (first_byte >> 6) & 1 rsv2 = (first_byte >> 5) & 1 rsv3 = (first_byte >> 4) & 1 opcode = first_byte & 0xf # frame-fin = %x0 ; more frames of this message follow # / %x1 ; final frame of this message # frame-rsv1 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise # frame-rsv2 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise # frame-rsv3 = %x0 ; 1 bit, MUST be 0 unless negotiated otherwise if rsv1 or rsv2 or rsv3: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Received frame with non-zero reserved bits') if opcode > 0x7 and fin == 0: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Received fragmented control frame') if fin == 0 and opcode == WSMsgType.CONTINUATION and not continuation: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, 'Received new fragment frame with non-zero ' 'opcode {!r}'.format(opcode)) has_mask = (second_byte >> 7) & 1 length = (second_byte) & 0x7f # Control frames MUST have a payload length of 125 bytes or less if opcode > 0x7 and length > 125: raise WebSocketError( WSCloseCode.PROTOCOL_ERROR, "Control frame payload cannot be larger than 125 bytes") # read payload if length == 126: data = yield from buf.read(2) length = UNPACK_LEN2(data)[0] elif length > 126: data = yield from buf.read(8) length = UNPACK_LEN3(data)[0] if has_mask: mask = yield from buf.read(4) if length: payload = yield from buf.read(length) else: payload = bytearray() if has_mask: payload = _websocket_mask(bytes(mask), payload) return fin, opcode, payload class WebSocketWriter: def __init__(self, writer, *, use_mask=False, random=random.Random()): self.writer = writer self.use_mask = use_mask self.randrange = random.randrange def _send_frame(self, message, opcode): """Send a frame over the websocket with message as its payload.""" msg_length = len(message) use_mask = self.use_mask if use_mask: mask_bit = 0x80 else: mask_bit = 0 if msg_length < 126: header = PACK_LEN1(0x80 | opcode, msg_length | mask_bit) elif msg_length < (1 << 16): header = PACK_LEN2(0x80 | opcode, 126 | mask_bit, msg_length) else: header = PACK_LEN3(0x80 | opcode, 127 | mask_bit, msg_length) if use_mask: mask = self.randrange(0, 0xffffffff) mask = mask.to_bytes(4, 'big') message = _websocket_mask(mask, bytearray(message)) self.writer.write(header + mask + message) else: if len(message) > MSG_SIZE: self.writer.write(header) self.writer.write(message) else: self.writer.write(header + message) def pong(self, message=b''): """Send pong message.""" if isinstance(message, str): message = message.encode('utf-8') self._send_frame(message, WSMsgType.PONG) def ping(self, message=b''): """Send ping message.""" if isinstance(message, str): message = message.encode('utf-8') self._send_frame(message, WSMsgType.PING) def send(self, message, binary=False): """Send a frame over the websocket with message as its payload.""" if isinstance(message, str): message = message.encode('utf-8') if binary: self._send_frame(message, WSMsgType.BINARY) else: self._send_frame(message, WSMsgType.TEXT) def close(self, code=1000, message=b''): """Close the websocket, sending the specified code and message.""" if isinstance(message, str): message = message.encode('utf-8') self._send_frame( PACK_CLOSE_CODE(code) + message, opcode=WSMsgType.CLOSE) def do_handshake(method, headers, transport, protocols=()): """Prepare WebSocket handshake. It return HTTP response code, response headers, websocket parser, websocket writer. It does not perform any IO. `protocols` is a sequence of known protocols. On successful handshake, the returned response headers contain the first protocol in this list which the server also knows. """ # WebSocket accepts only GET if method.upper() != hdrs.METH_GET: raise errors.HttpProcessingError( code=405, headers=((hdrs.ALLOW, hdrs.METH_GET),)) if 'websocket' != headers.get(hdrs.UPGRADE, '').lower().strip(): raise errors.HttpBadRequest( message='No WebSocket UPGRADE hdr: {}\n Can ' '"Upgrade" only to "WebSocket".'.format(headers.get(hdrs.UPGRADE))) if 'upgrade' not in headers.get(hdrs.CONNECTION, '').lower(): raise errors.HttpBadRequest( message='No CONNECTION upgrade hdr: {}'.format( headers.get(hdrs.CONNECTION))) # find common sub-protocol between client and server protocol = None if hdrs.SEC_WEBSOCKET_PROTOCOL in headers: req_protocols = [str(proto.strip()) for proto in headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(',')] for proto in req_protocols: if proto in protocols: protocol = proto break else: # No overlap found: Return no protocol as per spec ws_logger.warning( 'Client protocols %r don’t overlap server-known ones %r', req_protocols, protocols) # check supported version version = headers.get(hdrs.SEC_WEBSOCKET_VERSION, '') if version not in ('13', '8', '7'): raise errors.HttpBadRequest( message='Unsupported version: {}'.format(version), headers=((hdrs.SEC_WEBSOCKET_VERSION, '13'),)) # check client handshake for validity key = headers.get(hdrs.SEC_WEBSOCKET_KEY) try: if not key or len(base64.b64decode(key)) != 16: raise errors.HttpBadRequest( message='Handshake error: {!r}'.format(key)) except binascii.Error: raise errors.HttpBadRequest( message='Handshake error: {!r}'.format(key)) from None response_headers = [ (hdrs.UPGRADE, 'websocket'), (hdrs.CONNECTION, 'upgrade'), (hdrs.TRANSFER_ENCODING, 'chunked'), (hdrs.SEC_WEBSOCKET_ACCEPT, base64.b64encode( hashlib.sha1(key.encode() + WS_KEY).digest()).decode())] if protocol: response_headers.append((hdrs.SEC_WEBSOCKET_PROTOCOL, protocol)) # response code, headers, parser, writer, protocol return (101, response_headers, WebSocketParser, WebSocketWriter(transport), protocol)