codekingpro/portable-devtools

1"""2Implements the HPACK header compression algorithm as detailed by RFC 7541.3"""4from __future__ import annotations5 6import logging7from typing import TYPE_CHECKING, Any8 9from .exceptions import HPACKDecodingError, InvalidTableSizeError, OversizedHeaderListError10from .huffman import HuffmanEncoder11from .huffman_constants import REQUEST_CODES, REQUEST_CODES_LENGTH12from .huffman_table import decode_huffman13from .struct import HeaderTuple, HeaderWeaklyTyped, NeverIndexedHeaderTuple14from .table import HeaderTable, table_entry_size15 16if TYPE_CHECKING:17    from collections.abc import Iterable  # pragma: no cover18 19log = logging.getLogger(__name__)20 21INDEX_NONE = b"\x00"22INDEX_NEVER = b"\x10"23INDEX_INCREMENTAL = b"\x40"24 25# Precompute 2^i for 1-8 for use in prefix calcs.26# Zero index is not used but there to save a subtraction27# as prefix numbers are not zero indexed.28_PREFIX_BIT_MAX_NUMBERS = [(2 ** i) - 1 for i in range(9)]29 30# We default the maximum header list we're willing to accept to 64kB. That's a31# lot of headers, but if applications want to raise it they can do.32DEFAULT_MAX_HEADER_LIST_SIZE = 2 ** 1633 34 35def _unicode_if_needed(header: HeaderWeaklyTyped, raw: bool) -> HeaderTuple:36    """37    Provides a header as a unicode string if raw is False, otherwise returns38    it as a bytestring.39    """40    name = bytes(header[0])  # type: ignore41    value = bytes(header[1])  # type: ignore42 43    if not raw:44        return header.__class__(name.decode("utf-8"), value.decode("utf-8"))  # type: ignore45    return header.__class__(name, value)  # type: ignore46 47 48def encode_integer(integer: int, prefix_bits: int) -> bytearray:49    """50    Encodes an integer according to the wacky integer encoding rules51    defined in the HPACK spec.52    """53    log.debug("Encoding %d with %d bits", integer, prefix_bits)54 55    if integer < 0:56        msg = f"Can only encode positive integers, got {integer}"57        raise ValueError(msg)58 59    if prefix_bits < 1 or prefix_bits > 8:60        msg = f"Prefix bits must be between 1 and 8, got {prefix_bits}"61        raise ValueError(msg)62 63    max_number = _PREFIX_BIT_MAX_NUMBERS[prefix_bits]64 65    if integer < max_number:66        return bytearray([integer])  # Seriously?67    elements = [max_number]68    integer -= max_number69 70    while integer >= 128:71        elements.append((integer & 127) + 128)72        integer >>= 773 74    elements.append(integer)75 76    return bytearray(elements)77 78 79def decode_integer(data: bytes, prefix_bits: int) -> tuple[int, int]:80    """81    Decodes an integer according to the wacky integer encoding rules82    defined in the HPACK spec. Returns a tuple of the decoded integer and the83    number of bytes that were consumed from ``data`` in order to get that84    integer.85    """86    if prefix_bits < 1 or prefix_bits > 8:87        msg = f"Prefix bits must be between 1 and 8, got {prefix_bits}"88        raise ValueError(msg)89 90    max_number = _PREFIX_BIT_MAX_NUMBERS[prefix_bits]91    index = 192    shift = 093    mask = (0xFF >> (8 - prefix_bits))94 95    try:96        number = data[0] & mask97        if number == max_number:98            while True:99                next_byte = data[index]100                index += 1101 102                if next_byte >= 128:103                    number += (next_byte - 128) << shift104                else:105                    number += next_byte << shift106                    break107                shift += 7108 109    except IndexError as err:110        msg = f"Unable to decode HPACK integer representation from {data!r}"111        raise HPACKDecodingError(msg) from err112 113    log.debug("Decoded %d, consumed %d bytes", number, index)114 115    return number, index116 117 118def _dict_to_iterable(header_dict: dict[bytes | str, bytes | str]) \119        -> Iterable[tuple[bytes | str, bytes | str]]:120    """121    Converts a dictionary to an iterable of key-value tuples. This is a122    HPACK-specific function because it pulls "special-headers" out first and123    then emits them.124    """125    if not isinstance(header_dict, dict):  # pragma: no cover126        msg = f"header_dict not a dict, but {type(header_dict)}"127        raise TypeError(msg)128 129    keys = sorted(130        header_dict.keys(),131        key=lambda k: not _to_bytes(k).startswith(b":"),132    )133    for key in keys:134        yield key, header_dict[key]135 136 137def _to_bytes(value: bytes | str | Any) -> bytes:138    """139    Convert anything to bytes through a UTF-8 encoded string140    """141    t = type(value)142    if t is bytes:143        return value  # type: ignore144    if t is not str:145        value = str(value)146    return value.encode("utf-8")  # type: ignore147 148 149class Encoder:150    """151    An HPACK encoder object. This object takes HTTP headers and emits encoded152    HTTP/2 header blocks.153    """154 155    def __init__(self) -> None:156        self.header_table = HeaderTable()157        self.huffman_coder = HuffmanEncoder(158            REQUEST_CODES, REQUEST_CODES_LENGTH,159        )160        self.table_size_changes: list[int] = []161 162    @property163    def header_table_size(self) -> int:164        """165        Controls the size of the HPACK header table.166        """167        return self.header_table.maxsize168 169    @header_table_size.setter170    def header_table_size(self, value: int) -> None:171        self.header_table.maxsize = value172        if self.header_table.resized:173            self.table_size_changes.append(value)174 175    def encode(self,176               headers: Iterable[\177                   HeaderTuple | \178                   tuple[bytes | str, bytes | str] | \179                   tuple[bytes | str, bytes | str, bool | None]] | \180                   dict[bytes | str, bytes | str],181               huffman: bool = True) -> bytes:182        """183        Takes a set of headers and encodes them into a HPACK-encoded header184        block.185 186        :param headers: The headers to encode. Must be either an iterable of187                        tuples, an iterable of :class:`HeaderTuple188                        <hpack.HeaderTuple>`, or a ``dict``.189 190                        If an iterable of tuples, the tuples may be either191                        two-tuples or three-tuples. If they are two-tuples, the192                        tuples must be of the format ``(name, value)``. If they193                        are three-tuples, they must be of the format194                        ``(name, value, sensitive)``, where ``sensitive`` is a195                        boolean value indicating whether the header should be196                        added to header tables anywhere. If not present,197                        ``sensitive`` defaults to ``False``.198 199                        If an iterable of :class:`HeaderTuple200                        <hpack.HeaderTuple>`, the tuples must always be201                        two-tuples. Instead of using ``sensitive`` as a third202                        tuple entry, use :class:`NeverIndexedHeaderTuple203                        <hpack.NeverIndexedHeaderTuple>` to request that204                        the field never be indexed.205 206                        .. warning:: HTTP/2 requires that all special headers207                            (headers whose names begin with ``:`` characters)208                            appear at the *start* of the header block. While209                            this method will ensure that happens for ``dict``210                            subclasses, callers using any other iterable of211                            tuples **must** ensure they place their special212                            headers at the start of the iterable.213 214                            For efficiency reasons users should prefer to use215                            iterables of two-tuples: fixing the ordering of216                            dictionary headers is an expensive operation that217                            should be avoided if possible.218 219        :param huffman: (optional) Whether to Huffman-encode any header sent as220                        a literal value. Except for use when debugging, it is221                        recommended that this be left enabled.222 223        :returns: A bytestring containing the HPACK-encoded header block.224        """225        # Transforming the headers into a header block is a procedure that can226        # be modeled as a chain or pipe. First, the headers are encoded. This227        # encoding can be done a number of ways. If the header name-value pair228        # are already in the header table we can represent them using the229        # indexed representation: the same is true if they are in the static230        # table. Otherwise, a literal representation will be used.231        header_block = []232 233        # Before we begin, if the header table size has been changed we need234        # to signal all changes since last emission appropriately.235        if self.header_table.resized:236            header_block.append(self._encode_table_size_change())237            self.header_table.resized = False238 239        if isinstance(headers, dict):240            # Turn the headers into a list of tuples if possible. This is the241            # natural way to interact with them in HPACK. Because dictionaries are242            # un-ordered, we need to make sure we grab the "special" headers first.243            hpack_headers = _dict_to_iterable(headers)244        else:245            """246            Assume headers is an iterable of HeaderTuples, or plain 2-tuples, or plain 3-tuples:247 248            examples:249            [250                HeaderTuple(':method', 'GET'),251                NeverIndexedHeaderTuple('customkey', 'sensitiveinfo'),252            ]253            or254            [255                (':method', 'GET'),256                ('customkey', 'some-data'),257            ]258            or259            [260                (':method', 'GET', True),261                ('customkey', 'sensitiveinfo', True),262            ]263            """264            hpack_headers = iter(headers)  # type: ignore265 266        # Add each header to the header block267        for header in hpack_headers:268            sensitive = False269            if isinstance(header, HeaderTuple):270                # HeaderTuple implies it's a 2-tuple with the sensitive information stored as instance attribute271                sensitive = not header.indexable272            elif len(header) > 2:273                sensitive = header[2]274 275            new_header = (_to_bytes(header[0]), _to_bytes(header[1]))276            header_block.append(self.add(new_header, sensitive, huffman))277 278        encoded = b"".join(header_block)279 280        log.debug("Encoded header block to %s", encoded)281 282        return encoded283 284    def add(self, to_add: tuple[bytes, bytes], sensitive: bool, huffman: bool = False) -> bytes:285        """286        Serializes a header key-value tuple.287        """288        log.debug(289            "Adding %s to the header table, sensitive:%s, huffman:%s",290            to_add,291            sensitive,292            huffman,293        )294 295        name, value = to_add296 297        # Set our indexing mode298        indexbit = INDEX_INCREMENTAL if not sensitive else INDEX_NEVER299 300        # Search for a matching header in the header table.301        match = self.header_table.search(name, value)302 303        if match is None:304            # Not in the header table. Encode using the literal syntax,305            # and add it to the header table.306            encoded = self._encode_literal(name, value, indexbit, huffman)307            if not sensitive:308                self.header_table.add(name, value)309            return encoded310 311        # The header is in the table, break out the values. If we matched312        # perfectly, we can use the indexed representation: otherwise we313        # can use the indexed literal.314        index, name, perfect = match315 316        if perfect:317            # Indexed representation.318            encoded = self._encode_indexed(index)319        else:320            # Indexed literal. We are going to add header to the321            # header table unconditionally. It is a future todo to322            # filter out headers which are known to be ineffective for323            # indexing since they just take space in the table and324            # pushed out other valuable headers.325            encoded = self._encode_indexed_literal(326                index, value, indexbit, huffman,327            )328            if not sensitive:329                self.header_table.add(name, value)330 331        return encoded332 333    def _encode_indexed(self, index: int) -> bytes:334        """335        Encodes a header using the indexed representation.336        """337        field = encode_integer(index, 7)338        field[0] |= 0x80  # we set the top bit339        return bytes(field)340 341    def _encode_literal(self, name: bytes, value: bytes, indexbit: bytes, huffman: bool = False) -> bytes:342        """343        Encodes a header with a literal name and literal value. If ``indexing``344        is True, the header will be added to the header table: otherwise it345        will not.346        """347        if huffman:348            name = self.huffman_coder.encode(name)349            value = self.huffman_coder.encode(value)350 351        name_len = encode_integer(len(name), 7)352        value_len = encode_integer(len(value), 7)353 354        if huffman:355            name_len[0] |= 0x80356            value_len[0] |= 0x80357 358        return b"".join(359            [indexbit, bytes(name_len), name, bytes(value_len), value],360        )361 362    def _encode_indexed_literal(self, index: int, value: bytes, indexbit: bytes, huffman: bool = False) -> bytes:363        """364        Encodes a header with an indexed name and a literal value and performs365        incremental indexing.366        """367        if indexbit != INDEX_INCREMENTAL:368            prefix = encode_integer(index, 4)369        else:370            prefix = encode_integer(index, 6)371 372        prefix[0] |= ord(indexbit)373 374        if huffman:375            value = self.huffman_coder.encode(value)376 377        value_len = encode_integer(len(value), 7)378 379        if huffman:380            value_len[0] |= 0x80381 382        return b"".join([bytes(prefix), bytes(value_len), value])383 384    def _encode_table_size_change(self) -> bytes:385        """386        Produces the encoded form of all header table size change context387        updates.388        """389        block = b""390        for size_bytes in self.table_size_changes:391            b = encode_integer(size_bytes, 5)392            b[0] |= 0x20393            block += bytes(b)394        self.table_size_changes = []395        return block396 397 398class Decoder:399    """400    An HPACK decoder object.401 402    .. versionchanged:: 2.3.0403       Added ``max_header_list_size`` argument.404 405    :param max_header_list_size: The maximum decompressed size we will allow406        for any single header block. This is a protection against DoS attacks407        that attempt to force the application to expand a relatively small408        amount of data into a really large header list, allowing enormous409        amounts of memory to be allocated.410 411        If this amount of data is exceeded, a `OversizedHeaderListError412        <hpack.OversizedHeaderListError>` exception will be raised. At this413        point the connection should be shut down, as the HPACK state will no414        longer be usable.415 416        Defaults to 64kB.417    :type max_header_list_size: ``int``418    """419 420    def __init__(self, max_header_list_size: int = DEFAULT_MAX_HEADER_LIST_SIZE) -> None:421        self.header_table = HeaderTable()422 423        #: The maximum decompressed size we will allow for any single header424        #: block. This is a protection against DoS attacks that attempt to425        #: force the application to expand a relatively small amount of data426        #: into a really large header list, allowing enormous amounts of memory427        #: to be allocated.428        #:429        #: If this amount of data is exceeded, a `OversizedHeaderListError430        #: <hpack.OversizedHeaderListError>` exception will be raised. At this431        #: point the connection should be shut down, as the HPACK state will no432        #: longer be usable.433        #:434        #: Defaults to 64kB.435        #:436        #: .. versionadded:: 2.3.0437        self.max_header_list_size = max_header_list_size438 439        #: Maximum allowed header table size.440        #:441        #: A HTTP/2 implementation should set this to the most recent value of442        #: SETTINGS_HEADER_TABLE_SIZE that it sent *and has received an ACK443        #: for*. Once this setting is set, the actual header table size will be444        #: checked at the end of each decoding run and whenever it is changed,445        #: to confirm that it fits in this size.446        self.max_allowed_table_size = self.header_table.maxsize447 448    @property449    def header_table_size(self) -> int:450        """451        Controls the size of the HPACK header table.452        """453        return self.header_table.maxsize454 455    @header_table_size.setter456    def header_table_size(self, value: int) -> None:457        self.header_table.maxsize = value458 459    def decode(self, data: bytes, raw: bool = False) -> Iterable[HeaderTuple]:460        """461        Takes an HPACK-encoded header block and decodes it into a header set.462 463        :param data: A bytestring representing a complete HPACK-encoded header464                     block.465        :param raw: (optional) Whether to return the headers as tuples of raw466                    byte strings or to decode them as UTF-8 before returning467                    them. The default value is False, which returns tuples of468                    Unicode strings469        :returns: A list of two-tuples of ``(name, value)`` representing the470                  HPACK-encoded headers, in the order they were decoded.471        :raises HPACKDecodingError: If an error is encountered while decoding472                                    the header block.473        """474        log.debug("Decoding %s", data)475 476        data_mem = memoryview(data)477        headers: list[HeaderTuple] = []478        data_len = len(data)479        inflated_size = 0480        current_index = 0481 482        while current_index < data_len:483            # Work out what kind of header we're decoding.484            # If the high bit is 1, it's an indexed field.485            current = data[current_index]486            indexed = bool(current & 0x80)487 488            # Otherwise, if the second-highest bit is 1 it's a field that does489            # alter the header table.490            literal_index = bool(current & 0x40)491 492            # Otherwise, if the third-highest bit is 1 it's an encoding context493            # update.494            encoding_update = bool(current & 0x20)495 496            if indexed:497                header, consumed = self._decode_indexed(498                    data_mem[current_index:],499                )500            elif literal_index:501                # It's a literal header that does affect the header table.502                header, consumed = self._decode_literal_index(503                    data_mem[current_index:],504                )505            elif encoding_update:506                # It's an update to the encoding context. These are forbidden507                # in a header block after any actual header.508                if headers:509                    msg = "Table size update not at the start of the block"510                    raise HPACKDecodingError(msg)511                consumed = self._update_encoding_context(512                    data_mem[current_index:],513                )514                header = None515            else:516                # It's a literal header that does not affect the header table.517                header, consumed = self._decode_literal_no_index(518                    data_mem[current_index:],519                )520 521            if header:522                headers.append(header)523                inflated_size += table_entry_size(header[0], header[1])524 525                if inflated_size > self.max_header_list_size:526                    msg = f"A header list larger than {self.max_header_list_size} has been received"527                    raise OversizedHeaderListError(msg)528 529            current_index += consumed530 531        # Confirm that the table size is lower than the maximum. We do this532        # here to ensure that we catch when the max has been *shrunk* and the533        # remote peer hasn't actually done that.534        self._assert_valid_table_size()535 536        try:537            return [_unicode_if_needed(h, raw) for h in headers]538        except UnicodeDecodeError as err:539            msg = "Unable to decode headers as UTF-8"540            raise HPACKDecodingError(msg) from err541 542    def _assert_valid_table_size(self) -> None:543        """544        Check that the table size set by the encoder is lower than the maximum545        we expect to have.546        """547        if self.header_table_size > self.max_allowed_table_size:548            msg = "Encoder did not shrink table size to within the max"549            raise InvalidTableSizeError(msg)550 551    def _update_encoding_context(self, data: bytes) -> int:552        """553        Handles a byte that updates the encoding context.554        """555        # We've been asked to resize the header table.556        new_size, consumed = decode_integer(data, 5)557        if new_size > self.max_allowed_table_size:558            msg = "Encoder exceeded max allowable table size"559            raise InvalidTableSizeError(msg)560        self.header_table_size = new_size561        return consumed562 563    def _decode_indexed(self, data: bytes) -> tuple[HeaderTuple, int]:564        """565        Decodes a header represented using the indexed representation.566        """567        index, consumed = decode_integer(data, 7)568        header = HeaderTuple(*self.header_table.get_by_index(index))569        log.debug("Decoded %s, consumed %d", header, consumed)570        return header, consumed571 572    def _decode_literal_no_index(self, data: bytes) -> tuple[HeaderTuple, int]:573        return self._decode_literal(data, should_index=False)574 575    def _decode_literal_index(self, data: bytes) -> tuple[HeaderTuple, int]:576        return self._decode_literal(data, should_index=True)577 578    def _decode_literal(self, data: bytes, should_index: bool) -> tuple[HeaderTuple, int]:579        """580        Decodes a header represented with a literal.581        """582        total_consumed = 0583 584        # When should_index is true, if the low six bits of the first byte are585        # nonzero, the header name is indexed.586        # When should_index is false, if the low four bits of the first byte587        # are nonzero the header name is indexed.588        if should_index:589            indexed_name = data[0] & 0x3F590            name_len = 6591            not_indexable = False592        else:593            high_byte = data[0]594            indexed_name = high_byte & 0x0F595            name_len = 4596            not_indexable = bool(high_byte & 0x10)597 598        if indexed_name:599            # Indexed header name.600            index, consumed = decode_integer(data, name_len)601            name = self.header_table.get_by_index(index)[0]602 603            total_consumed = consumed604            length = 0605        else:606            # Literal header name. The first byte was consumed, so we need to607            # move forward.608            data = data[1:]609 610            length, consumed = decode_integer(data, 7)611            name = data[consumed:consumed + length]612            if len(name) != length:613                msg = "Truncated header block"614                raise HPACKDecodingError(msg)615 616            if data[0] & 0x80:617                name = decode_huffman(name)618            total_consumed = consumed + length + 1  # Since we moved forward 1.619 620        data = data[consumed + length:]621 622        # The header value is definitely length-based.623        length, consumed = decode_integer(data, 7)624        value = data[consumed:consumed + length]625        if len(value) != length:626            msg = "Truncated header block"627            raise HPACKDecodingError(msg)628 629        if data[0] & 0x80:630            value = decode_huffman(value)631 632        # Updated the total consumed length.633        total_consumed += length + consumed634 635        # If we have been told never to index the header field, encode that in636        # the tuple we use.637        header: HeaderTuple638        if not_indexable:639            header = NeverIndexedHeaderTuple(name, value)640        else:641            header = HeaderTuple(name, value)642 643        # If we've been asked to index this, add it to the header table.644        if should_index:645            self.header_table.add(name, value)646 647        log.debug(648            "Decoded %s, total consumed %d bytes, indexed %s",649            header,650            total_consumed,651            should_index,652        )653 654        return header, total_consumed655