Encode long Huffman values without growing-integer shifts

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**
 
 

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
 

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