diff --git a/HISTORY.rst b/HISTORY.rst index f463918..b5b0baa 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -3,6 +3,24 @@ 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``. +* The pure Python reader now rejects invalid search tree sizes when opening + a database. +* Improved pure Python lookup performance. + 3.1.1 (2026-03-05) ++++++++++++++++++ diff --git a/README.rst b/README.rst index 87ba401..ccdc4eb 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. Python's recursion limit may stop decoding sooner. + 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..9ef985a 100644 --- a/maxminddb/decoder.py +++ b/maxminddb/decoder.py @@ -3,7 +3,8 @@ from __future__ import annotations import struct -from typing import TYPE_CHECKING, ClassVar, cast +from dataclasses import dataclass +from typing import TYPE_CHECKING try: import mmap @@ -13,12 +14,50 @@ 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 value limit recommended by the MaxMind DB specification. It stops +# pointer fan-out, where nested containers share targets that would otherwise +# cost 2**depth decode operations. The root costs one value. Arrays charge each +# element, maps charge each key and value, and pointers cost no extra value. +# Real records decode a few hundred values, leaving a wide margin. +# An explicit depth limit catches container cycles and overly nested data. +# Python's recursion limit may fire first, which decode converts to the same +# error. The explicit limit also applies when callers raise Python's 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)" +) + + +@dataclass +class _DecodeBudget: + """Shared counters for one record or metadata decode.""" + + values_left: int + depth: int + payload_left: int class Decoder: @@ -42,35 +81,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: _DecodeBudget, + ) -> tuple[list[Record], int]: + remaining = budget.values_left - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_MANY_VALUES) + budget.values_left = remaining + depth = budget.depth + 1 + if depth > _MAX_DEPTH: + raise InvalidDatabaseError(_TOO_DEEP) + budget.depth = depth array = [] + decode = self._decode for _ in range(size): - (value, offset) = self.decode(offset) + (value, offset) = decode(offset, budget, False) # noqa: FBT003 array.append(value) + budget.depth -= 1 return array, offset - def _decode_boolean(self, size: int, offset: int) -> tuple[bool, int]: + def _decode_boolean( + self, + size: int, + offset: int, + _budget: _DecodeBudget, + ) -> 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: _DecodeBudget, + ) -> 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.payload_left - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_LARGE) + budget.payload_left = 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: _DecodeBudget, + ) -> 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: _DecodeBudget, + ) -> 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: _DecodeBudget, + ) -> tuple[int, int]: + if size > _MAX_INT32_BYTES: + raise InvalidDatabaseError(_BAD_DATA) if size == 0: return 0, offset new_offset = offset + size @@ -81,62 +168,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: _DecodeBudget, + ) -> tuple[dict[str, Record], int]: + # A map entry decodes a key and a value, so it costs two values. + remaining = budget.values_left - size * 2 + if remaining < 0: + raise InvalidDatabaseError(_TOO_MANY_VALUES) + budget.values_left = remaining + depth = budget.depth + 1 + if depth > _MAX_DEPTH: + raise InvalidDatabaseError(_TOO_DEEP) + budget.depth = 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, False) # noqa: FBT003 + (value, offset) = decode(offset, budget, False) # noqa: FBT003 + container[key] = value # type: ignore[index] + budget.depth -= 1 return container, offset - def _decode_pointer(self, size: int, offset: int) -> tuple[Record, int]: + def _decode_pointer( + self, + size: int, + offset: int, + budget: _DecodeBudget, + ) -> 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.depth + 1 + if depth > _MAX_DEPTH: + raise InvalidDatabaseError(_TOO_DEEP) + budget.depth = depth + (value, _) = self._decode(pointer, budget, True) # noqa: FBT003 + budget.depth -= 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: _DecodeBudget, + ) -> 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 +245,30 @@ def decode(self, offset: int) -> tuple[Record, int]: offset: the location of the data structure to decode """ + # Each call gets its own budget, shared by recursive calls. Charge the + # root here. + try: + return self._decode( + offset, + _DecodeBudget(_MAX_VALUES - 1, 0, _MAX_PAYLOAD_BYTES), + False, # noqa: FBT003 + ) + except RecursionError as ex: + raise InvalidDatabaseError(_TOO_DEEP) from ex + except (IndexError, struct.error) as ex: + # Convert failed buffer indexing and fixed-width unpacking. + raise InvalidDatabaseError(_BAD_DATA) from ex + + # Keep type dispatch inline to avoid another call for every decoded value. + # The positional booleans are intentional: keywords and omitted defaults + # prevented CPython from using its fastest call path in our benchmarks. + # pointer_target rejects pointers to other pointers. + def _decode( # noqa: C901, PLR0911, PLR0912 + self, + offset: int, + budget: _DecodeBudget, + pointer_target: bool, # noqa: FBT001 + ) -> tuple[Record, int]: new_offset = offset + 1 ctrl_byte = self._buffer[offset] type_num = ctrl_byte >> 5 @@ -151,16 +276,47 @@ 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.payload_left - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_LARGE) + budget.payload_left = remaining + end = new_offset + size + return self._buffer[new_offset:end].decode("utf-8"), end + case 1: + if pointer_target: + raise InvalidDatabaseError(_BAD_DATA) + 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 +334,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..32fe961 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__( @@ -67,51 +67,76 @@ def __init__( """ filename = self._load_buffer(database, mode) - metadata_start = self._buffer.rfind( - self._METADATA_START_MARKER, - max(0, self._buffer_size - 128 * 1024), - ) - - if metadata_start == -1: - self.close() - msg = ( - f"Error opening database file ({filename}). " - "Is this a valid MaxMind DB file?" - ) - raise InvalidDatabaseError( - msg, + # Include validation errors in this cleanup scope. TRY301 is suppressed + # because the handler only closes the buffer and re-raises the error. + try: + metadata_start = self._buffer.rfind( + self._METADATA_START_MARKER, + max(0, self._buffer_size - 128 * 1024), ) - metadata_start += len(self._METADATA_START_MARKER) - metadata_decoder = Decoder(self._buffer, metadata_start) - (metadata, _) = metadata_decoder.decode(metadata_start) - - if not isinstance(metadata, dict): - msg = f"Error reading metadata in database file ({filename})." - raise InvalidDatabaseError( - msg, + if metadata_start == -1: + msg = ( + f"Error opening database file ({filename}). " + "Is this a valid MaxMind DB file?" + ) + raise InvalidDatabaseError( # noqa: TRY301 + msg, + ) + + metadata_start += len(self._METADATA_START_MARKER) + metadata_decoder = Decoder(self._buffer, metadata_start) + (metadata, _) = metadata_decoder.decode(metadata_start) + + if not isinstance(metadata, dict): + msg = f"Error reading metadata in database file ({filename})." + raise InvalidDatabaseError( # noqa: TRY301 + msg, + ) + + self._metadata = Metadata(**metadata) + self._record_size = self._metadata.record_size + if self._record_size not in (24, 28, 32): + msg = f"Unknown record size: {self._record_size}" + raise InvalidDatabaseError(msg) # noqa: TRY301 + if self._metadata.node_count < 0: + msg = f"Invalid node count: {self._metadata.node_count}" + raise InvalidDatabaseError(msg) # noqa: TRY301 + + # Traversal reads nodes below node_count. Once the tree fits, those + # reads need no length checks of their own. + tree_end = ( + self._metadata.search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE ) - - self._metadata = Metadata(**metadata) - - self._decoder = Decoder( - self._buffer, - self._metadata.search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE, - ) - self.closed = False - - ipv4_start = 0 - if self._metadata.ip_version == 6: - # We store the IPv4 starting node as an optimization for IPv4 lookups - # in IPv6 trees. This allows us to skip over the first 96 nodes in - # this case. - node = 0 - for _ in range(96): - if node >= self._metadata.node_count: - break - node = self._read_node(node, 0) - ipv4_start = node - self._ipv4_start = ipv4_start + if tree_end > self._buffer_size: + msg = ( + f"Error opening database file ({filename}). The search tree " + "extends past the end of the file." + ) + raise InvalidDatabaseError(msg) # noqa: TRY301 + + self._decoder = Decoder( + self._buffer, + self._metadata.search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE, + ) + self.closed = False + + ipv4_start = 0 + if self._metadata.ip_version == 6: + # We store the IPv4 starting node as an optimization for IPv4 lookups + # in IPv6 trees. This allows us to skip over the first 96 nodes in + # this case. + node = 0 + for _ in range(96): + if node >= self._metadata.node_count: + break + node = self._read_node(node, 0) + ipv4_start = node + self._ipv4_start = ipv4_start + except BaseException: + # Release the buffer on any initialization failure. + self.close() + raise def metadata(self) -> Metadata: """Return the metadata associated with the MaxMind DB file.""" @@ -217,28 +242,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..f4760c3 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): @@ -102,6 +142,9 @@ def test_pointer(self) -> None: b"\x37\xff\xff\xff": 134744063, b"\x38\x7f\xff\xff\xff": 2147483647, b"\x38\xff\xff\xff\xff": 4294967295, + b"\x3d\xff\xff\xff\xff": 4294967295, + b"\x3e\xff\xff\xff\xff": 4294967295, + b"\x3f\xff\xff\xff\xff": 4294967295, } self.validate_type_decoding("pointers", pointers) @@ -232,3 +275,254 @@ 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_pointer_to_pointer_is_rejected(self) -> None: + # The root array shares a pointer chain that would bypass value counting. + buf = b"\xa0" + self._pointer(0) + b"\x02\x04" + self._pointer(1) * 2 + with self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"): + Decoder(buf).decode(3) + + def test_cyclic_pointer_raises(self) -> None: + with self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"): + Decoder(self._pointer(0)).decode(0) + + def test_pointer_to_container_with_pointer(self) -> None: + # A pointer may target an array that contains another pointer. + buf = b"\xa0\x01\x04" + self._pointer(0) + self._pointer(1) + self.assertEqual(Decoder(buf).decode(5), ([0], 7)) + + def test_cyclic_container_hits_depth_limit(self) -> None: + # An array containing a pointer to itself still needs a depth limit. + cyclic = b"\x01\x04" + self._pointer(0) + old_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(_DEPTH_TEST_RECURSION_LIMIT) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_DEEP) as cm: + Decoder(cyclic).decode(0) + self.assertIsNone(cm.exception.__cause__) + finally: + sys.setrecursionlimit(old_recursion_limit) + + def test_python_recursion_limit_raises_database_error(self) -> None: + # This nesting fits the decoder's limit but exceeds Python's lower limit. + buf = b"\x01\x04" * 128 + b"\xa0" + old_recursion_limit = sys.getrecursionlimit() + try: + sys.setrecursionlimit(200) + with self.assertRaisesRegex(InvalidDatabaseError, _TOO_DEEP) as cm: + Decoder(buf).decode(0) + self.assertIsInstance(cm.exception.__cause__, RecursionError) + finally: + sys.setrecursionlimit(old_recursion_limit) + + def test_sibling_maps_restore_depth(self) -> None: + # An array of 600 empty maps has depth two, regardless of its length. + buf = b"\x1e\x04" + (600 - 285).to_bytes(2, "big") + b"\xe0" * 600 + self.assertEqual(Decoder(buf).decode(0), ([{}] * 600, len(buf))) + + 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..ce95714 100644 --- a/tests/reader_test.py +++ b/tests/reader_test.py @@ -1,10 +1,13 @@ from __future__ import annotations +import contextlib import io import ipaddress import multiprocessing import os import pathlib +import sys +import tempfile import threading import unittest from typing import TYPE_CHECKING, cast @@ -28,9 +31,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 +110,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 +125,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 +573,32 @@ 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. + if self.reader_class is maxminddb.reader.Reader: + with ( + self.assertRaisesRegex( + InvalidDatabaseError, + "The search tree extends past the end of the file", + ), + open_database( + f"{_TEST_DATA_DIR}/GeoIP2-City-Test-Invalid-Node-Count.mmdb", + self.mode, + ), + ): + pass + return + 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", @@ -733,6 +986,96 @@ def setUp(self) -> None: reader_class = maxminddb.reader.Reader +class TestReaderInitialization(unittest.TestCase): + def test_empty_search_tree_is_accepted(self) -> None: + data = pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" + ).read_bytes() + original = b"node_count\xc1\xa3" + self.assertEqual(data.count(original), 1) + with ( + io.BytesIO(data.replace(original, b"node_count\xc0")) as database, + maxminddb.reader.Reader(database, MODE_FD) as reader, + ): + self.assertIsNone(reader.get("1.1.1.1")) + self.assertEqual(list(reader), []) + + def test_invalid_tree_metadata_is_rejected_on_open(self) -> None: + data = pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" + ).read_bytes() + cases = ( + (b"record_size\xa1\x18", b"record_size\xa1\x1e", "Unknown record size: 30"), + ( + b"node_count\xc1\xa3", + b"node_count\x04\x01\xff\xff\xff\xff", + "Invalid node count: -1", + ), + ) + for original, replacement, message in cases: + with self.subTest(message=message): + self.assertEqual(data.count(original), 1) + with ( + io.BytesIO(data.replace(original, replacement)) as database, + self.assertRaisesRegex(InvalidDatabaseError, message), + maxminddb.reader.Reader(database, MODE_FD), + ): + pass + + def test_failed_initialization_closes_buffer(self) -> None: + reader_class = maxminddb.reader.Reader + marker = b"\xab\xcd\xefMaxMind.com" + cases = ( + (b"not a database", InvalidDatabaseError, "Is this a valid MaxMind DB"), + (marker + b"\x40", InvalidDatabaseError, "Error reading metadata"), + (marker + b"\xe0", TypeError, "required keyword-only arguments"), + ( + pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-metadata-payload-limit.mmdb" + ).read_bytes(), + InvalidDatabaseError, + _PAYLOAD_TOO_LARGE, + ), + ) + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "invalid.mmdb" + for mode in (MODE_FILE, MODE_MMAP): + for data, error, message in cases: + with self.subTest(mode=mode, message=message): + path.write_bytes(data) + with ( + _bounded(), + mock.patch.object( + reader_class, + "close", + autospec=True, + side_effect=reader_class.close, + ) as close, + self.assertRaisesRegex(error, message), + ): + reader_class(path, mode) + close.assert_called_once() + reader = close.call_args.args[0] + self.assertTrue(reader.closed) + if mode == MODE_FILE: + self.assertTrue(reader._buffer._handle.closed) # noqa: SLF001 + else: + self.assertTrue(reader._buffer.closed) # noqa: SLF001 + + +class TestSearchTreeNodes(unittest.TestCase): + def test_28_bit_records_preserve_high_nibbles(self) -> None: + # Node decoding needs only the record size and buffer, not a database. + reader = object.__new__(maxminddb.reader.Reader) + reader._record_size = 28 # noqa: SLF001 + # The middle byte holds the left record's high nibble, then the right's. + reader._buffer = bytes.fromhex("aabbcc de ff0011 123456 f8 789abc") # noqa: SLF001 + self.assertEqual(reader._read_node(0, 0), 0xDAABBCC) # noqa: SLF001 + self.assertEqual(reader._read_node(0, 1), 0xEFF0011) # noqa: SLF001 + self.assertEqual(reader._read_node(1, 0), 0xF123456) # noqa: SLF001 + self.assertEqual(reader._read_node(1, 1), 0x8789ABC) # noqa: SLF001 + + class TestOldReader(unittest.TestCase): def test_old_reader(self) -> None: reader = maxminddb.Reader("tests/data/test-data/MaxMind-DB-test-decoder.mmdb")