296 lines
8.7 KiB
Python
296 lines
8.7 KiB
Python
|
"""
|
||
|
|
The :mod:`websockets.framing` module implements data framing as specified in
|
||
|
|
`section 5 of RFC 6455`_.
|
||
|
|
|
||
|
|
It deals with a single frame at a time. Anything that depends on the sequence
|
||
|
|
of frames is implemented in :mod:`websockets.protocol`.
|
||
|
|
|
||
|
|
.. _section 5 of RFC 6455: http://tools.ietf.org/html/rfc6455#section-5
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import collections
|
||
|
|
import io
|
||
|
|
import random
|
||
|
|
import struct
|
||
|
|
|
||
|
|
from .exceptions import PayloadTooBig, WebSocketProtocolError
|
||
|
|
|
||
|
|
|
||
|
|
try:
|
||
|
|
from .speedups import apply_mask
|
||
|
|
except ImportError: # pragma: no cover
|
||
|
|
from .utils import apply_mask
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
'DATA_OPCODES', 'CTRL_OPCODES',
|
||
|
|
'OP_CONT', 'OP_TEXT', 'OP_BINARY', 'OP_CLOSE', 'OP_PING', 'OP_PONG',
|
||
|
|
'Frame', 'encode_data', 'parse_close', 'serialize_close'
|
||
|
|
]
|
||
|
|
|
||
|
|
DATA_OPCODES = OP_CONT, OP_TEXT, OP_BINARY = 0x00, 0x01, 0x02
|
||
|
|
CTRL_OPCODES = OP_CLOSE, OP_PING, OP_PONG = 0x08, 0x09, 0x0a
|
||
|
|
|
||
|
|
# Close code that are allowed in a close frame.
|
||
|
|
# Using a list optimizes `code in EXTERNAL_CLOSE_CODES`.
|
||
|
|
EXTERNAL_CLOSE_CODES = [
|
||
|
|
1000,
|
||
|
|
1001,
|
||
|
|
1002,
|
||
|
|
1003,
|
||
|
|
1007,
|
||
|
|
1008,
|
||
|
|
1009,
|
||
|
|
1010,
|
||
|
|
1011,
|
||
|
|
]
|
||
|
|
|
||
|
|
FrameData = collections.namedtuple(
|
||
|
|
'FrameData',
|
||
|
|
['fin', 'opcode', 'data', 'rsv1', 'rsv2', 'rsv3'],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class Frame(FrameData):
|
||
|
|
"""
|
||
|
|
WebSocket frame.
|
||
|
|
|
||
|
|
* ``fin`` is the FIN bit
|
||
|
|
* ``rsv1`` is the RSV1 bit
|
||
|
|
* ``rsv2`` is the RSV2 bit
|
||
|
|
* ``rsv3`` is the RSV3 bit
|
||
|
|
* ``opcode`` is the opcode
|
||
|
|
* ``data`` is the payload data
|
||
|
|
|
||
|
|
Only these fields are needed by higher level code. The MASK bit, payload
|
||
|
|
length and masking-key are handled on the fly by :meth:`read` and
|
||
|
|
:meth:`write`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
def __new__(cls, fin, opcode, data, rsv1=False, rsv2=False, rsv3=False):
|
||
|
|
return FrameData.__new__(cls, fin, opcode, data, rsv1, rsv2, rsv3)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
@asyncio.coroutine
|
||
|
|
def read(cls, reader, *, mask, max_size=None, extensions=None):
|
||
|
|
"""
|
||
|
|
Read a WebSocket frame and return a :class:`Frame` object.
|
||
|
|
|
||
|
|
``reader`` is a coroutine taking an integer argument and reading
|
||
|
|
exactly this number of bytes, unless the end of file is reached.
|
||
|
|
|
||
|
|
``mask`` is a :class:`bool` telling whether the frame should be masked
|
||
|
|
i.e. whether the read happens on the server side.
|
||
|
|
|
||
|
|
If ``max_size`` is set and the payload exceeds this size in bytes,
|
||
|
|
:exc:`~websockets.exceptions.PayloadTooBig` is raised.
|
||
|
|
|
||
|
|
If ``extensions`` is provided, it's a list of classes with an
|
||
|
|
``decode()`` method that transform the frame and return a new frame.
|
||
|
|
They are applied in reverse order.
|
||
|
|
|
||
|
|
This function validates the frame before returning it and raises
|
||
|
|
:exc:`~websockets.exceptions.WebSocketProtocolError` if it contains
|
||
|
|
incorrect values.
|
||
|
|
|
||
|
|
"""
|
||
|
|
# Read the header.
|
||
|
|
data = yield from reader(2)
|
||
|
|
head1, head2 = struct.unpack('!BB', data)
|
||
|
|
|
||
|
|
# While not Pythonic, this is marginally faster than calling bool().
|
||
|
|
fin = True if head1 & 0b10000000 else False
|
||
|
|
rsv1 = True if head1 & 0b01000000 else False
|
||
|
|
rsv2 = True if head1 & 0b00100000 else False
|
||
|
|
rsv3 = True if head1 & 0b00010000 else False
|
||
|
|
opcode = head1 & 0b00001111
|
||
|
|
|
||
|
|
if (True if head2 & 0b10000000 else False) != mask:
|
||
|
|
raise WebSocketProtocolError("Incorrect masking")
|
||
|
|
|
||
|
|
length = head2 & 0b01111111
|
||
|
|
if length == 126:
|
||
|
|
data = yield from reader(2)
|
||
|
|
length, = struct.unpack('!H', data)
|
||
|
|
elif length == 127:
|
||
|
|
data = yield from reader(8)
|
||
|
|
length, = struct.unpack('!Q', data)
|
||
|
|
if max_size is not None and length > max_size:
|
||
|
|
raise PayloadTooBig(
|
||
|
|
"Payload length exceeds size limit ({} > {} bytes)"
|
||
|
|
.format(length, max_size))
|
||
|
|
if mask:
|
||
|
|
mask_bits = yield from reader(4)
|
||
|
|
|
||
|
|
# Read the data.
|
||
|
|
data = yield from reader(length)
|
||
|
|
if mask:
|
||
|
|
data = apply_mask(data, mask_bits)
|
||
|
|
|
||
|
|
frame = cls(fin, opcode, data, rsv1, rsv2, rsv3)
|
||
|
|
|
||
|
|
if extensions is None:
|
||
|
|
extensions = []
|
||
|
|
for extension in reversed(extensions):
|
||
|
|
frame = extension.decode(frame, max_size=max_size)
|
||
|
|
|
||
|
|
frame.check()
|
||
|
|
|
||
|
|
return frame
|
||
|
|
|
||
|
|
def write(frame, writer, *, mask, extensions=None):
|
||
|
|
"""
|
||
|
|
Write a WebSocket frame.
|
||
|
|
|
||
|
|
``frame`` is the :class:`Frame` object to write.
|
||
|
|
|
||
|
|
``writer`` is a function accepting bytes.
|
||
|
|
|
||
|
|
``mask`` is a :class:`bool` telling whether the frame should be masked
|
||
|
|
i.e. whether the write happens on the client side.
|
||
|
|
|
||
|
|
If ``extensions`` is provided, it's a list of classes with an
|
||
|
|
``encode()`` method that transform the frame and return a new frame.
|
||
|
|
They are applied in order.
|
||
|
|
|
||
|
|
This function validates the frame before sending it and raises
|
||
|
|
:exc:`~websockets.exceptions.WebSocketProtocolError` if it contains
|
||
|
|
incorrect values.
|
||
|
|
|
||
|
|
"""
|
||
|
|
# The first parameter is called `frame` rather than `self`,
|
||
|
|
# but it's the instance of class to which this method is bound.
|
||
|
|
|
||
|
|
frame.check()
|
||
|
|
|
||
|
|
if extensions is None:
|
||
|
|
extensions = []
|
||
|
|
for extension in extensions:
|
||
|
|
frame = extension.encode(frame)
|
||
|
|
|
||
|
|
output = io.BytesIO()
|
||
|
|
|
||
|
|
# Prepare the header.
|
||
|
|
head1 = (
|
||
|
|
(0b10000000 if frame.fin else 0) |
|
||
|
|
(0b01000000 if frame.rsv1 else 0) |
|
||
|
|
(0b00100000 if frame.rsv2 else 0) |
|
||
|
|
(0b00010000 if frame.rsv3 else 0) |
|
||
|
|
frame.opcode
|
||
|
|
)
|
||
|
|
|
||
|
|
head2 = 0b10000000 if mask else 0
|
||
|
|
|
||
|
|
length = len(frame.data)
|
||
|
|
if length < 126:
|
||
|
|
output.write(struct.pack('!BB', head1, head2 | length))
|
||
|
|
elif length < 65536:
|
||
|
|
output.write(struct.pack('!BBH', head1, head2 | 126, length))
|
||
|
|
else:
|
||
|
|
output.write(struct.pack('!BBQ', head1, head2 | 127, length))
|
||
|
|
|
||
|
|
if mask:
|
||
|
|
mask_bits = struct.pack('!I', random.getrandbits(32))
|
||
|
|
output.write(mask_bits)
|
||
|
|
|
||
|
|
# Prepare the data.
|
||
|
|
if mask:
|
||
|
|
data = apply_mask(frame.data, mask_bits)
|
||
|
|
else:
|
||
|
|
data = frame.data
|
||
|
|
output.write(data)
|
||
|
|
|
||
|
|
# Send the frame.
|
||
|
|
|
||
|
|
# The frame is written in a single call to writer in order to prevent
|
||
|
|
# TCP fragmentation. See #68 for details. This also makes it safe to
|
||
|
|
# send frames concurrently from multiple coroutines.
|
||
|
|
writer(output.getvalue())
|
||
|
|
|
||
|
|
def check(frame):
|
||
|
|
"""
|
||
|
|
Check that this frame contains acceptable values.
|
||
|
|
|
||
|
|
Raise :exc:`~websockets.exceptions.WebSocketProtocolError` if this
|
||
|
|
frame contains incorrect values.
|
||
|
|
|
||
|
|
"""
|
||
|
|
# The first parameter is called `frame` rather than `self`,
|
||
|
|
# but it's the instance of class to which this method is bound.
|
||
|
|
|
||
|
|
if frame.rsv1 or frame.rsv2 or frame.rsv3:
|
||
|
|
raise WebSocketProtocolError("Reserved bits must be 0")
|
||
|
|
|
||
|
|
if frame.opcode in DATA_OPCODES:
|
||
|
|
return
|
||
|
|
elif frame.opcode in CTRL_OPCODES:
|
||
|
|
if len(frame.data) > 125:
|
||
|
|
raise WebSocketProtocolError("Control frame too long")
|
||
|
|
if not frame.fin:
|
||
|
|
raise WebSocketProtocolError("Fragmented control frame")
|
||
|
|
else:
|
||
|
|
raise WebSocketProtocolError(
|
||
|
|
"Invalid opcode: {}".format(frame.opcode))
|
||
|
|
|
||
|
|
|
||
|
|
def encode_data(data):
|
||
|
|
"""
|
||
|
|
Helper that converts :class:`str` or :class:`bytes` to :class:`bytes`.
|
||
|
|
|
||
|
|
:class:`str` are encoded with UTF-8.
|
||
|
|
|
||
|
|
"""
|
||
|
|
# Expect str or bytes, return bytes.
|
||
|
|
if isinstance(data, str):
|
||
|
|
return data.encode('utf-8')
|
||
|
|
elif isinstance(data, bytes):
|
||
|
|
return data
|
||
|
|
else:
|
||
|
|
raise TypeError("data must be bytes or str")
|
||
|
|
|
||
|
|
|
||
|
|
def parse_close(data):
|
||
|
|
"""
|
||
|
|
Parse the data in a close frame.
|
||
|
|
|
||
|
|
Return ``(code, reason)`` when ``code`` is an :class:`int` and ``reason``
|
||
|
|
a :class:`str`.
|
||
|
|
|
||
|
|
Raise :exc:`~websockets.exceptions.WebSocketProtocolError` or
|
||
|
|
:exc:`UnicodeDecodeError` if the data is invalid.
|
||
|
|
|
||
|
|
"""
|
||
|
|
length = len(data)
|
||
|
|
if length >= 2:
|
||
|
|
code, = struct.unpack('!H', data[:2])
|
||
|
|
check_close(code)
|
||
|
|
reason = data[2:].decode('utf-8')
|
||
|
|
return code, reason
|
||
|
|
elif length == 0:
|
||
|
|
return 1005, ''
|
||
|
|
else:
|
||
|
|
assert length == 1
|
||
|
|
raise WebSocketProtocolError("Close frame too short")
|
||
|
|
|
||
|
|
|
||
|
|
def serialize_close(code, reason):
|
||
|
|
"""
|
||
|
|
Serialize the data for a close frame.
|
||
|
|
|
||
|
|
This is the reverse of :func:`parse_close`.
|
||
|
|
|
||
|
|
"""
|
||
|
|
check_close(code)
|
||
|
|
return struct.pack('!H', code) + reason.encode('utf-8')
|
||
|
|
|
||
|
|
|
||
|
|
def check_close(code):
|
||
|
|
"""
|
||
|
|
Check the close code for a close frame.
|
||
|
|
|
||
|
|
"""
|
||
|
|
if not (code in EXTERNAL_CLOSE_CODES or 3000 <= code < 5000):
|
||
|
|
raise WebSocketProtocolError("Invalid status code")
|