diff --git a/HISTORY.rst b/HISTORY.rst index f463918..3c4dfd2 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -3,6 +3,22 @@ History ------- +3.2.0 ++++++ + +* Added limits to pure Python record and metadata decoding to prevent denial + of service from crafted databases: 65,536 values, 512 nesting levels, and + 2 MiB of string and bytes payload. Exceeding a limit raises + ``InvalidDatabaseError``. CPython may reach its recursion limit earlier, + which raises the same error. +* Rejected unsigned integers longer than 16 bytes and signed integers longer + than 4 bytes before reading their payload. +* Updated the vendored libmaxminddb to 1.14.0, which adds the same resource + limits to the C extension. +* Truncated reads that previously raised ``IndexError`` or ``struct.error`` + now raise ``InvalidDatabaseError``. +* Improved pure Python lookup performance. + 3.1.1 (2026-03-05) ++++++++++++++++++ diff --git a/README.rst b/README.rst index 87ba401..5a84285 100644 --- a/README.rst +++ b/README.rst @@ -94,6 +94,11 @@ The module will return an ``InvalidDatabaseError`` if the database is corrupt or otherwise invalid. A ``ValueError`` will be thrown if you look up an invalid IP address or an IPv6 address in an IPv4 database. +The reader also raises ``InvalidDatabaseError`` when one record, or the +database metadata, exceeds its resource limits: 65,536 decoded values, 512 +levels of nesting, or 2 MiB of string and bytes data. Real databases stay far +below these limits. + Thread Safety ------------- diff --git a/extension/libmaxminddb b/extension/libmaxminddb index 09a0540..0077fd7 160000 --- a/extension/libmaxminddb +++ b/extension/libmaxminddb @@ -1 +1 @@ -Subproject commit 09a0540fea89a16e5c6a9e21e93ee9aece6639e3 +Subproject commit 0077fd76d00a1656b9cb3028d467736504794f41 diff --git a/maxminddb/decoder.py b/maxminddb/decoder.py index 8f67a7d..09d0665 100644 --- a/maxminddb/decoder.py +++ b/maxminddb/decoder.py @@ -3,7 +3,7 @@ from __future__ import annotations import struct -from typing import TYPE_CHECKING, ClassVar, cast +from typing import TYPE_CHECKING try: import mmap @@ -13,12 +13,45 @@ from maxminddb.errors import InvalidDatabaseError if TYPE_CHECKING: - from collections.abc import Callable - from maxminddb.file import FileBuffer from maxminddb.types import Record - DecoderFunc = Callable[["Decoder", int, int], tuple[Record, int]] + +# Per-lookup limit on the number of values decoded, recommended by the MaxMind +# DB specification. It stops a pointer fan-out, where nested pointers to shared +# targets would otherwise cost 2**depth decode operations. The count follows +# the specification's flat rule: the root is one value, each array and map +# charges its declared children, and a pointer costs nothing beyond the value +# it resolves to, which its container already charged. The largest real +# records decode a few hundred values, so the limit leaves a wide margin. +# Pointer cycles and over-deep data are caught by an explicit, call-local depth +# limit (see ``decode``). Each level costs about two interpreter frames, so +# under CPython's default recursion limit RecursionError can fire first; decode +# converts it to the same error. The explicit limit matters when a caller has +# raised the recursion limit. +_MAX_VALUES = 1 << 16 +_MAX_DEPTH = 512 +# Per-lookup limit on the total string and bytes payload materialized, matching +# libmaxminddb and the Go reader. It stops a payload amplification, where many +# pointers to one large value would otherwise materialize N * size bytes from a +# small file. Each string or bytes value is charged its length wherever it is +# decoded, so re-decoding a shared target through another pointer recharges. +_MAX_PAYLOAD_BYTES = 1 << 21 +# The widest fixed-width integer the format defines is the 16-byte uint128; a +# declared size past that is malformed and could copy attacker-controlled bytes. +_MAX_UINT_BYTES = 16 +_MAX_INT32_BYTES = 4 +# Added to a pointer value, by pointer size. A 4-byte pointer adds nothing. +_POINTER_VALUE_OFFSETS = (0, 0, 2048, 526336) +_TOO_MANY_VALUES = ( + "The MaxMind DB file's data section exceeds the maximum number of values" +) +_TOO_DEEP = "The MaxMind DB file's data section exceeds the maximum depth" +_TOO_LARGE = "The MaxMind DB file's data section exceeds the maximum payload size" +_BAD_DATA = ( + "The MaxMind DB file's data section contains bad data " + "(unknown data type or corrupt data)" +) class Decoder: @@ -42,35 +75,83 @@ def __init__( self._buffer = database_buffer self._pointer_base = pointer_base - def _decode_array(self, size: int, offset: int) -> tuple[list[Record], int]: + def _decode_array( + self, + size: int, + offset: int, + budget: list[int], + ) -> tuple[list[Record], int]: + remaining = budget[0] - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_MANY_VALUES) + budget[0] = remaining + depth = budget[1] + 1 + if depth > _MAX_DEPTH: + raise InvalidDatabaseError(_TOO_DEEP) + budget[1] = depth array = [] + decode = self._decode for _ in range(size): - (value, offset) = self.decode(offset) + (value, offset) = decode(offset, budget) array.append(value) + budget[1] -= 1 return array, offset - def _decode_boolean(self, size: int, offset: int) -> tuple[bool, int]: + def _decode_boolean( + self, + size: int, + offset: int, + _budget: list[int], + ) -> tuple[bool, int]: return size != 0, offset - def _decode_bytes(self, size: int, offset: int) -> tuple[bytes, int]: + def _decode_bytes( + self, + size: int, + offset: int, + budget: list[int], + ) -> tuple[bytes, int]: + # Charge the payload before copying so a crafted size cannot force a + # large allocation, and so pointers reusing one target recharge. + remaining = budget[2] - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_LARGE) + budget[2] = remaining new_offset = offset + size return self._buffer[offset:new_offset], new_offset - def _decode_double(self, size: int, offset: int) -> tuple[float, int]: + def _decode_double( + self, + size: int, + offset: int, + _budget: list[int], + ) -> tuple[float, int]: self._verify_size(size, 8) new_offset = offset + size packed_bytes = self._buffer[offset:new_offset] (value,) = struct.unpack(b"!d", packed_bytes) return value, new_offset - def _decode_float(self, size: int, offset: int) -> tuple[float, int]: + def _decode_float( + self, + size: int, + offset: int, + _budget: list[int], + ) -> tuple[float, int]: self._verify_size(size, 4) new_offset = offset + size packed_bytes = self._buffer[offset:new_offset] (value,) = struct.unpack(b"!f", packed_bytes) return value, new_offset - def _decode_int32(self, size: int, offset: int) -> tuple[int, int]: + def _decode_int32( + self, + size: int, + offset: int, + _budget: list[int], + ) -> tuple[int, int]: + if size > _MAX_INT32_BYTES: + raise InvalidDatabaseError(_BAD_DATA) if size == 0: return 0, offset new_offset = offset + size @@ -81,62 +162,76 @@ def _decode_int32(self, size: int, offset: int) -> tuple[int, int]: (value,) = struct.unpack(b"!i", packed_bytes) return value, new_offset - def _decode_map(self, size: int, offset: int) -> tuple[dict[str, Record], int]: + def _decode_map( + self, + size: int, + offset: int, + budget: list[int], + ) -> tuple[dict[str, Record], int]: + # A map entry decodes a key and a value, so it costs two values. + remaining = budget[0] - size * 2 + if remaining < 0: + raise InvalidDatabaseError(_TOO_MANY_VALUES) + budget[0] = remaining + depth = budget[1] + 1 + if depth > _MAX_DEPTH: + raise InvalidDatabaseError(_TOO_DEEP) + budget[1] = depth container: dict[str, Record] = {} + decode = self._decode for _ in range(size): - (key, offset) = self.decode(offset) - (value, offset) = self.decode(offset) - container[cast("str", key)] = value + (key, offset) = decode(offset, budget) + (value, offset) = decode(offset, budget) + container[key] = value # type: ignore[index] + budget[1] -= 1 return container, offset - def _decode_pointer(self, size: int, offset: int) -> tuple[Record, int]: + def _decode_pointer( + self, + size: int, + offset: int, + budget: list[int], + ) -> tuple[Record, int]: pointer_size = (size >> 3) + 1 - - buf = self._buffer[offset : offset + pointer_size] new_offset = offset + pointer_size - - if pointer_size == 1: - buf = bytes([size & 0x7]) + buf - pointer = struct.unpack(b"!H", buf)[0] + self._pointer_base - elif pointer_size == 2: - buf = b"\x00" + bytes([size & 0x7]) + buf - pointer = struct.unpack(b"!I", buf)[0] + 2048 + self._pointer_base - elif pointer_size == 3: - buf = bytes([size & 0x7]) + buf - pointer = struct.unpack(b"!I", buf)[0] + 526336 + self._pointer_base - else: - pointer = struct.unpack(b"!I", buf)[0] + self._pointer_base + pointer_bytes = self._buffer[offset:new_offset] + if len(pointer_bytes) != pointer_size: + raise InvalidDatabaseError(_BAD_DATA) + pointer = int.from_bytes(pointer_bytes, "big") + if pointer_size < 4: + # The low three bits of the ctrl byte are the high bits of the + # pointer, and sizes 2 and 3 add a fixed offset. + pointer |= (size & 0x7) << (pointer_size << 3) + pointer += _POINTER_VALUE_OFFSETS[pointer_size] + pointer += self._pointer_base if self._pointer_test: return pointer, new_offset - (value, _) = self.decode(pointer) + + # The value at the pointer's position was charged by its containing + # array or map, so the target costs nothing more. Only the depth changes. + depth = budget[1] + 1 + if depth > _MAX_DEPTH: + raise InvalidDatabaseError(_TOO_DEEP) + budget[1] = depth + (value, _) = self._decode(pointer, budget) + budget[1] -= 1 return value, new_offset - def _decode_uint(self, size: int, offset: int) -> tuple[int, int]: + def _decode_uint( + self, + size: int, + offset: int, + _budget: list[int], + ) -> tuple[int, int]: + # Reject a declared size past the widest defined unsigned integer before + # copying, so a crafted size cannot force a large allocation. + if size > _MAX_UINT_BYTES: + raise InvalidDatabaseError(_BAD_DATA) new_offset = offset + size uint_bytes = self._buffer[offset:new_offset] return int.from_bytes(uint_bytes, "big"), new_offset - def _decode_utf8_string(self, size: int, offset: int) -> tuple[str, int]: - new_offset = offset + size - return self._buffer[offset:new_offset].decode("utf-8"), new_offset - - _type_decoder: ClassVar[dict[int, DecoderFunc]] = { - 1: _decode_pointer, - 2: _decode_utf8_string, - 3: _decode_double, - 4: _decode_bytes, - 5: _decode_uint, # uint16 - 6: _decode_uint, # uint32 - 7: _decode_map, - 8: _decode_int32, - 9: _decode_uint, # uint64 - 10: _decode_uint, # uint128 - 11: _decode_array, - 14: _decode_boolean, - 15: _decode_float, - } - def decode(self, offset: int) -> tuple[Record, int]: """Decode a section of the data section starting at offset. @@ -144,6 +239,26 @@ def decode(self, offset: int) -> tuple[Record, int]: offset: the location of the data structure to decode """ + # Bound the work per lookup so a crafted database cannot exhaust CPU or + # memory. ``budget`` carries the remaining value count, the current + # nested decode depth, and the remaining string and bytes payload, so + # all three are shared across the recursion. It is call-local, which + # keeps the decoder safe for concurrent reads. The root value is charged + # here; containers charge their children. The explicit depth limit + # is independent of Python's process-wide recursion limit; RecursionError + # remains a fallback on interpreters whose stack limit is reached first. + try: + return self._decode(offset, [_MAX_VALUES - 1, 0, _MAX_PAYLOAD_BYTES]) + except RecursionError as ex: + raise InvalidDatabaseError(_TOO_DEEP) from ex + except (IndexError, struct.error) as ex: + # Truncated data: a ctrl, size, or pointer read ran off the buffer. + raise InvalidDatabaseError(_BAD_DATA) from ex + + # Keep type dispatch inline to avoid another call for every decoded value. + def _decode( # noqa: C901, PLR0911, PLR0912 + self, offset: int, budget: list[int] + ) -> tuple[Record, int]: new_offset = offset + 1 ctrl_byte = self._buffer[offset] type_num = ctrl_byte >> 5 @@ -151,16 +266,45 @@ def decode(self, offset: int) -> tuple[Record, int]: if not type_num: (type_num, new_offset) = self._read_extended(new_offset) - try: - decoder = self._type_decoder[type_num] - except KeyError as ex: - msg = f"Unexpected type number ({type_num}) encountered" - raise InvalidDatabaseError( - msg, - ) from ex - - (size, new_offset) = self._size_from_ctrl_byte(ctrl_byte, new_offset, type_num) - return decoder(self, size, new_offset) + size = ctrl_byte & 0x1F + # Sizes under 29 are stored in the ctrl byte, and a pointer's size bits + # are not a size. Skip the call for that common case. + if size >= 29 and type_num != 1: + (size, new_offset) = self._size_from_ctrl_byte(size, new_offset) + # Put common types first to reduce comparisons during real lookups. + match type_num: + case 2: + # Strings are most of the values in a real database. Decode them + # here to save a method call. + # Charge the payload before copying so a crafted size cannot force + # a large allocation, and so pointers reusing one target recharge. + remaining = budget[2] - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_LARGE) + budget[2] = remaining + end = new_offset + size + return self._buffer[new_offset:end].decode("utf-8"), end + case 1: + return self._decode_pointer(size, new_offset, budget) + case 7: + return self._decode_map(size, new_offset, budget) + case 6 | 5 | 9 | 10: # uint32, uint16, uint64, uint128 + return self._decode_uint(size, new_offset, budget) + case 11: + return self._decode_array(size, new_offset, budget) + case 3: + return self._decode_double(size, new_offset, budget) + case 4: + return self._decode_bytes(size, new_offset, budget) + case 8: + return self._decode_int32(size, new_offset, budget) + case 14: + return self._decode_boolean(size, new_offset, budget) + case 15: + return self._decode_float(size, new_offset, budget) + case _: + msg = f"Unexpected type number ({type_num}) encountered" + raise InvalidDatabaseError(msg) def _read_extended(self, offset: int) -> tuple[int, int]: next_byte = self._buffer[offset] @@ -178,24 +322,10 @@ def _read_extended(self, offset: int) -> tuple[int, int]: @staticmethod def _verify_size(expected: int, actual: int) -> None: if expected != actual: - msg = ( - "The MaxMind DB file's data section contains bad data " - "(unknown data type or corrupt data)" - ) - raise InvalidDatabaseError( - msg, - ) - - def _size_from_ctrl_byte( - self, - ctrl_byte: int, - offset: int, - type_num: int, - ) -> tuple[int, int]: - size = ctrl_byte & 0x1F - if type_num == 1 or size < 29: - return size, offset + raise InvalidDatabaseError(_BAD_DATA) + def _size_from_ctrl_byte(self, size: int, offset: int) -> tuple[int, int]: + # Called only for size codes 29 to 31, which are followed by size bytes. if size == 29: size = 29 + self._buffer[offset] return size, offset + 1 diff --git a/maxminddb/reader.py b/maxminddb/reader.py index e65d995..829b4bc 100644 --- a/maxminddb/reader.py +++ b/maxminddb/reader.py @@ -9,7 +9,6 @@ import contextlib import ipaddress -import struct from dataclasses import dataclass from ipaddress import IPv4Address, IPv6Address from typing import IO, TYPE_CHECKING, Any, AnyStr @@ -44,6 +43,7 @@ class Reader: closed: bool _decoder: Decoder _metadata: Metadata + _record_size: int _ipv4_start: int def __init__( @@ -93,6 +93,18 @@ def __init__( ) self._metadata = Metadata(**metadata) + self._record_size = self._metadata.record_size + + # Every node read stays inside the buffer once the tree fits, so the + # node reads below need no length checks of their own. + tree_end = self._metadata.search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE + if tree_end > self._buffer_size: + self.close() + msg = ( + f"Error opening database file ({filename}). The search tree " + "extends past the end of the file." + ) + raise InvalidDatabaseError(msg) self._decoder = Decoder( self._buffer, @@ -217,28 +229,25 @@ def _start_node(self, length: int) -> int: return 0 def _read_node(self, node_number: int, index: int) -> int: - base_offset = node_number * self._metadata.node_byte_size - - record_size = self._metadata.record_size - node_bytes: bytes | bytearray - if record_size == 24: - offset = base_offset + index * 3 - node_bytes = b"\x00" + self._buffer[offset : offset + 3] - elif record_size == 28: - offset = base_offset + 3 * index - node_bytes = bytearray(self._buffer[offset : offset + 4]) + record_size = self._record_size + if record_size == 28: + # Two 28-bit records share the middle byte: its high nibble + # belongs to the left record and its low nibble to the right. + base_offset = node_number * 7 if index: - node_bytes[0] = 0x0F & node_bytes[0] - else: - middle = (0xF0 & node_bytes.pop()) >> 4 - node_bytes.insert(0, middle) - elif record_size == 32: - offset = base_offset + index * 4 - node_bytes = self._buffer[offset : offset + 4] - else: - msg = f"Unknown record size: {record_size}" - raise InvalidDatabaseError(msg) - return struct.unpack(b"!I", node_bytes)[0] + offset = base_offset + 3 + record = int.from_bytes(self._buffer[offset : offset + 4], "big") + return record & 0x0FFFFFFF + record = int.from_bytes(self._buffer[base_offset : base_offset + 4], "big") + return (record >> 8) | ((record & 0xF0) << 20) + if record_size == 24: + offset = node_number * 6 + index * 3 + return int.from_bytes(self._buffer[offset : offset + 3], "big") + if record_size == 32: + offset = node_number * 8 + index * 4 + return int.from_bytes(self._buffer[offset : offset + 4], "big") + msg = f"Unknown record size: {record_size}" + raise InvalidDatabaseError(msg) def _resolve_data_pointer(self, pointer: int) -> Record: resolved = pointer - self._metadata.node_count + self._metadata.search_tree_size diff --git a/tests/data b/tests/data index b2a3df1..363086b 160000 --- a/tests/data +++ b/tests/data @@ -1 +1 @@ -Subproject commit b2a3df13c0e274d7a2dca3d5415465a3a9670e23 +Subproject commit 363086b7d90650100e91f954937794c6a090c2a0 diff --git a/tests/decoder_test.py b/tests/decoder_test.py index b755b5d..243592a 100644 --- a/tests/decoder_test.py +++ b/tests/decoder_test.py @@ -1,13 +1,53 @@ from __future__ import annotations import mmap +import sys +import threading import unittest -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, SupportsIndex from maxminddb.decoder import Decoder +from maxminddb.errors import InvalidDatabaseError if TYPE_CHECKING: from _typeshed import SizedBuffer + from typing_extensions import Self + +# Each structural level uses about two Python frames. This lets the 513-level +# cases reach the decoder's explicit limit with ample test-harness headroom. +_DEPTH_TEST_RECURSION_LIMIT = 2_000 + +_TOO_MANY_VALUES = ( + "^The MaxMind DB file's data section exceeds the maximum number of values$" +) +_TOO_DEEP = "^The MaxMind DB file's data section exceeds the maximum depth$" + + +class _HeaderOnlyBuffer(bytes): + """A buffer that fails any read past its first ``header_len`` bytes. + + It proves that a check runs before the decoder touches child or payload + bytes, rather than only that the check eventually fires. + """ + + header_len: int + + def __new__(cls, data: bytes, header_len: int) -> Self: + buf = super().__new__(cls, data) + buf.header_len = header_len + return buf + + def __getitem__(self, index: SupportsIndex | slice) -> int | bytes: # type: ignore[override] + stop = index.stop if isinstance(index, slice) else int(index) + 1 + if stop > self.header_len: + msg = f"decoder read past the {self.header_len}-byte header" + raise AssertionError(msg) + return bytes.__getitem__(self, index) + + +_PAYLOAD_TOO_LARGE = ( + "^The MaxMind DB file's data section exceeds the maximum payload size$" +) class TestDecoder(unittest.TestCase): @@ -232,3 +272,222 @@ def test_real_pointers(self) -> None: self.assertEqual(({"long_key2": "long_value2"}, 59), decoder.decode(57)) mm.close() + + @staticmethod + def _pointer(target: int) -> bytes: + # One-byte-payload pointer (type 1, pointer_size 1) with base 0. + return bytes([(1 << 5) | ((target >> 8) & 0x7), target & 0xFF]) + + def test_pointer_fan_out_is_bounded(self) -> None: + # A data section of nested arrays, each holding two pointers to the + # node below, would cost 2**depth decode operations. The decoder bounds + # the number of values it decodes per lookup and rejects the database. + depth = 100 + buf = bytearray([0xA0]) # leaf: uint16 with value 0 + prev = 0 + for _ in range(depth): + offset = len(buf) + buf += bytes([0x02, 0x04]) + self._pointer(prev) + self._pointer(prev) + prev = offset + + with self.assertRaises(InvalidDatabaseError): + Decoder(bytes(buf), pointer_base=0).decode(prev) + + @classmethod + def _scalar_pointer_array(cls, pointer_count: int) -> bytes: + # A uint16 leaf at offset 0 and, at offset 1, an array of pointers to + # it. 0x1e: extended type with size code 30; 0x04: array. + header = bytes([0xA0, 0x1E, 0x04]) + (pointer_count - 285).to_bytes(2, "big") + return header + cls._pointer(0) * pointer_count + + def test_value_limit_follows_the_flat_rule(self) -> None: + # The specification charges the root as one value and each pointer as + # the value it resolves to, not as a separate value. An array of 65,535 + # pointers to a scalar is therefore 65,536 values, exactly the limit, + # and decodes. One more pointer exceeds it. + (decoded, _) = Decoder( + self._scalar_pointer_array(65_535), pointer_base=0 + ).decode(1) + self.assertEqual(decoded, [0] * 65_535) + + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_MANY_VALUES): + Decoder(self._scalar_pointer_array(65_536), pointer_base=0).decode(1) + + def test_cyclic_pointer_raises(self) -> None: + # A pointer to itself must hit the decoder's own depth limit even when + # Python's process-wide recursion limit is much higher. + cyclic = bytes([0x20, 0x00]) # pointer (base 0) to offset 0, itself + old_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(_DEPTH_TEST_RECURSION_LIMIT) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_DEEP): + Decoder(cyclic, pointer_base=0).decode(0) + finally: + sys.setrecursionlimit(old_recursion_limit) + + def test_container_depth_is_bounded_independently_of_recursion_limit(self) -> None: + # Each prefix is an array with one element. Raising Python's global + # recursion limit proves that the decoder's call-local limit is what + # accepts 512 containers and rejects the 513th. + at_limit = bytes([0x01, 0x04]) * 512 + bytes([0xA0]) + over_limit = bytes([0x01, 0x04]) * 513 + bytes([0xA0]) + old_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(_DEPTH_TEST_RECURSION_LIMIT) + Decoder(at_limit, pointer_base=0).decode(0) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_DEEP): + Decoder(over_limit, pointer_base=0).decode(0) + finally: + sys.setrecursionlimit(old_recursion_limit) + + @classmethod + def _pointer_chain(cls, levels: int) -> tuple[bytes, int]: + # Each level is a one-element array whose element is a pointer to the + # level below, so each level costs two depth units: the array and the + # pointer follow. + buf = bytearray([0xA0]) + prev = 0 + for _ in range(levels): + offset = len(buf) + buf += bytes([0x01, 0x04]) + cls._pointer(prev) + prev = offset + return bytes(buf), prev + + def test_depth_counts_pointer_follows(self) -> None: + # 256 array-plus-pointer levels are exactly 512 depth units and decode. + # 257 exceed the limit through the decoder's own counter, not the + # interpreter's, so the error has no RecursionError cause. + old_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(_DEPTH_TEST_RECURSION_LIMIT) + buf, start = self._pointer_chain(256) + Decoder(buf, pointer_base=0).decode(start) + buf, start = self._pointer_chain(257) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_DEEP) as cm: + Decoder(buf, pointer_base=0).decode(start) + self.assertIsNone(cm.exception.__cause__) + finally: + sys.setrecursionlimit(old_recursion_limit) + + def test_budget_is_local_to_each_decode(self) -> None: + # Decoding an at-limit value twice on one Decoder, and from several + # threads at once, must succeed every time. A budget stored on the + # decoder would drain after the first call. + decoder = Decoder(self._scalar_pointer_array(65_535), pointer_base=0) + expected = [0] * 65_535 + self.assertEqual(decoder.decode(1)[0], expected) + self.assertEqual(decoder.decode(1)[0], expected) + + # Each thread writes its own slot, so the test itself has no shared + # mutable state under free threading. + results: list[object] = [None] * 8 + + def run(index: int) -> None: + results[index] = decoder.decode(1)[0] + + threads = [threading.Thread(target=run, args=(i,)) for i in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + self.assertEqual(results, [expected] * 8) + + def test_map_depth_is_bounded(self) -> None: + # Each prefix is a one-entry map whose key is the string "a" and + # whose value is the next map, so every level goes through + # _decode_map's own depth check. 0xe1: map, size 1; 0x41 0x61: "a". + at_limit = bytes([0xE1, 0x41, 0x61]) * 512 + bytes([0xA0]) + over_limit = bytes([0xE1, 0x41, 0x61]) * 513 + bytes([0xA0]) + old_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(_DEPTH_TEST_RECURSION_LIMIT) + Decoder(at_limit, pointer_base=0).decode(0) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_DEEP): + Decoder(over_limit, pointer_base=0).decode(0) + finally: + sys.setrecursionlimit(old_recursion_limit) + + def test_oversized_array_is_rejected_before_reading_children(self) -> None: + # A root array that declares 65,536 elements is 65,537 values. The + # buffer fails any read past the header, so the test proves the check + # runs before the first element. 0x1e: extended type with size code + # 30; 0x04: array; 0xfee3: 65,536 - 285. + header = _HeaderOnlyBuffer(bytes([0x1E, 0x04, 0xFE, 0xE3]), 4) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_MANY_VALUES): + Decoder(header, pointer_base=0).decode(0) + + def test_oversized_map_is_rejected_before_reading_keys(self) -> None: + # A map entry decodes a key and a value, so 32,769 entries cost 65,538 + # values, just past the limit. 0xfe: map with size code 30, then the + # two size bytes for 32,769 - 285 = 32,484 (0x7ee4). + header = _HeaderOnlyBuffer(bytes([0xFE, 0x7E, 0xE4]), 3) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_MANY_VALUES): + Decoder(header, pointer_base=0).decode(0) + + def test_oversized_string_payload_is_bounded(self) -> None: + # A single string that declares one byte more than the 2 MiB payload + # limit is rejected before its bytes are read. 0x5f: string with size + # code 31; 0x1efee4: 2,097,153 - 65,821, one byte over 2 MiB. + oversized_string = _HeaderOnlyBuffer(bytes([0x5F, 0x1E, 0xFE, 0xE4]), 4) + with self.assertRaisesRegex(InvalidDatabaseError, _PAYLOAD_TOO_LARGE): + Decoder(oversized_string, pointer_base=0).decode(0) + + def test_oversized_bytes_payload_is_bounded(self) -> None: + # As above for the bytes type. 0x9f: bytes with size code 31. + oversized_bytes = _HeaderOnlyBuffer(bytes([0x9F, 0x1E, 0xFE, 0xE4]), 4) + with self.assertRaisesRegex(InvalidDatabaseError, _PAYLOAD_TOO_LARGE): + Decoder(oversized_bytes, pointer_base=0).decode(0) + + def test_oversized_uint_is_bounded(self) -> None: + # A uint128 that declares 17 bytes exceeds the 16-byte format maximum + # and is rejected before the declared bytes are copied. 0x11: extended + # type, size 17; 0x03: extended type number 10 (uint128). + oversized_uint = _HeaderOnlyBuffer(bytes([0x11, 0x03]), 2) + with self.assertRaises(InvalidDatabaseError): + Decoder(oversized_uint, pointer_base=0).decode(0) + + def test_oversized_int32_is_bounded(self) -> None: + # An int32 that declares 5 bytes exceeds its 4-byte maximum and is + # rejected before the declared bytes are copied. 0x05: extended type, + # size 5; 0x01: extended type number 8 (int32). + oversized_int32 = _HeaderOnlyBuffer(bytes([0x05, 0x01]), 2) + with self.assertRaises(InvalidDatabaseError): + Decoder(oversized_int32, pointer_base=0).decode(0) + + def test_truncated_data_raises_invalid_database_error(self) -> None: + # A ctrl byte past the buffer end, a string header missing its size + # bytes, and a pointer missing its offset byte must not escape as + # IndexError or struct.error. + for truncated in (b"", bytes([0x5F]), bytes([0x20])): + with self.assertRaisesRegex(InvalidDatabaseError, "bad data"): + Decoder(truncated, pointer_base=0).decode(0) + + @classmethod + def _wrapped_string_pointers(cls, pointer_count: int) -> tuple[bytes, int]: + # Offset 0: a one-element array holding an inline 1 MiB string. After + # it: an array of pointers to that array. The string is inline in a + # pointed-to container, so only a charge at the string decoder itself + # catches the amplification. 0x5f: string with size code 31. + size = 1 << 20 + leaf = bytes([0x01, 0x04, 0x5F]) + (size - 65_821).to_bytes(3, "big") + leaf += b"a" * size + outer = bytes([pointer_count, 0x04]) + cls._pointer(0) * pointer_count + return leaf + outer, len(leaf) + + def test_wrapped_payload_is_charged(self) -> None: + # Two pointers materialize 2 MiB, exactly the limit. Three exceed it. + buf, start = self._wrapped_string_pointers(2) + (decoded, _) = Decoder(buf, pointer_base=0).decode(start) + self.assertEqual(decoded, [["a" * (1 << 20)]] * 2) + buf, start = self._wrapped_string_pointers(3) + with self.assertRaisesRegex(InvalidDatabaseError, _PAYLOAD_TOO_LARGE): + Decoder(buf, pointer_base=0).decode(start) + + def test_pointer_backed_map_key_is_charged(self) -> None: + # Offset 0: a string one byte over 2 MiB. Offset 4: a one-entry map + # whose key is a pointer to it. The key is decoded through the string + # decoder, so it is rejected before its bytes are read. + key = bytes([0x5F, 0x1E, 0xFE, 0xE4]) + buf = key + bytes([0xE1]) + self._pointer(0) + bytes([0xA0]) + with self.assertRaisesRegex(InvalidDatabaseError, _PAYLOAD_TOO_LARGE): + Decoder(_HeaderOnlyBuffer(buf, len(buf)), pointer_base=0).decode(len(key)) diff --git a/tests/reader_test.py b/tests/reader_test.py index 59cb574..bc24074 100644 --- a/tests/reader_test.py +++ b/tests/reader_test.py @@ -1,10 +1,12 @@ from __future__ import annotations +import contextlib import io import ipaddress import multiprocessing import os import pathlib +import sys import threading import unittest from typing import TYPE_CHECKING, cast @@ -28,9 +30,69 @@ ) if TYPE_CHECKING: + from collections.abc import Iterator + from maxminddb.reader import Reader +# Directory holding the shared MaxMind DB test fixtures. +_TEST_DATA_DIR = "tests/data/test-data" +_PAYLOAD_TOO_LARGE = ( + "^The MaxMind DB file's data section exceeds the maximum payload size$" +) +_TOO_MANY_VALUES = ( + "^The MaxMind DB file's data section exceeds the maximum number of values$" +) +_TOO_DEEP = "^The MaxMind DB file's data section exceeds the maximum depth$" +_EXTENSION_LIMIT_MESSAGE = "exceeds the configured resource limits" + + +@contextlib.contextmanager +def _bounded(seconds: int = 60, address_space: int = 2 << 30) -> Iterator[None]: + """Fail, rather than hang or exhaust memory, if a limit regresses. + + POSIX only. macOS refuses to lower RLIMIT_AS, and a process that already + uses more address space than the cap, such as one under AddressSanitizer, + would die on its next allocation; only the alarm applies in those cases. + """ + if sys.platform == "win32": + yield + return + import resource # noqa: PLC0415 + import signal # noqa: PLC0415 + + def on_alarm(*_: object) -> None: + msg = f"hostile decode did not stop within {seconds}s" + raise TimeoutError(msg) + + def address_space_in_use() -> int: + # Linux only; elsewhere the size is unknown and the cap applies. + try: + with open("/proc/self/statm") as statm: + return int(statm.read().split()[0]) * resource.getpagesize() + except (OSError, ValueError): + return 0 + + cap_memory = sys.platform != "darwin" and address_space_in_use() < address_space + if cap_memory: + soft, hard = resource.getrlimit(resource.RLIMIT_AS) + limit = ( + address_space + if hard == resource.RLIM_INFINITY + else min(address_space, hard) + ) + resource.setrlimit(resource.RLIMIT_AS, (limit, hard)) + old_handler = signal.signal(signal.SIGALRM, on_alarm) + signal.alarm(seconds) + try: + yield + finally: + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + if cap_memory: + resource.setrlimit(resource.RLIMIT_AS, (soft, hard)) + + def get_reader_from_file_descriptor(filepath: str, mode: int) -> Reader: """Patches open_database() for class TestFDReader().""" if mode == MODE_FD: @@ -47,6 +109,10 @@ class BaseTestReader(unittest.TestCase): mode: int reader_class: type[maxminddb.extension.Reader | maxminddb.reader.Reader] use_ip_objects = False + payload_error = _PAYLOAD_TOO_LARGE + value_count_error = _TOO_MANY_VALUES + metadata_error = _PAYLOAD_TOO_LARGE + fan_out_error = f"{_TOO_MANY_VALUES}|{_TOO_DEEP}" # fork doesn't work on Windows and spawn would involve pickling the reader, # which isn't possible. @@ -58,6 +124,166 @@ def ipf(self, ip: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | str: return ipaddress.ip_address(ip) return ip + def _require_resource_limits(self) -> None: + # Only resource-limit tests call this, so older system libraries still + # run the other reader tests. reader_class also handles MODE_AUTO. + if self.reader_class is maxminddb.reader.Reader: + return + self.payload_error = _EXTENSION_LIMIT_MESSAGE + self.value_count_error = _EXTENSION_LIMIT_MESSAGE + self.fan_out_error = _EXTENSION_LIMIT_MESSAGE + # libmaxminddb reports metadata rejection as a generic open failure. + self.metadata_error = "Error opening" + + # Probe with a fixture one byte over the 2 MiB payload limit, which is + # small and safe to decode even without the limits. The bundled + # libmaxminddb has them, so it must reject the probe with the + # decoder-limit message; anything else is a failure. A system library + # selected with MAXMINDDB_USE_SYSTEM_LIBMAXMINDDB may predate the + # limits and decode the probe. Skip then, rather than run the large + # DoS fixtures through a decoder that would exhaust memory. + try: + self._lookup_resource_record( + "MaxMind-DB-test-decoder-payload-limit-over.mmdb" + ) + except InvalidDatabaseError as exc: + if _EXTENSION_LIMIT_MESSAGE in str(exc): + return + raise + if not os.environ.get("MAXMINDDB_USE_SYSTEM_LIBMAXMINDDB"): + self.fail( + "the bundled libmaxminddb decoded a record over the payload limit" + ) + self.skipTest( + "system libmaxminddb predates the decoder resource limits " + "(needs the release that adds MMDB_DECODER_LIMIT_ERROR)", + ) + + def _lookup_resource_record(self, filename: str, ip: str = "0.0.0.1") -> object: + # Each DoS fixture resolves any address to its single crafted record. + with open_database(f"{_TEST_DATA_DIR}/{filename}", self.mode) as reader: + return reader.get(self.ipf(ip)) + + def test_payload_amplification_is_rejected(self) -> None: + self._require_resource_limits() + # An array of 8,192 pointers to one 65,535-byte value. The value count + # stays low, but copying each target would materialize about 512 MiB. + with ( + _bounded(), + self.assertRaisesRegex(InvalidDatabaseError, self.payload_error), + ): + self._lookup_resource_record( + "MaxMind-DB-test-payload-amplification-dos.mmdb" + ) + + def test_payload_amplification_string_is_rejected(self) -> None: + self._require_resource_limits() + # The UTF-8 string variant, so the decode path for strings is exercised. + with ( + _bounded(), + self.assertRaisesRegex(InvalidDatabaseError, self.payload_error), + ): + self._lookup_resource_record( + "MaxMind-DB-test-payload-amplification-dos-string.mmdb" + ) + + def test_payload_amplification_worst_case_is_rejected(self) -> None: + self._require_resource_limits() + # 65,535 pointers to one 65,535-byte value. The record is exactly + # 65,536 values under the flat rule, so only the payload budget can + # reject it. + with ( + _bounded(), + self.assertRaisesRegex(InvalidDatabaseError, self.payload_error), + ): + self._lookup_resource_record( + "MaxMind-DB-test-payload-amplification-dos-worst-case.mmdb" + ) + + def test_value_count_boundary(self) -> None: + self._require_resource_limits() + # The at-limit fixture decodes to exactly 65,536 values and must decode. + # The pointer-heavy fixture reaches 65,535 values through pointers, + # which cost nothing beyond the values they resolve to. One value more + # than the limit is rejected. + self.assertIsInstance( + self._lookup_resource_record("MaxMind-DB-test-decoder-value-limit.mmdb"), + list, + ) + self.assertIsInstance( + self._lookup_resource_record( + "MaxMind-DB-test-decoder-value-limit-pointer-heavy.mmdb" + ), + list, + ) + with self.assertRaisesRegex(InvalidDatabaseError, self.value_count_error): + self._lookup_resource_record( + "MaxMind-DB-test-decoder-value-limit-over.mmdb" + ) + + def test_pointer_fan_out_fixture_is_rejected(self) -> None: + self._require_resource_limits() + # A full database whose record nests arrays of pointers to the level + # below, the classic 2**depth fan-out. + with ( + _bounded(), + self.assertRaisesRegex(InvalidDatabaseError, self.fan_out_error), + ): + self._lookup_resource_record("MaxMind-DB-test-pointer-decoder-dos.mmdb") + + def test_pointer_fan_out_ipv6_fixture_is_rejected(self) -> None: + self._require_resource_limits() + # The same fan-out in a conventional IPv6 database that maps the whole + # address space to the record, so the IPv6 tree path is covered too. + with ( + _bounded(), + self.assertRaisesRegex(InvalidDatabaseError, self.fan_out_error), + ): + self._lookup_resource_record( + "MaxMind-DB-test-pointer-decoder-dos-ipv6.mmdb", "2001:db8::1" + ) + + def test_payload_at_limit_is_accepted(self) -> None: + self._require_resource_limits() + # References totaling exactly 2 MiB of payload decode successfully, so + # the limit does not reject a record at the boundary. + self.assertIsInstance( + self._lookup_resource_record("MaxMind-DB-test-decoder-payload-limit.mmdb"), + list, + ) + + def test_payload_one_over_limit_is_rejected(self) -> None: + self._require_resource_limits() + # One byte more than 2 MiB is rejected, catching an off-by-one. + with self.assertRaisesRegex(InvalidDatabaseError, self.payload_error): + self._lookup_resource_record( + "MaxMind-DB-test-decoder-payload-limit-over.mmdb" + ) + + def test_metadata_payload_limit_is_enforced_on_open(self) -> None: + self._require_resource_limits() + # Metadata must stay within the payload limit when the database is opened. + with ( + _bounded(), + self.assertRaisesRegex(InvalidDatabaseError, self.metadata_error), + open_database( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-metadata-payload-limit.mmdb", + self.mode, + ), + ): + pass + + def test_normal_record_still_decodes(self) -> None: + self._require_resource_limits() + # A record with ordinary string and bytes values, which the payload + # budget also charges, decodes unchanged. + record = cast( + "dict", + self._lookup_resource_record("MaxMind-DB-test-decoder.mmdb", "::1.1.1.0"), + ) + self.assertEqual(record["utf8_string"], "unicode! ☯ - ♫") + self.assertEqual(record["bytes"], b"\x00\x00\x00*") + def test_reader(self) -> None: for record_size in [24, 28, 32]: for ip_version in [4, 6]: @@ -346,6 +572,19 @@ def test_broken_database(self) -> None: reader.get(self.ipf("2001:220::")) reader.close() + def test_search_tree_past_end_of_file(self) -> None: + # The metadata claims more nodes than the file holds. The pure Python + # reader rejects this when the database is opened; libmaxminddb does + # the same or fails the first lookup. + with ( + self.assertRaises(InvalidDatabaseError), + open_database( + "tests/data/test-data/GeoIP2-City-Test-Invalid-Node-Count.mmdb", + self.mode, + ) as reader, + ): + reader.get(self.ipf("1.1.1.1")) + def test_ip_validation(self) -> None: reader = open_database( "tests/data/test-data/MaxMind-DB-test-decoder.mmdb",