diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a4544ef..f767222 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -10,6 +10,8 @@ dev **API Changes (Backward Compatible)** +- Encoding long header values with Huffman coding is no longer quadratic in the value length. + **Bugfixes** diff --git a/src/hpack/huffman.py b/src/hpack/huffman.py index 1c55de0..e259ac4 100644 --- a/src/hpack/huffman.py +++ b/src/hpack/huffman.py @@ -24,6 +24,24 @@ def encode(self, bytes_to_encode: bytes | None) -> bytes: if not bytes_to_encode: return b"" + # For longer byte strings, assemble the bits in C instead of repeatedly + # shifting an ever-growing Python integer. Build the symbol map for this + # call so changes to caller-owned coding tables are always observed. + if type(bytes_to_encode) is bytes and len(bytes_to_encode) >= 256 and type(self) is HuffmanEncoder: + codes = self.huffman_code_list + lengths = self.huffman_code_list_lengths + if type(codes) is list and type(lengths) is list and len(codes) >= 256 and len(lengths) >= 256: + symbols = {} + for symbol in set(bytes_to_encode): + code, length = codes[symbol], lengths[symbol] + if type(code) is not int or type(length) is not int or not 1 <= length <= 30 or not 0 <= code < (1 << length): + break + symbols[symbol] = f"{code:0{length}b}" + else: + bits = "".join(map(symbols.__getitem__, bytes_to_encode)) + bits += "1" * (-len(bits) % 8) + return int(bits, 2).to_bytes(len(bits) // 8, "big") + final_num = 0 final_int_len = 0 diff --git a/tests/test_huffman.py b/tests/test_huffman.py index 0ea789f..232795e 100644 --- a/tests/test_huffman.py +++ b/tests/test_huffman.py @@ -52,3 +52,29 @@ def test_huffman_decoder_properly_handles_all_bytestrings(self, data): result = b'' assert isinstance(result, bytes) + + +def test_long_huffman_encoding_observes_mutable_tables(): + codes, lengths = REQUEST_CODES[:], REQUEST_CODES_LENGTH[:] + encoder = HuffmanEncoder(codes, lengths) + assert decode_huffman(encoder.encode(b'a' * 512)) == b'a' * 512 + codes[ord('a')], lengths[ord('a')] = codes[ord('b')], lengths[ord('b')] + assert decode_huffman(encoder.encode(b'a' * 512)) == b'b' * 512 + + +def test_long_huffman_encoding_custom_table_fallbacks(): + # Over-wide code words preserve the original masking and overlap behavior. + encoder = HuffmanEncoder([3] * 256, [1] * 256) + assert encoder.encode(b'x' * 256) == b'\x01' + b'\xff' * 32 + encoder = HuffmanEncoder([0] * 256, [0] * 256) + assert encoder.encode(b'x' * 256) == b'\x00' + class Code(int): + pass + encoder = HuffmanEncoder([Code(0)] * 256, [1] * 256) + assert encoder.encode(b'x' * 256) == b'\x00' * 32 + class Codes(list): + pass + encoder = HuffmanEncoder(Codes([0] * 256), [1] * 256) + assert encoder.encode(b'x' * 256) == b'\x00' * 32 + encoder = HuffmanEncoder([0] * 120, [1] * 120) + assert encoder.encode(b'a' * 256) == b'\x00' * 32