diff --git a/HISTORY.rst b/HISTORY.rst index c566ae2..5d787d5 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -23,9 +23,14 @@ History C extension. Before, the iterator walked the new database with node numbers from the old one. A failed ``__init__`` keeps the old database. After ``close()``, an iterator raises ``ValueError`` in every mode. -* For a record map with a key that cannot be hashed, such as a list, the - pure Python reader raises ``InvalidDatabaseError`` instead of ``TypeError``, - as the C extension does. +* The pure Python reader now raises ``InvalidDatabaseError`` when a value, + search tree record or pointer goes past the end of the data section, when + the search tree overlaps the metadata, or when the file shrinks after open in + ``MODE_FILE``. Previously, it could return a shorter or wrong value, or + decode bytes from the metadata section. +* The pure Python reader now raises ``InvalidDatabaseError`` when a record map + key is not a string, as the C extension does. Previously, it returned the + non-string key or raised ``TypeError``. * C extension: * Fixed segmentation faults from invalid use of ``Metadata``, ``Reader`` and diff --git a/maxminddb/decoder.py b/maxminddb/decoder.py index 7fc3444..ac71e3d 100644 --- a/maxminddb/decoder.py +++ b/maxminddb/decoder.py @@ -49,6 +49,8 @@ "The MaxMind DB file's data section contains bad data " "(unknown data type or corrupt data)" ) +# The same message as the C extension. +_NON_STRING_KEY = "Invalid map key: the key is not a string." @dataclass @@ -68,6 +70,7 @@ def __init__( database_buffer: FileBuffer | mmap.mmap | bytes | bytearray, pointer_base: int = 0, pointer_test: bool = False, # noqa: FBT001, FBT002 + data_end: int | None = None, ) -> None: """Create a Decoder for a MaxMind DB. @@ -75,11 +78,18 @@ def __init__( database_buffer: an mmap'd MaxMind DB file. pointer_base: the base number to use when decoding a pointer pointer_test: used for internal unit testing of pointer code + data_end: the offset where the decoded section ends. No payload + may extend past it, and no pointer may target it or a later + offset. Control and size bytes are not checked. The default + is the end of the buffer. """ self._pointer_test = pointer_test self._buffer = database_buffer self._pointer_base = pointer_base + if data_end is None: + data_end = len(database_buffer) + self._data_end = data_end def _decode_array( self, @@ -124,6 +134,8 @@ def _decode_bytes( raise InvalidDatabaseError(_TOO_LARGE) budget.payload_left = remaining new_offset = offset + size + if new_offset > self._data_end: + raise InvalidDatabaseError(_BAD_DATA) return self._buffer[offset:new_offset], new_offset def _decode_double( @@ -134,6 +146,8 @@ def _decode_double( ) -> tuple[float, int]: self._verify_size(size, 8) new_offset = offset + size + if new_offset > self._data_end: + raise InvalidDatabaseError(_BAD_DATA) packed_bytes = self._buffer[offset:new_offset] (value,) = struct.unpack(b"!d", packed_bytes) return value, new_offset @@ -146,6 +160,8 @@ def _decode_float( ) -> tuple[float, int]: self._verify_size(size, 4) new_offset = offset + size + if new_offset > self._data_end: + raise InvalidDatabaseError(_BAD_DATA) packed_bytes = self._buffer[offset:new_offset] (value,) = struct.unpack(b"!f", packed_bytes) return value, new_offset @@ -161,6 +177,8 @@ def _decode_int32( if size == 0: return 0, offset new_offset = offset + size + if new_offset > self._data_end: + raise InvalidDatabaseError(_BAD_DATA) packed_bytes = self._buffer[offset:new_offset] if size != 4: @@ -185,17 +203,57 @@ def _decode_map( budget.depth = depth container: dict[str, Record] = {} decode = self._decode + decode_key = self._decode_key for _ in range(size): - (key, offset) = decode(offset, budget, False) # noqa: FBT003 + (key, offset) = decode_key(offset, budget) (value, offset) = decode(offset, budget, False) # noqa: FBT003 - try: - container[key] = value # type: ignore[index] - except TypeError as ex: - # The key cannot be hashed, such as a list. - raise InvalidDatabaseError(_BAD_DATA) from ex + container[key] = value budget.depth -= 1 return container, offset + def _decode_key(self, offset: int, budget: _DecodeBudget) -> tuple[str, int]: + # A key is a string or a pointer to one. Real databases share most keys + # through pointers, so decode both here without the generic dispatch. + # A string cannot nest, so a pointer key does not count toward the + # depth. This repeats the pointer arithmetic of _decode_pointer to save + # a call per key. pointer_test applies to pointer values only. + buffer = self._buffer + data_end = self._data_end + ctrl_byte = buffer[offset] + new_offset = offset + 1 + is_pointer = ctrl_byte >> 5 == 1 + if is_pointer: + size = ctrl_byte & 0x1F + pointer_size = (size >> 3) + 1 + pointer_end = new_offset + pointer_size + if pointer_end > data_end: + raise InvalidDatabaseError(_BAD_DATA) + pointer = int.from_bytes(buffer[new_offset:pointer_end], "big") + if pointer_size < 4: + pointer |= (size & 0x7) << (pointer_size << 3) + pointer += _POINTER_VALUE_OFFSETS[pointer_size] + pointer += self._pointer_base + # The end check below also rejects a target past data_end. + ctrl_byte = buffer[pointer] + new_offset = pointer + 1 + if ctrl_byte >> 5 != 2: + raise InvalidDatabaseError(_NON_STRING_KEY) + size = ctrl_byte & 0x1F + if size >= 29: + (size, new_offset) = self._size_from_ctrl_byte(size, new_offset) + remaining = budget.payload_left - size + if remaining < 0: + raise InvalidDatabaseError(_TOO_LARGE) + budget.payload_left = remaining + end = new_offset + size + if end > data_end: + raise InvalidDatabaseError(_BAD_DATA) + key = buffer[new_offset:end].decode("utf-8") + if is_pointer: + # The map continues after the pointer, not after its target. + return key, pointer_end + return key, end + def _decode_pointer( self, size: int, @@ -204,10 +262,9 @@ def _decode_pointer( ) -> tuple[Record, int]: pointer_size = (size >> 3) + 1 new_offset = offset + pointer_size - pointer_bytes = self._buffer[offset:new_offset] - if len(pointer_bytes) != pointer_size: + if new_offset > self._data_end: raise InvalidDatabaseError(_BAD_DATA) - pointer = int.from_bytes(pointer_bytes, "big") + pointer = int.from_bytes(self._buffer[offset:new_offset], "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. @@ -218,6 +275,11 @@ def _decode_pointer( if self._pointer_test: return pointer, new_offset + # The payload checks do not cover a target with no payload, such as + # an empty map. + if pointer >= self._data_end: + raise InvalidDatabaseError(_BAD_DATA) + # 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 @@ -239,8 +301,9 @@ def _decode_uint( 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 + if new_offset > self._data_end: + raise InvalidDatabaseError(_BAD_DATA) + return int.from_bytes(self._buffer[offset:new_offset], "big"), new_offset def decode(self, offset: int) -> tuple[Record, int]: """Decode a section of the data section starting at offset. @@ -297,6 +360,8 @@ def _decode( # noqa: C901, PLR0911, PLR0912 raise InvalidDatabaseError(_TOO_LARGE) budget.payload_left = remaining end = new_offset + size + if end > self._data_end: + raise InvalidDatabaseError(_BAD_DATA) return self._buffer[new_offset:end].decode("utf-8"), end case 1: if pointer_target: diff --git a/maxminddb/file.py b/maxminddb/file.py index 62ddb62..6295a34 100644 --- a/maxminddb/file.py +++ b/maxminddb/file.py @@ -5,6 +5,8 @@ import os from typing import TYPE_CHECKING, overload +from maxminddb.errors import InvalidDatabaseError + try: from multiprocessing import Lock except ImportError: @@ -13,6 +15,13 @@ if TYPE_CHECKING: from maxminddb.types import StrOrBytesPath +# The readers check offsets against the size at open. A short read means that +# the file is corrupt or that it shrank after open. +_SHORT_READ = ( + "The read went past the end of the database file. The file is corrupt, or " + "it changed after the reader opened it." +) + class FileBuffer: """A slice-able file reader.""" @@ -33,9 +42,16 @@ def __getitem__(self, index: slice) -> bytes: ... def __getitem__(self, index: slice | int) -> bytes | int: """Get item by index.""" if isinstance(index, slice): - return self._read(index.stop - index.start, index.start) + size = index.stop - index.start + data = self._read(size, index.start) + if len(data) != size: + raise InvalidDatabaseError(_SHORT_READ) + return data if isinstance(index, int): - return self._read(1, index)[0] + data = self._read(1, index) + if not data: + raise InvalidDatabaseError(_SHORT_READ) + return data[0] msg = "Invalid argument type." raise TypeError(msg) @@ -46,7 +62,7 @@ def rfind(self, needle: bytes, start: int) -> int: return pos return start + pos - def size(self) -> int: + def __len__(self) -> int: """Size of file.""" return self._size diff --git a/maxminddb/reader.py b/maxminddb/reader.py index 8cf65dc..e6c1a40 100644 --- a/maxminddb/reader.py +++ b/maxminddb/reader.py @@ -51,7 +51,8 @@ class Reader: _METADATA_START_MARKER = b"\xab\xcd\xefMaxMind.com" _buffer: bytes | bytearray | FileBuffer | "mmap.mmap" # noqa: UP037 - _buffer_size: int + # The offset of the metadata marker, where the data section ends. + _data_end: int # No database is open until __init__ succeeds. closed: bool = True _decoder: Decoder @@ -125,7 +126,7 @@ def _load( filename = self._load_buffer(database, mode) metadata_start = self._buffer.rfind( self._METADATA_START_MARKER, - max(0, self._buffer_size - 128 * 1024), + max(0, len(self._buffer) - 128 * 1024), ) if metadata_start == -1: @@ -137,6 +138,7 @@ def _load( msg, ) + self._data_end = metadata_start metadata_start += len(self._METADATA_START_MARKER) metadata_decoder = Decoder(self._buffer, metadata_start) # For a repeated key, the decoder keeps the last value, but @@ -167,14 +169,13 @@ def _load( # Traversal reads nodes below node_count. Once the tree fits, those # reads need no length checks of their own. - if self._data_start > 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 + _check_search_tree_fits( + self._data_start, self._data_end, len(self._buffer), filename + ) - self._decoder = Decoder(self._buffer, self._data_start) + self._decoder = Decoder( + self._buffer, self._data_start, data_end=self._data_end + ) self.closed = False ipv4_start = 0 @@ -362,9 +363,9 @@ def _read_node(self, node_number: int, index: int) -> int: def _resolve_data_pointer(self, pointer: int) -> Record: resolved = pointer + self._pointer_base - # A pointer into the separator between the tree and the data section - # is as corrupt as one past the end, as libmaxminddb checks. - if resolved < self._data_start or resolved >= self._buffer_size: + # A record must point into the data section. libmaxminddb also rejects + # a pointer into the separator between the tree and the data section. + if resolved < self._data_start or resolved >= self._data_end: raise InvalidDatabaseError(_CORRUPT_TREE) (data, _) = self._decoder.decode(resolved) @@ -414,15 +415,11 @@ def _load_path(self, database: StrOrBytesPath | int, mode: int) -> object: if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP: with open(database, "rb") as db_file: self._buffer = mmap.mmap(db_file.fileno(), 0, access=mmap.ACCESS_READ) - self._buffer_size = self._buffer.size() elif mode in (MODE_AUTO, MODE_FILE): self._buffer = FileBuffer(database) - self._buffer_size = self._buffer.size() elif mode == MODE_MEMORY: with open(database, "rb") as db_file: - buf = db_file.read() - self._buffer = buf - self._buffer_size = len(buf) + self._buffer = db_file.read() else: msg = f"Unsupported open mode ({mode}) for a path or a file descriptor." raise ValueError(msg) @@ -444,7 +441,6 @@ def _load_file_object(self, database: object, read: Callable[[], object]) -> obj msg += " Open it in binary mode." raise TypeError(msg) self._buffer = buf - self._buffer_size = len(buf) # io buffers are not guaranteed to have a name attribute if hasattr(database, "name"): return database.name @@ -501,10 +497,6 @@ def _metadata_fields(metadata: RecordDict, filename: object) -> dict[str, Any]: A new minor version of the format can add keys. This ignores them. """ prefix = f"Error reading metadata in database file ({filename})." - # The C extension also rejects a key that is not a string. - if not all(type(k) is str for k in metadata): - msg = f"{prefix} A metadata key is not a string." - raise InvalidDatabaseError(msg) fields: dict[str, Any] = {} for key, value_type in _METADATA_TYPES.items(): value = metadata.get(key) @@ -513,7 +505,7 @@ def _metadata_fields(metadata: RecordDict, filename: object) -> dict[str, Any]: if valid and isinstance(value, list): valid = all(type(v) is str for v in value) elif valid and isinstance(value, dict): - valid = all(type(k) is str and type(v) is str for k, v in value.items()) + valid = all(type(v) is str for v in value.values()) if not valid: msg = f"{prefix} The {key} value is missing or has the wrong type." raise InvalidDatabaseError(msg) @@ -549,6 +541,23 @@ def _check_metadata_ranges(fields: dict[str, Any], prefix: str) -> None: raise InvalidDatabaseError(msg) +def _check_search_tree_fits( + data_start: int, + data_end: int, + file_size: int, + filename: object, +) -> None: + """Raise InvalidDatabaseError if the search tree overlaps the metadata.""" + if data_start <= data_end: + return + if data_start > file_size: + problem = "extends past the end of the file" + else: + problem = "extends into the metadata section" + msg = f"Error opening database file ({filename}). The search tree {problem}." + raise InvalidDatabaseError(msg) + + @dataclass(kw_only=True, frozen=True) class Metadata: """Metadata for the MaxMind DB reader.""" diff --git a/tests/decoder_test.py b/tests/decoder_test.py index ceb80e8..22691f8 100644 --- a/tests/decoder_test.py +++ b/tests/decoder_test.py @@ -1,13 +1,18 @@ from __future__ import annotations +import contextlib import mmap +import os +import pathlib import sys +import tempfile import threading import unittest from typing import TYPE_CHECKING, Any, ClassVar, SupportsIndex from maxminddb.decoder import Decoder from maxminddb.errors import InvalidDatabaseError +from maxminddb.file import FileBuffer if TYPE_CHECKING: from _typeshed import SizedBuffer @@ -316,18 +321,59 @@ def test_value_limit_follows_the_flat_rule(self) -> None: with self.assertRaisesRegex(InvalidDatabaseError, _TOO_MANY_VALUES): Decoder(self._scalar_pointer_array(65_536), pointer_base=0).decode(1) - def test_map_key_that_cannot_be_hashed_is_rejected(self) -> None: - # A map with one entry, whose key is the array [1]. - decoder = Decoder(bytes.fromhex("e10104a1014178")) - with self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"): - decoder.decode(0) - 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_pointer_past_data_end_is_rejected(self) -> None: + for buf, offset, data_end, expected in ( + # The pointer targets the empty map at offset 2. An empty map has + # no payload, so only the target check can reject it. + (self._pointer(2) + b"\xe0", 0, 2, ({}, 2)), + # The pointer at offset 2 targets the string "a" at offset 0, but + # its offset byte is past data_end. + (b"\x41a" + self._pointer(0), 2, 3, ("a", 4)), + ): + with self.subTest(buf=buf.hex()): + self.assertEqual(expected, Decoder(buf).decode(offset)) + with self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"): + Decoder(buf, data_end=data_end).decode(offset) + + def test_map_key_past_data_end_is_rejected(self) -> None: + for buf, data_end, expected in ( + # A map whose inline key "ab" ends past data_end. Its value, an + # empty map, has no payload to check. + ("e1426162e0", 3, ({"ab": {}}, 5)), + # A map whose key is a pointer to the string "a" at offset 4. + ("e12004a04161", 4, ({"a": 0}, 4)), + ): + with self.subTest(buf=buf): + self.assertEqual(expected, Decoder(bytes.fromhex(buf)).decode(0)) + with self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"): + Decoder(bytes.fromhex(buf), data_end=data_end).decode(0) + + def test_read_from_a_shrunk_file_is_rejected(self) -> None: + # Offset 0 holds an empty map, and the pointer at offset 1 targets the + # string at offset 3. After the file shrinks to 2 bytes, a short read + # of the pointer byte would give a pointer to offset 0. After it + # shrinks to 1 byte, the pointer's ctrl byte is past the end. + data = b"\xe0" + self._pointer(3) + b"\x41a" + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "data" + path.write_bytes(data) + with contextlib.closing(FileBuffer(str(path))) as file_buffer: + decoder = Decoder(file_buffer) + self.assertEqual(("a", 3), decoder.decode(1)) + for size in (2, 1): + os.truncate(path, size) + with ( + self.subTest(size=size), + self.assertRaisesRegex(InvalidDatabaseError, "changed after"), + ): + decoder.decode(1) + def test_cyclic_pointer_raises(self) -> None: with self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"): Decoder(self._pointer(0)).decode(0) @@ -503,6 +549,53 @@ def test_truncated_data_raises_invalid_database_error(self) -> None: with self.assertRaisesRegex(InvalidDatabaseError, "bad data"): Decoder(truncated, pointer_base=0).decode(0) + def test_truncated_payload_raises_invalid_database_error(self) -> None: + for truncated in ( + "c4ffff", # uint32 declaring 4 bytes + "0303ffff", # uint128 declaring 3 bytes + "456162", # string declaring 5 bytes + "856162", # bytes declaring 5 bytes + "0301ff", # int32 declaring 3 bytes + "040100ffff", # int32 declaring 4 bytes + "680000", # double + "040800", # float + "e120", # map key pointer missing its byte + ): + self.assert_bad_data_in_each_buffer(bytes.fromhex(truncated)) + + def assert_bad_data_in_each_buffer(self, data: bytes) -> None: + # A slice of bytes or an mmap past its end returns fewer bytes, and + # FileBuffer raises its own error. Check that the default data_end of + # each buffer type stops the decode first. + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "data" + path.write_bytes(data) + with ( + contextlib.closing(FileBuffer(str(path))) as file_buffer, + mmap.mmap(-1, len(data)) as mm, + ): + mm.write(data) + for buffer in (data, mm, file_buffer): + with ( + self.subTest(data=data.hex(), buffer=type(buffer).__name__), + self.assertRaisesRegex(InvalidDatabaseError, "bad data"), + ): + Decoder(buffer, pointer_base=0).decode(0) + + def test_non_string_map_key_raises_invalid_database_error(self) -> None: + for bad_key in ( + "e18141a101", # bytes key + "e1a101a101", # uint16 key + "e10104a1014178", # array key, which cannot be hashed + "e12003a0a101", # pointer to a uint16 key + "e12004a02000", # pointer to a pointer key + ): + with ( + self.subTest(bad_key=bad_key), + self.assertRaisesRegex(InvalidDatabaseError, "not a string"), + ): + Decoder(bytes.fromhex(bad_key)).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 diff --git a/tests/reader_test.py b/tests/reader_test.py index 74a9144..378ec3d 100644 --- a/tests/reader_test.py +++ b/tests/reader_test.py @@ -150,12 +150,20 @@ def _database_with_metadata( return data[:start] + _encode_control(7, len(entries)) + items -def _database(records: tuple[int, ...], *, ip_version: int) -> bytes: +def _database( + records: tuple[int, ...], + *, + ip_version: int, + data: bytes | None = None, +) -> bytes: """Return a database with a 24-bit search tree and one data record. The tree has two records per node, and a record of node_count + 16 points - at the data record, the string "net". + at the data record. data is the encoded data section, by default the + string "net". """ + if data is None: + data = _encode_value("net") metadata = { "binary_format_major_version": 2, "binary_format_minor_version": 0, @@ -168,13 +176,7 @@ def _database(records: tuple[int, ...], *, ip_version: int) -> bytes: "record_size": 24, } tree = b"".join(record.to_bytes(3, "big") for record in records) - return ( - tree - + bytes(16) - + _encode_value("net") - + _METADATA_START_MARKER - + _encode_value(metadata) - ) + return tree + bytes(16) + data + _METADATA_START_MARKER + _encode_value(metadata) def _encode_value(value: object, key: str = "") -> bytes: @@ -813,6 +815,8 @@ def test_metadata_that_does_not_decode_is_rejected(self) -> None: cm.exception.__cause__, UnicodeDecodeError, ) + else: + self.assertIn("not a string", str(cm.exception)) def test_exhausted_iterator_stops_after_close(self) -> None: with open_database( @@ -962,6 +966,75 @@ def test_cyclic_search_tree_is_rejected(self) -> None: ): list(reader) + def test_value_past_data_section_is_rejected(self) -> None: + if self.reader_class is not maxminddb.reader.Reader: + self.skipTest("libmaxminddb's data section bound includes the metadata") + # One node whose records both point to the start of the data section. + # The value is the whole data section, so its payload overruns into + # the metadata marker. + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "overrun.mmdb" + for value in ( + "c4ffff", # uint32 declaring 4 bytes + "0303ffff", # uint128 declaring 3 bytes + "456162", # string declaring 5 bytes + "856162", # bytes declaring 5 bytes + "0301ff", # int32 declaring 3 bytes + "680000", # double + "040800", # float + "2800", # pointer declaring 2 bytes + ): + data = bytes.fromhex(value) + path.write_bytes(_database((17, 17), ip_version=4, data=data)) + with ( + self.subTest(value=value), + open_database(str(path), self.mode) as reader, + self.assertRaisesRegex(InvalidDatabaseError, "contains bad data"), + ): + reader.get(self.ipf("1.1.1.1")) + + def test_record_that_points_into_the_metadata_is_rejected(self) -> None: + if self.reader_class is not maxminddb.reader.Reader: + self.skipTest("libmaxminddb's data section bound includes the metadata") + data = bytearray(_database((17, 17), ip_version=4)) + # binary_format_minor_version is a uint16 with no payload, so only the + # record check can reject it. + key = _encode_value("binary_format_minor_version") + target = data.index(key) + len(key) + self.assertEqual(data[target], 0xA0) + # A record resolves to itself plus the tree size minus node_count. + data[:6] = (target - 5).to_bytes(3, "big") * 2 + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "metadata-record.mmdb" + path.write_bytes(data) + with ( + open_database(str(path), self.mode) as reader, + self.assertRaisesRegex(InvalidDatabaseError, "search tree is corrupt"), + ): + reader.get(self.ipf("1.1.1.1")) + + def test_search_tree_that_overlaps_the_metadata_is_rejected(self) -> None: + if self.reader_class is not maxminddb.reader.Reader: + self.skipTest("libmaxminddb's data section bound includes the metadata") + # The file holds a 6-byte tree, but node_count 3 needs 18 bytes. The + # tree and separator then end past the metadata marker, but not past + # the end of the file. + node_count = _encode_value("node_count") + data = _database((17, 17), ip_version=4) + self.assertEqual(data.count(node_count + b"\xc1\x01"), 1) + data = data.replace(node_count + b"\xc1\x01", node_count + b"\xc1\x03") + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "overlap.mmdb" + path.write_bytes(data) + with ( + self.assertRaisesRegex( + InvalidDatabaseError, + "The search tree extends into the metadata section", + ), + open_database(str(path), self.mode), + ): + pass + def test_ip_validation(self) -> None: reader = open_database( "tests/data/test-data/MaxMind-DB-test-decoder.mmdb", @@ -995,6 +1068,24 @@ def test_database_with_invalid_utf8_key(self) -> None: with self.assertRaises(UnicodeDecodeError): reader.get_with_prefix_len("163.254.149.39") + def test_non_string_map_key_is_rejected(self) -> None: + data = pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" + ).read_bytes() + # Each record map refers to the shared "ip" key through a pointer. + self.assertEqual(data.count(b"\x42ip"), 1) + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "non-string-map-key.mmdb" + # Change the key type from string (0x42) to bytes or uint16. + for ctrl in (b"\x82", b"\xa2"): + path.write_bytes(data.replace(b"\x42ip", ctrl + b"ip")) + with ( + self.subTest(ctrl=ctrl.hex()), + open_database(str(path), self.mode) as reader, + self.assertRaisesRegex(InvalidDatabaseError, "not a string"), + ): + reader.get(self.ipf("1.1.1.1")) + def test_too_many_constructor_args(self) -> None: with self.assertRaises(TypeError): self.reader_class("README.md", self.mode, 1) # type: ignore[arg-type,call-arg] @@ -1372,21 +1463,6 @@ class TestExtensionReader(BaseTestReader): if has_maxminddb_extension(): reader_class = maxminddb.extension.Reader - def test_map_key_that_is_not_a_string_is_rejected(self) -> None: - data = bytearray( - pathlib.Path(f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb").read_bytes(), - ) - # Change the type of the "ip" key from a string to a uint16. - data[data.index(b"\x42ip")] = 0xA2 - with tempfile.TemporaryDirectory() as directory: - path = pathlib.Path(directory) / "int-key.mmdb" - path.write_bytes(data) - with ( - maxminddb.extension.Reader(path) as reader, - self.assertRaisesRegex(InvalidDatabaseError, "not a string"), - ): - reader.get("1.1.1.1") - def test_invalid_utf8_key_does_not_leak(self) -> None: def fail_to_decode(count: int) -> None: for _ in range(count): @@ -2218,8 +2294,7 @@ def test_metadata_decode_error_names_the_file(self) -> None: self.assertEqual( str(cm.exception), f"Error reading metadata in database file ({path}). " - "The MaxMind DB file's data section contains bad data " - "(unknown data type or corrupt data)", + "Invalid map key: the key is not a string.", ) def test_failed_initialization_closes_buffer(self) -> None: