"""
h2/frame_buffer
~~~~~~~~~~~~~~~
A data structure that provides a way to iterate over a byte buffer in terms of
frames.
"""
from __future__ import annotations
from hyperframe.exceptions import InvalidDataError, InvalidFrameError
from hyperframe.frame import ContinuationFrame, Frame, HeadersFrame, PushPromiseFrame
from .exceptions import FrameDataMissingError, FrameTooLargeError, ProtocolError
# To avoid a DOS attack based on sending loads of continuation frames, we limit
# the maximum number we're perpared to receive. In this case, we'll set the
# limit to 64, which means the largest encoded header block we can receive by
# default is 262144 bytes long, and the largest possible *at all* is 1073741760
# bytes long.
#
# This value seems reasonable for now, but in future we may want to evaluate
# making it configurable.
CONTINUATION_BACKLOG = 64
class FrameBuffer:
"""
A buffer data structure for HTTP/2 data that allows iteraton in terms of
H2 frames.
"""
def __init__(self, server: bool = False) -> None:
self.data = b""
self.max_frame_size = 0
self._preamble = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n" if server else b""
self._preamble_len = len(self._preamble)
self._headers_buffer: list[HeadersFrame | ContinuationFrame | PushPromiseFrame] = []
def add_data(self, data: bytes) -> None:
"""
Add more data to the frame buffer.
:param data: A bytestring containing the byte buffer.
"""
if self._preamble_len:
data_len = len(data)
of_which_preamble = min(self._preamble_len, data_len)
if self._preamble[:of_which_preamble] != data[:of_which_preamble]:
msg = "Invalid HTTP/2 preamble."
raise ProtocolError(msg)
data = data[of_which_preamble:]
self._preamble_len -= of_which_preamble
self._preamble = self._preamble[of_which_preamble:]
self.data += data
def _validate_frame_length(self, length: int) -> None:
"""
Confirm that the frame is an appropriate length.
"""
if length > self.max_frame_size:
msg = f"Received overlong frame: length {length}, max {self.max_frame_size}"
raise FrameTooLargeError(msg)
def _update_header_buffer(self, f: Frame | None) -> Frame | None:
"""
Updates the internal header buffer. Returns a frame that should replace
the current one. May throw exceptions if this frame is invalid.
"""
# Check if we're in the middle of a headers block. If we are, this
# frame *must* be a CONTINUATION frame with the same stream ID as the
# leading HEADERS or PUSH_PROMISE frame. Anything else is a
# ProtocolError. If the frame *is* valid, append it to the header
# buffer.
if self._headers_buffer:
stream_id = self._headers_buffer[0].stream_id
valid_frame = (
f is not None and
isinstance(f, ContinuationFrame) and
f.stream_id == stream_id
)
if not valid_frame:
msg = "Invalid frame during header block."
raise ProtocolError(msg)
assert isinstance(f, ContinuationFrame)
# Append the frame to the buffer.
self._headers_buffer.append(f)
if len(self._headers_buffer) > CONTINUATION_BACKLOG:
msg = "Too many continuation frames received."
raise ProtocolError(msg)
# If this is the end of the header block, then we want to build a
# mutant HEADERS frame that's massive. Use the original one we got,
# then set END_HEADERS and set its data appopriately. If it's not
# the end of the block, lose the current frame: we can't yield it.
if "END_HEADERS" in f.flags:
f = self._headers_buffer[0]
f.flags.add("END_HEADERS")
f.data = b"".join(x.data for x in self._headers_buffer)
self._headers_buffer = []
else:
f = None
elif (isinstance(f, (HeadersFrame, PushPromiseFrame)) and
"END_HEADERS" not in f.flags):
# This is the start of a headers block! Save the frame off and then
# act like we didn't receive one.
self._headers_buffer.append(f)
f = None
return f
# The methods below support the iterator protocol.
def __iter__(self) -> FrameBuffer:
return self
def __next__(self) -> Frame:
# First, check that we have enough data to successfully parse the
# next frame header. If not, bail. Otherwise, parse it.
if len(self.data) < 9:
raise StopIteration
try:
f, length = Frame.parse_frame_header(memoryview(self.data[:9]))
except (InvalidDataError, InvalidFrameError) as err: # pragma: no cover
msg = f"Received frame with invalid header: {err!s}"
raise ProtocolError(msg) from err
# Next, check that we have enough length to parse the frame body. If
# not, bail, leaving the frame header data in the buffer for next time.
if len(self.data) < length + 9:
raise StopIteration
# Confirm the frame has an appropriate length.
self._validate_frame_length(length)
# Try to parse the frame body
try:
f.parse_body(memoryview(self.data[9:9+length]))
except InvalidDataError as err:
msg = "Received frame with non-compliant data"
raise ProtocolError(msg) from err
except InvalidFrameError as err:
msg = "Frame data missing or invalid"
raise FrameDataMissingError(msg) from err
# At this point, as we know we'll use or discard the entire frame, we
# can update the data.
self.data = self.data[9+length:]
# Pass the frame through the header buffer.
new_frame = self._update_header_buffer(f)
# If we got a frame we didn't understand or shouldn't yield, rather
# than return None it'd be better if we just tried to get the next
# frame in the sequence instead. Recurse back into ourselves to do
# that. This is safe because the amount of work we have to do here is
# strictly bounded by the length of the buffer.
return new_frame if new_frame is not None else self.__next__()
