codekingpro/portable-devtools

1"""2h2/frame_buffer3~~~~~~~~~~~~~~~4 5A data structure that provides a way to iterate over a byte buffer in terms of6frames.7"""8from __future__ import annotations9 10from hyperframe.exceptions import InvalidDataError, InvalidFrameError11from hyperframe.frame import ContinuationFrame, Frame, HeadersFrame, PushPromiseFrame12 13from .exceptions import FrameDataMissingError, FrameTooLargeError, ProtocolError14 15# To avoid a DOS attack based on sending loads of continuation frames, we limit16# the maximum number we're perpared to receive. In this case, we'll set the17# limit to 64, which means the largest encoded header block we can receive by18# default is 262144 bytes long, and the largest possible *at all* is 107374176019# bytes long.20#21# This value seems reasonable for now, but in future we may want to evaluate22# making it configurable.23CONTINUATION_BACKLOG = 6424 25 26class FrameBuffer:27    """28    A buffer data structure for HTTP/2 data that allows iteraton in terms of29    H2 frames.30    """31 32    def __init__(self, server: bool = False) -> None:33        self._data = bytearray()34        self.max_frame_size = 035        self._preamble = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n" if server else b""36        self._preamble_len = len(self._preamble)37        self._headers_buffer: list[HeadersFrame | ContinuationFrame | PushPromiseFrame] = []38 39    def add_data(self, data: bytes) -> None:40        """41        Add more data to the frame buffer.42 43        :param data: A bytestring containing the byte buffer.44        """45        if self._preamble_len:46            data_len = len(data)47            of_which_preamble = min(self._preamble_len, data_len)48 49            if self._preamble[:of_which_preamble] != data[:of_which_preamble]:50                msg = "Invalid HTTP/2 preamble."51                raise ProtocolError(msg)52 53            data = data[of_which_preamble:]54            self._preamble_len -= of_which_preamble55            self._preamble = self._preamble[of_which_preamble:]56 57        self._data += data58 59    def _validate_frame_length(self, length: int) -> None:60        """61        Confirm that the frame is an appropriate length.62        """63        if length > self.max_frame_size:64            msg = f"Received overlong frame: length {length}, max {self.max_frame_size}"65            raise FrameTooLargeError(msg)66 67    def _update_header_buffer(self, f: Frame | None) -> Frame | None:68        """69        Updates the internal header buffer. Returns a frame that should replace70        the current one. May throw exceptions if this frame is invalid.71        """72        # Check if we're in the middle of a headers block. If we are, this73        # frame *must* be a CONTINUATION frame with the same stream ID as the74        # leading HEADERS or PUSH_PROMISE frame. Anything else is a75        # ProtocolError. If the frame *is* valid, append it to the header76        # buffer.77        if self._headers_buffer:78            stream_id = self._headers_buffer[0].stream_id79            valid_frame = (80                f is not None and81                isinstance(f, ContinuationFrame) and82                f.stream_id == stream_id83            )84            if not valid_frame:85                msg = "Invalid frame during header block."86                raise ProtocolError(msg)87            assert isinstance(f, ContinuationFrame)88 89            # Append the frame to the buffer.90            self._headers_buffer.append(f)91            if len(self._headers_buffer) > CONTINUATION_BACKLOG:92                msg = "Too many continuation frames received."93                raise ProtocolError(msg)94 95            # If this is the end of the header block, then we want to build a96            # mutant HEADERS frame that's massive. Use the original one we got,97            # then set END_HEADERS and set its data appopriately. If it's not98            # the end of the block, lose the current frame: we can't yield it.99            if "END_HEADERS" in f.flags:100                f = self._headers_buffer[0]101                f.flags.add("END_HEADERS")102                f.data = b"".join(x.data for x in self._headers_buffer)103                self._headers_buffer = []104            else:105                f = None106        elif (isinstance(f, (HeadersFrame, PushPromiseFrame)) and107                "END_HEADERS" not in f.flags):108            # This is the start of a headers block! Save the frame off and then109            # act like we didn't receive one.110            self._headers_buffer.append(f)111            f = None112 113        return f114 115    # The methods below support the iterator protocol.116    def __iter__(self) -> FrameBuffer:117        return self118 119    def __next__(self) -> Frame:120        # First, check that we have enough data to successfully parse the121        # next frame header. If not, bail. Otherwise, parse it.122        if len(self._data) < 9:123            raise StopIteration124 125        try:126            f, length = Frame.parse_frame_header(memoryview(self._data[:9]))127        except (InvalidDataError, InvalidFrameError) as err:  # pragma: no cover128            msg = f"Received frame with invalid header: {err!s}"129            raise ProtocolError(msg) from err130 131        # Next, check that we have enough length to parse the frame body. If132        # not, bail, leaving the frame header data in the buffer for next time.133        if len(self._data) < length + 9:134            raise StopIteration135 136        # Confirm the frame has an appropriate length.137        self._validate_frame_length(length)138 139        # Try to parse the frame body140        try:141            f.parse_body(memoryview(self._data[9:9+length]))142        except InvalidDataError as err:143            msg = "Received frame with non-compliant data"144            raise ProtocolError(msg) from err145        except InvalidFrameError as err:146            msg = "Frame data missing or invalid"147            raise FrameDataMissingError(msg) from err148 149        # At this point, as we know we'll use or discard the entire frame, we150        # can update the data.151        self._data = self._data[9+length:]152 153        # Pass the frame through the header buffer.154        new_frame = self._update_header_buffer(f)155 156        # If we got a frame we didn't understand or shouldn't yield, rather157        # than return None it'd be better if we just tried to get the next158        # frame in the sequence instead. Recurse back into ourselves to do159        # that. This is safe because the amount of work we have to do here is160        # strictly bounded by the length of the buffer.161        return new_frame if new_frame is not None else self.__next__()162