diff --git a/HISTORY.rst b/HISTORY.rst index 50839c9a..8f08019a 100644 --- a/HISTORY.rst +++ b/HISTORY.rst @@ -3,6 +3,73 @@ History ------- +3.3.0 +++++++++++++++++++ + +* Fixed iteration over an IPv6 database with a network shorter than /96 + whose first bits are zero, such as ``::/1``. The readers raised + ``ValueError`` or skipped networks. +* The pure Python reader now raises ``InvalidDatabaseError`` for a search tree + record that points before the data section. Previously, it returned an + empty map. +* Iterating over a database with a corrupt search tree, such as one with a + cycle, now raises ``InvalidDatabaseError``. Previously, the C extension + could corrupt memory, and the pure Python reader raised ``RecursionError`` + or returned part of the networks. + +* C extension: + + * Fixed segmentation faults from invalid use of ``Metadata``, ``Reader`` and + the internal iterator type. + * Fixed a segmentation fault on a database with a map key that is not a + string. Such a database now raises ``InvalidDatabaseError``. + * Fixed large ``uint32`` values, which came back negative on Windows. + * Fixed memory leaks and a use-after-free. Reinitializing a ``Reader`` or a + ``Metadata`` now raises ``ValueError``. + * Fixed a deadlock on free-threaded Python when a ``Reader`` was closed + during iteration, from another thread or from a signal handler. + * Fixed a crash on free-threaded Python when two threads advanced the same + iterator. + * An exhausted iterator now raises ``StopIteration`` after its ``Reader`` + closes, not ``ValueError``. + * The iterator now stops after it raises an error, as the pure Python + iterator does. + * Added the ``node_byte_size`` and ``search_tree_size`` properties to + ``Metadata``, as the pure Python ``Metadata`` has. + +* Metadata: + + * The pure Python reader ignores unknown keys, which a new minor version of + the format can add. It raises ``InvalidDatabaseError`` for a missing key, a + value of the wrong type, an invalid ``ip_version`` or format version, or a + ``build_epoch`` of 0. + * The C extension ignores unknown keys. Previously, ``Reader.metadata()`` + crashed on them. + +* Added ``maxminddb.Mode``, which the README already described. +* ``maxminddb.types.Record`` and ``Primitive`` are no longer generic type + aliases. Remove any subscript, such as ``Record[str]``. Pull request by Adam + Hitchcock. GitHub #464. + +* Type hints: + + * ``Primitive`` includes ``bytearray``, which the C extension returns for + the ``bytes`` type. + * ``Reader.__iter__`` declares its item type. + * Added the ``StrOrBytesPath`` and ``DatabaseSource`` aliases to + ``maxminddb.types``. + * The ``maxminddb.extension.Reader`` stub accepts only a path, as the + extension does. + * ``MODE_FD`` accepts any object whose ``read()`` method returns ``bytes``, + such as a ``gzip.GzipFile``. ``maxminddb.types.SupportsRead`` describes + this type. + +* ``MODE_AUTO`` now accepts a binary file object. Previously, this could + raise ``TypeError``. It reads the file object into memory with the pure + Python reader, as ``MODE_FD`` does, so lookups are slower than with a path. +* The pure Python reader raises ``TypeError`` when the database argument does + not suit the mode. + 3.2.0 (2026-09-10) ++++++++++++++++++ diff --git a/README.rst b/README.rst index ccdc4ebc..e4df5544 100644 --- a/README.rst +++ b/README.rst @@ -39,7 +39,7 @@ provide `free GeoLite databases files must be decompressed with ``gunzip``. After you have obtained a database and imported the module, call -``open_database`` with a path, or file descriptor (in the case of ``Mode.FD``), +``open_database`` with a path, or binary file object (with ``Mode.FD`` or ``Mode.AUTO``), to the database as the first argument. Optionally, you may pass a mode as the second argument. The modes are available from ``maxminddb.Mode``. Valid modes are: @@ -47,13 +47,14 @@ second argument. The modes are available from ``maxminddb.Mode``. Valid modes ar * ``Mode.MMAP`` - read from memory map. Pure Python. * ``Mode.FILE`` - read database as standard file. Pure Python. * ``Mode.MEMORY`` - load database into memory. Pure Python. -* ``Mode.FD`` - load database into memory from a file descriptor. Pure Python. +* ``Mode.FD`` - load database into memory from a binary file object. Pure Python. * ``Mode.AUTO`` - try ``Mode.MMAP_EXT``, ``Mode.MMAP``, ``Mode.FILE`` in that - order. Default. + order. A file object is read into memory with the pure Python reader, as + with ``Mode.FD``. Pass a path to use the faster C extension. Default. -**NOTE**: When using ``Mode.FD``, it is the *caller's* responsibility to be -sure that the file descriptor gets closed properly. The caller may close the -file descriptor immediately after the ``Reader`` object is created. +**NOTE**: When using a file object, it is the *caller's* responsibility to be +sure that the file object gets closed properly. The caller may close the +file object immediately after the ``Reader`` object is created. The ``open_database`` function returns a ``Reader`` object. To look up an IP address, use the ``get`` method on this object. The method will return the diff --git a/docs/index.rst b/docs/index.rst index b062a0d9..5eca7cb3 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -35,6 +35,14 @@ Database Reader :undoc-members: :show-inheritance: +===== +Types +===== + +.. automodule:: maxminddb.types + :members: + :undoc-members: + ================== Indices and tables ================== diff --git a/extension/maxminddb.c b/extension/maxminddb.c index c2dd73e7..9fb43cda 100644 --- a/extension/maxminddb.c +++ b/extension/maxminddb.c @@ -129,6 +129,8 @@ static inline maxminddb_state *get_maxminddb_state_from_self(PyObject *self) { static bool can_read(const char *path); static int get_record(PyObject *self, PyObject *args, PyObject **record); +static PyObject *reader_iter_next(PyObject *self); +static void free_records(struct record *next); static bool format_sockaddr(struct sockaddr *addr, char *dst); static PyObject *from_entry_data_list(maxminddb_state *state, MMDB_entry_data_list_s **entry_data_list); @@ -308,6 +310,28 @@ static void reader_release_write_lock(Reader_obj *reader) { // Reader implementation // ============================================================================= +static PyObject * +Reader_new(PyTypeObject *type, PyObject *UNUSED(args), PyObject *UNUSED(kwds)) { + PyObject *self = type->tp_alloc(type, 0); + if (self == NULL) { + return NULL; + } + + // Initialize the lock once, for the whole lifetime of the object. + // Reader_dealloc destroys it. A bare __new__ or a failed Reader_init then + // still leaves a valid, unlocked lock, so no path uses or destroys an + // uninitialized lock. + if (reader_lock_init(&((Reader_obj *)self)->rwlock) != 0) { + // Skip Reader_dealloc, which would destroy the failed lock. An + // instance of a heap type holds a reference to its type. + PyObject_Del(self); + Py_DECREF(type); + return NULL; + } + + return self; +} + static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds) { maxminddb_state *state = get_maxminddb_state_from_self(self); if (state == NULL) { @@ -349,31 +373,36 @@ static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds) { return -1; } - MMDB_s *mmdb = (MMDB_s *)malloc(sizeof(MMDB_s)); - if (mmdb == NULL) { + Reader_obj *mmdb_obj = (Reader_obj *)self; + if (reader_acquire_write_lock(mmdb_obj) != 0) { Py_XDECREF(filepath); - PyErr_NoMemory(); return -1; } - Reader_obj *mmdb_obj = (Reader_obj *)self; - if (!mmdb_obj) { + // Refuse a second init. closed is NULL until the first init or close(), + // and the write lock stops two threads from both passing this check. A + // second init would leak the open database. After close(), it would leave + // an existing iterator pointing at freed memory. + if (mmdb_obj->closed != NULL) { + reader_release_write_lock(mmdb_obj); Py_XDECREF(filepath); - free(mmdb); - PyErr_NoMemory(); + PyErr_SetString(PyExc_ValueError, + "Attempt to reinitialize a MaxMind DB reader."); return -1; } - if (reader_lock_init(&mmdb_obj->rwlock) != 0) { - free(mmdb); + MMDB_s *mmdb = (MMDB_s *)malloc(sizeof(MMDB_s)); + if (mmdb == NULL) { + reader_release_write_lock(mmdb_obj); Py_XDECREF(filepath); + PyErr_NoMemory(); return -1; } int const status = MMDB_open(filename, MMDB_MODE_MMAP, mmdb); if (status != MMDB_SUCCESS) { - reader_lock_destroy(&mmdb_obj->rwlock); + reader_release_write_lock(mmdb_obj); free(mmdb); PyErr_Format(state->MaxMindDB_error, "Error opening database file (%s). Is this a valid " @@ -387,6 +416,7 @@ static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds) { mmdb_obj->mmdb = mmdb; mmdb_obj->closed = Py_False; + reader_release_write_lock(mmdb_obj); return 0; } @@ -604,6 +634,18 @@ static bool format_sockaddr(struct sockaddr *sa, char *dst) { return false; } +// The keys that Metadata accepts, in argument order. +static char *metadata_keys[] = {"binary_format_major_version", + "binary_format_minor_version", + "build_epoch", + "database_type", + "description", + "ip_version", + "languages", + "node_count", + "record_size", + NULL}; + static PyObject *Reader_metadata(PyObject *self, PyObject *UNUSED(args)) { maxminddb_state *state = get_maxminddb_state_from_self(self); if (state == NULL) { @@ -637,24 +679,51 @@ static PyObject *Reader_metadata(PyObject *self, PyObject *UNUSED(args)) { PyObject *metadata_dict = from_entry_data_list(state, &entry_data_list); MMDB_free_entry_data_list(original_entry_data_list); - if (metadata_dict == NULL || !PyDict_Check(metadata_dict)) { + if (metadata_dict == NULL) { + reader_release_read_lock(mmdb_obj); + return NULL; + } + if (!PyDict_Check(metadata_dict)) { reader_release_read_lock(mmdb_obj); PyErr_SetString(state->MaxMindDB_error, "Error decoding metadata."); + Py_DECREF(metadata_dict); return NULL; } reader_release_read_lock(mmdb_obj); - PyObject *args = PyTuple_New(0); + // A newer minor version of the format can add metadata keys. Pass only + // the keys that Metadata accepts. + PyObject *args = PyTuple_New(Py_ARRAY_LENGTH(metadata_keys) - 1); if (args == NULL) { Py_DECREF(metadata_dict); return NULL; } - - PyObject *metadata = - PyObject_Call(state->Metadata_Type, args, metadata_dict); - + for (Py_ssize_t i = 0; metadata_keys[i] != NULL; i++) { + PyObject *key = PyUnicode_FromString(metadata_keys[i]); + if (key == NULL) { + Py_DECREF(args); + Py_DECREF(metadata_dict); + return NULL; + } + PyObject *value = PyDict_GetItemWithError(metadata_dict, key); + Py_DECREF(key); + if (value == NULL) { + if (!PyErr_Occurred()) { + PyErr_Format(state->MaxMindDB_error, + "Error decoding metadata. The %s value is " + "missing.", + metadata_keys[i]); + } + Py_DECREF(args); + Py_DECREF(metadata_dict); + return NULL; + } + PyTuple_SET_ITEM(args, i, Py_NewRef(value)); + } Py_DECREF(metadata_dict); + + PyObject *metadata = PyObject_CallObject(state->Metadata_Type, args); Py_DECREF(args); return metadata; } @@ -686,7 +755,7 @@ static PyObject *Reader__enter__(PyObject *self, PyObject *UNUSED(args)) { return NULL; } - if (mmdb_obj->closed == Py_True) { + if (mmdb_obj->mmdb == NULL) { reader_release_read_lock(mmdb_obj); PyErr_SetString(PyExc_ValueError, "Attempt to reopen a closed MaxMind DB."); @@ -700,19 +769,29 @@ static PyObject *Reader__enter__(PyObject *self, PyObject *UNUSED(args)) { } static PyObject *Reader__exit__(PyObject *self, PyObject *UNUSED(args)) { - Reader_close(self, NULL); - Py_RETURN_NONE; + return Reader_close(self, NULL); } static void Reader_dealloc(PyObject *self) { Reader_obj *obj = (Reader_obj *)self; if (obj->mmdb != NULL) { - Reader_close(self, NULL); + PyObject *result = Reader_close(self, NULL); + if (result == NULL) { + // dealloc cannot raise, so report the error and continue. Pass + // NULL, because the hook would take a reference to self, whose + // count is already 0. + PyErr_WriteUnraisable(NULL); + } else { + Py_DECREF(result); + } } reader_lock_destroy(&obj->rwlock); + // An instance of a heap type holds a reference to its type. + PyTypeObject *type = Py_TYPE(self); PyObject_Del(self); + Py_DECREF(type); } static PyObject *Reader_iter(PyObject *obj) { @@ -727,7 +806,7 @@ static PyObject *Reader_iter(PyObject *obj) { return NULL; } - if (reader->closed == Py_True) { + if (reader->mmdb == NULL) { reader_release_read_lock(reader); PyErr_SetString(PyExc_ValueError, "Attempt to iterate over a closed MaxMind DB."); @@ -766,6 +845,27 @@ static bool is_ipv6(char ip[16]) { } static PyObject *ReaderIter_next(PyObject *self) { + PyObject *result; +#ifdef Py_GIL_DISABLED + // The iterator's list of pending records is not thread-safe, so let only + // one thread at a time advance an iterator. The read lock is shared, so + // it does not do this. + Py_BEGIN_CRITICAL_SECTION(self); +#endif + result = reader_iter_next(self); + // Stop after an error, as a generator does. + if (result == NULL && PyErr_Occurred()) { + ReaderIter_obj *ri = (ReaderIter_obj *)self; + free_records(ri->next); + ri->next = NULL; + } +#ifdef Py_GIL_DISABLED + Py_END_CRITICAL_SECTION(); +#endif + return result; +} + +static PyObject *reader_iter_next(PyObject *self) { maxminddb_state *state = get_maxminddb_state_from_self((PyObject *)self); if (state == NULL) { return NULL; @@ -773,11 +873,16 @@ static PyObject *ReaderIter_next(PyObject *self) { ReaderIter_obj *ri = (ReaderIter_obj *)self; + // An exhausted iterator stays exhausted, even after the reader closes. + if (ri->next == NULL) { + return NULL; + } + if (reader_acquire_read_lock(ri->reader) != 0) { return NULL; } - if (ri->reader->closed == Py_True) { + if (ri->reader->mmdb == NULL) { reader_release_read_lock(ri->reader); PyErr_SetString(PyExc_ValueError, "Attempt to iterate over a closed MaxMind DB."); @@ -791,8 +896,11 @@ static PyObject *ReaderIter_next(PyObject *self) { switch (cur->type) { case MMDB_RECORD_TYPE_INVALID: reader_release_read_lock(ri->reader); + // libmaxminddb before 1.14 returns this type for a record + // that points to the root. Later versions fail in + // MMDB_read_node instead. PyErr_SetString(state->MaxMindDB_error, - "Invalid record when reading node"); + "The MaxMind DB file's search tree is corrupt"); free(cur); return NULL; case MMDB_RECORD_TYPE_SEARCH_NODE: { @@ -802,6 +910,17 @@ static PyObject *ReaderIter_next(PyObject *self) { // These are aliased networks. Skip them. break; } + // A node at the full address depth would write its children + // past the end of ip_packed. Only a corrupt tree, such as one + // with a cycle, has one. + if (cur->depth >= ri->reader->mmdb->depth) { + reader_release_read_lock(ri->reader); + PyErr_SetString(state->MaxMindDB_error, + "The MaxMind DB file's search tree is " + "corrupt"); + free(cur); + return NULL; + } MMDB_search_node_s node; int status = MMDB_read_node( ri->reader->mmdb, (uint32_t)cur->record, &node); @@ -851,6 +970,9 @@ static PyObject *ReaderIter_next(PyObject *self) { case MMDB_RECORD_TYPE_EMPTY: break; case MMDB_RECORD_TYPE_DATA: { + // Read this before any Python code runs, which could close + // the reader. + uint16_t const depth = ri->reader->mmdb->depth; MMDB_entry_data_list_s *entry_data_list = NULL; int status = MMDB_get_entry_data_list(&cur->entry, &entry_data_list); @@ -870,16 +992,22 @@ static PyObject *ReaderIter_next(PyObject *self) { PyObject *record = from_entry_data_list(state, &entry_data_list); MMDB_free_entry_data_list(original_entry_data_list); + + // The rest uses only cur, which this call owns. Release the + // lock before ip_network runs Python code, which could close + // the reader on this thread. + reader_release_read_lock(ri->reader); if (record == NULL) { - reader_release_read_lock(ri->reader); free(cur); return NULL; } int ip_start = 0; - int ip_length = 4; - if (ri->reader->mmdb->depth == 128) { - if (is_ipv6(cur->ip_packed)) { + Py_ssize_t ip_length = 4; + if (depth == 128) { + // A network shorter than /96 is IPv6, even if its first + // 96 bits are zero. + if (is_ipv6(cur->ip_packed) || cur->depth < 96) { // IPv6 address ip_length = 16; } else { @@ -892,37 +1020,28 @@ static PyObject *ReaderIter_next(PyObject *self) { &(cur->ip_packed[ip_start]), ip_length, cur->depth - ip_start * 8); + free(cur); if (network_tuple == NULL) { - reader_release_read_lock(ri->reader); Py_DECREF(record); - free(cur); return NULL; } PyObject *args = PyTuple_Pack(1, network_tuple); Py_DECREF(network_tuple); if (args == NULL) { - reader_release_read_lock(ri->reader); Py_DECREF(record); - free(cur); return NULL; } PyObject *network = PyObject_CallObject(state->ipaddress_ip_network, args); Py_DECREF(args); if (network == NULL) { - reader_release_read_lock(ri->reader); Py_DECREF(record); - free(cur); return NULL; } PyObject *rv = PyTuple_Pack(2, network, record); Py_DECREF(network); Py_DECREF(record); - - reader_release_read_lock(ri->reader); - - free(cur); return rv; } default: @@ -939,18 +1058,23 @@ static PyObject *ReaderIter_next(PyObject *self) { return NULL; } -static void ReaderIter_dealloc(PyObject *self) { - ReaderIter_obj *ri = (ReaderIter_obj *)self; - - Py_DECREF(ri->reader); - - struct record *next = ri->next; +static void free_records(struct record *next) { while (next != NULL) { struct record *cur = next; next = cur->next; free(cur); } +} + +static void ReaderIter_dealloc(PyObject *self) { + ReaderIter_obj *ri = (ReaderIter_obj *)self; + + Py_DECREF(ri->reader); + + free_records(ri->next); + PyTypeObject *type = Py_TYPE(self); PyObject_Del(self); + Py_DECREF(type); } static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) { @@ -959,21 +1083,10 @@ static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) { *build_epoch, *database_type, *description, *ip_version, *languages, *node_count, *record_size; - static char *kwlist[] = {"binary_format_major_version", - "binary_format_minor_version", - "build_epoch", - "database_type", - "description", - "ip_version", - "languages", - "node_count", - "record_size", - NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, - "|OOOOOOOOO", - kwlist, + "OOOOOOOOO", + metadata_keys, &binary_format_major_version, &binary_format_minor_version, &build_epoch, @@ -988,41 +1101,51 @@ static int Metadata_init(PyObject *self, PyObject *args, PyObject *kwds) { Metadata_obj *obj = (Metadata_obj *)self; - obj->binary_format_major_version = binary_format_major_version; - obj->binary_format_minor_version = binary_format_minor_version; - obj->build_epoch = build_epoch; - obj->database_type = database_type; - obj->description = description; - obj->ip_version = ip_version; - obj->languages = languages; - obj->node_count = node_count; - obj->record_size = record_size; - - Py_INCREF(obj->binary_format_major_version); - Py_INCREF(obj->binary_format_minor_version); - Py_INCREF(obj->build_epoch); - Py_INCREF(obj->database_type); - Py_INCREF(obj->description); - Py_INCREF(obj->ip_version); - Py_INCREF(obj->languages); - Py_INCREF(obj->node_count); - Py_INCREF(obj->record_size); + // Refuse a second init, as Reader_init does. Replacing a field would leak + // the old value or free it while a getter uses it. On free-threaded + // builds, the critical section makes the check and the stores atomic. + int status = 0; +#ifdef Py_GIL_DISABLED + Py_BEGIN_CRITICAL_SECTION(self); +#endif + if (obj->binary_format_major_version != NULL) { + PyErr_SetString(PyExc_ValueError, + "Attempt to reinitialize a MaxMind DB Metadata."); + status = -1; + } else { + obj->binary_format_major_version = + Py_NewRef(binary_format_major_version); + obj->binary_format_minor_version = + Py_NewRef(binary_format_minor_version); + obj->build_epoch = Py_NewRef(build_epoch); + obj->database_type = Py_NewRef(database_type); + obj->description = Py_NewRef(description); + obj->ip_version = Py_NewRef(ip_version); + obj->languages = Py_NewRef(languages); + obj->node_count = Py_NewRef(node_count); + obj->record_size = Py_NewRef(record_size); + } +#ifdef Py_GIL_DISABLED + Py_END_CRITICAL_SECTION(); +#endif - return 0; + return status; } static void Metadata_dealloc(PyObject *self) { Metadata_obj *obj = (Metadata_obj *)self; - Py_DECREF(obj->binary_format_major_version); - Py_DECREF(obj->binary_format_minor_version); - Py_DECREF(obj->build_epoch); - Py_DECREF(obj->database_type); - Py_DECREF(obj->description); - Py_DECREF(obj->ip_version); - Py_DECREF(obj->languages); - Py_DECREF(obj->node_count); - Py_DECREF(obj->record_size); + Py_XDECREF(obj->binary_format_major_version); + Py_XDECREF(obj->binary_format_minor_version); + Py_XDECREF(obj->build_epoch); + Py_XDECREF(obj->database_type); + Py_XDECREF(obj->description); + Py_XDECREF(obj->ip_version); + Py_XDECREF(obj->languages); + Py_XDECREF(obj->node_count); + Py_XDECREF(obj->record_size); + PyTypeObject *type = Py_TYPE(self); PyObject_Del(self); + Py_DECREF(type); } static PyObject * @@ -1057,7 +1180,8 @@ from_entry_data_list(maxminddb_state *state, case MMDB_DATA_TYPE_UINT16: return PyLong_FromLong((*entry_data_list)->entry_data.uint16); case MMDB_DATA_TYPE_UINT32: - return PyLong_FromLong((*entry_data_list)->entry_data.uint32); + return PyLong_FromUnsignedLong( + (*entry_data_list)->entry_data.uint32); case MMDB_DATA_TYPE_BOOLEAN: return PyBool_FromLong((*entry_data_list)->entry_data.boolean); case MMDB_DATA_TYPE_UINT64: @@ -1090,12 +1214,21 @@ static PyObject *from_map(maxminddb_state *state, for (i = 0; i < map_size && *entry_data_list; i++) { *entry_data_list = (*entry_data_list)->next; + // libmaxminddb does not check the key type, and the union holds a + // string only for a string key. + if ((*entry_data_list)->entry_data.type != MMDB_DATA_TYPE_UTF8_STRING) { + PyErr_SetString(state->MaxMindDB_error, + "Invalid map key: the key is not a string."); + Py_DECREF(py_obj); + return NULL; + } PyObject *key = PyUnicode_FromStringAndSize( (*entry_data_list)->entry_data.utf8_string, (*entry_data_list)->entry_data.data_size); if (!key) { // PyUnicode_FromStringAndSize will set an appropriate exception // in this case. + Py_DECREF(py_obj); return NULL; } @@ -1107,9 +1240,13 @@ static PyObject *from_map(maxminddb_state *state, Py_DECREF(py_obj); return NULL; } - PyDict_SetItem(py_obj, key, value); + int const status = PyDict_SetItem(py_obj, key, value); Py_DECREF(value); Py_DECREF(key); + if (status < 0) { + Py_DECREF(py_obj); + return NULL; + } } return py_obj; @@ -1256,6 +1393,52 @@ static PyMemberDef Metadata_members[] = { NULL}, {NULL, 0, 0, 0, NULL}}; +static PyObject *Metadata_node_byte_size(PyObject *self, + void *UNUSED(closure)) { + Metadata_obj *obj = (Metadata_obj *)self; + if (obj->record_size == NULL) { + PyErr_SetString(PyExc_AttributeError, "record_size is not set"); + return NULL; + } + PyObject *four = PyLong_FromLong(4); + if (four == NULL) { + return NULL; + } + PyObject *node_byte_size = PyNumber_FloorDivide(obj->record_size, four); + Py_DECREF(four); + return node_byte_size; +} + +static PyObject *Metadata_search_tree_size(PyObject *self, + void *UNUSED(closure)) { + Metadata_obj *obj = (Metadata_obj *)self; + if (obj->node_count == NULL) { + PyErr_SetString(PyExc_AttributeError, "node_count is not set"); + return NULL; + } + PyObject *node_byte_size = Metadata_node_byte_size(self, NULL); + if (node_byte_size == NULL) { + return NULL; + } + PyObject *search_tree_size = + PyNumber_Multiply(obj->node_count, node_byte_size); + Py_DECREF(node_byte_size); + return search_tree_size; +} + +// These match the properties of the pure Python Metadata class. +static PyGetSetDef Metadata_getset[] = {{"node_byte_size", + Metadata_node_byte_size, + NULL, + "The size of a node in bytes.", + NULL}, + {"search_tree_size", + Metadata_search_tree_size, + NULL, + "The size of the search tree.", + NULL}, + {NULL, NULL, NULL, NULL, NULL}}; + // ============================================================================= // Type specs for heap type conversion (PEP 489) // ============================================================================= @@ -1263,6 +1446,7 @@ static PyMemberDef Metadata_members[] = { static PyType_Slot Reader_Type_slots[] = { {Py_tp_doc, "Reader object"}, {Py_tp_dealloc, Reader_dealloc}, + {Py_tp_new, Reader_new}, {Py_tp_init, Reader_init}, {Py_tp_iter, Reader_iter}, {Py_tp_methods, Reader_methods}, @@ -1283,6 +1467,7 @@ static PyType_Slot Metadata_Type_slots[] = { {Py_tp_init, Metadata_init}, {Py_tp_methods, Metadata_methods}, {Py_tp_members, Metadata_members}, + {Py_tp_getset, Metadata_getset}, {0, NULL}, }; @@ -1305,7 +1490,7 @@ static PyType_Slot ReaderIter_Type_slots[] = { static PyType_Spec ReaderIter_Type_spec = { .name = "maxminddb.extension.ReaderIter", .basicsize = sizeof(ReaderIter_obj), - .flags = Py_TPFLAGS_DEFAULT, + .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION, .slots = ReaderIter_Type_slots, }; diff --git a/maxminddb/__init__.py b/maxminddb/__init__.py index b3ef0ef6..54a83249 100644 --- a/maxminddb/__init__.py +++ b/maxminddb/__init__.py @@ -3,7 +3,7 @@ from __future__ import annotations from importlib.metadata import version -from typing import IO, TYPE_CHECKING, cast +from typing import TYPE_CHECKING, cast from .const import ( MODE_AUTO, @@ -12,12 +12,13 @@ MODE_MEMORY, MODE_MMAP, MODE_MMAP_EXT, + Mode, ) -from .decoder import InvalidDatabaseError -from .reader import Reader +from .errors import InvalidDatabaseError +from .reader import _PATH_TYPES, Reader if TYPE_CHECKING: - import os + from .types import DatabaseSource try: from . import extension as _extension @@ -33,29 +34,35 @@ "MODE_MMAP", "MODE_MMAP_EXT", "InvalidDatabaseError", + "Mode", "Reader", "open_database", ] def open_database( - database: str | bytes | int | os.PathLike[str] | os.PathLike[bytes] | IO[bytes], + database: DatabaseSource, mode: int = MODE_AUTO, ) -> Reader: """Open a MaxMind DB database. Arguments: database: A path to a valid MaxMind DB file such as a GeoIP database - file, or a file descriptor in the case of MODE_FD. + file, or a binary file object for MODE_FD or MODE_AUTO. + MODE_MMAP, MODE_FILE and MODE_MEMORY also accept the file + descriptor of a regular file. MODE_MEMORY reads it from its + current offset, the others from the start. The reader + closes it, even when the file is not a valid database. mode: mode to open the database with. Valid mode are: * MODE_MMAP_EXT - use the C extension with memory map. * MODE_MMAP - read from memory map. Pure Python. * MODE_FILE - read database as standard file. Pure Python. * MODE_MEMORY - load database into memory. Pure Python. - * MODE_FD - the param passed via database is a file descriptor, not - a path. This mode implies MODE_MEMORY. + * MODE_FD - the param passed via database is a binary file + object, not a path. This mode implies MODE_MEMORY. * MODE_AUTO - tries MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that - order. Default mode. + order. Reads a file object into memory, as MODE_FD + does. Default mode. """ if mode not in ( @@ -70,22 +77,34 @@ def open_database( raise ValueError(msg) has_extension = _extension and hasattr(_extension, "Reader") - use_extension = has_extension if mode == MODE_AUTO else mode == MODE_MMAP_EXT - if not use_extension: - return Reader(database, mode) - - if not has_extension: + if mode == MODE_MMAP_EXT and not has_extension: msg = "MODE_MMAP_EXT requires the maxminddb.extension module to be available" raise ValueError( msg, ) - # The C type exposes the same API as the Python Reader, so for type - # checking purposes, pretend it is one. (Ideally this would be a subclass - # of, or share a common parent class with, the Python Reader - # implementation.) - return cast("Reader", _extension.Reader(database, mode)) + # The extension accepts only a path, so MODE_AUTO gives a file object to + # the pure Python reader. It still refuses a file descriptor, as before. + # The cast pretends the C reader is the pure Python Reader, which has the + # same API. + if mode in (MODE_AUTO, MODE_MMAP_EXT) and has_extension: + if isinstance(database, _PATH_TYPES): + return cast("Reader", _extension.Reader(database, mode)) + # An object with __index__, such as an int, is a file descriptor. + # Leave a bool to the pure Python reader, which refuses it. + is_descriptor = not isinstance(database, bool) and hasattr( + type(database), "__index__" + ) + if mode == MODE_MMAP_EXT or is_descriptor: + msg = ( + f"The C extension requires a path ({type(database).__name__} " + "given). Use MODE_FD for a file object, or MODE_MMAP for a " + "file descriptor." + ) + raise TypeError(msg) + + return Reader(database, mode) __version__ = version("maxminddb") diff --git a/maxminddb/const.py b/maxminddb/const.py index 0f7e9826..2b1669d7 100644 --- a/maxminddb/const.py +++ b/maxminddb/const.py @@ -10,7 +10,10 @@ class Mode(IntEnum): """ AUTO = 0 - """Try MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that order. Default mode.""" + """Try MODE_MMAP_EXT, MODE_MMAP, MODE_FILE in that order. Default mode. + + A file object uses MODE_FD. + """ MMAP_EXT = 1 """Use the C extension with memory map.""" @@ -25,7 +28,7 @@ class Mode(IntEnum): """Load database into memory. Pure Python.""" FD = 16 - """Database is a file descriptor, not a path. This mode implies MODE_MEMORY.""" + """Database is a binary file object, not a path. This mode implies MODE_MEMORY.""" # Backward compatibility: export both enum members and old-style constants diff --git a/maxminddb/extension.pyi b/maxminddb/extension.pyi index 1a5e88b0..15d3a4c4 100644 --- a/maxminddb/extension.pyi +++ b/maxminddb/extension.pyi @@ -2,12 +2,10 @@ from collections.abc import Iterator from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network -from os import PathLike -from typing import IO, Any from typing_extensions import Self -from maxminddb.types import Record +from maxminddb.types import Record, StrOrBytesPath class Reader: """A C extension implementation of a reader for the MaxMind DB format. @@ -19,14 +17,14 @@ class Reader: def __init__( self, - database: str | bytes | int | PathLike[str] | PathLike[bytes] | IO[bytes], + database: StrOrBytesPath, mode: int = ..., ) -> None: """Reader for the MaxMind DB file format. Arguments: database: A path to a valid MaxMind DB file such as a GeoIP database - file, or a file descriptor in the case of MODE_FD. + file. mode: mode to open the database with. The only supported modes are MODE_AUTO and MODE_MMAP_EXT. @@ -59,7 +57,7 @@ class Reader: def __iter__(self) -> Iterator[tuple[IPv4Network | IPv6Network, Record]]: ... def __enter__(self) -> Self: ... - def __exit__(self, *args) -> None: ... # noqa: ANN002 + def __exit__(self, *args: object) -> None: ... class Metadata: """Metadata for the MaxMind DB reader.""" @@ -113,5 +111,24 @@ class Metadata: The bit size of a record in the search tree. """ - def __init__(self, **kwargs: Any) -> None: # noqa: ANN401 - """Create new Metadata object. kwargs are key/value pairs from spec.""" + def __init__( + self, + binary_format_major_version: int, + binary_format_minor_version: int, + build_epoch: int, + database_type: str, + description: dict[str, str], + ip_version: int, + languages: list[str], + node_count: int, + record_size: int, + ) -> None: + """Create new Metadata object from the metadata fields in the spec.""" + + @property + def node_byte_size(self) -> int: + """The size of a node in bytes.""" + + @property + def search_tree_size(self) -> int: + """The size of the search tree.""" diff --git a/maxminddb/file.py b/maxminddb/file.py index 1901d9ba..62ddb62c 100644 --- a/maxminddb/file.py +++ b/maxminddb/file.py @@ -3,18 +3,21 @@ from __future__ import annotations import os -from typing import overload +from typing import TYPE_CHECKING, overload try: from multiprocessing import Lock except ImportError: from threading import Lock # type: ignore[assignment] +if TYPE_CHECKING: + from maxminddb.types import StrOrBytesPath + class FileBuffer: """A slice-able file reader.""" - def __init__(self, database: str) -> None: + def __init__(self, database: StrOrBytesPath | int) -> None: """Create FileBuffer.""" self._handle = open(database, "rb") # noqa: SIM115 self._size = os.fstat(self._handle.fileno()).st_size @@ -51,7 +54,7 @@ def close(self) -> None: """Close file.""" self._handle.close() - if hasattr(os, "pread"): # type: ignore[attr-defined] + if hasattr(os, "pread"): def _read(self, buffersize: int, offset: int) -> bytes: """Read that uses pread.""" diff --git a/maxminddb/reader.py b/maxminddb/reader.py index b4def914..ed48ef89 100644 --- a/maxminddb/reader.py +++ b/maxminddb/reader.py @@ -9,9 +9,11 @@ import contextlib import ipaddress +import operator +import os from dataclasses import dataclass -from ipaddress import IPv4Address, IPv6Address -from typing import IO, TYPE_CHECKING, Any +from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network +from typing import TYPE_CHECKING, Any from maxminddb.const import MODE_AUTO, MODE_FD, MODE_FILE, MODE_MEMORY, MODE_MMAP from maxminddb.decoder import Decoder @@ -20,13 +22,21 @@ if TYPE_CHECKING: from collections.abc import Iterator - from os import PathLike from typing_extensions import Self - from maxminddb.types import Record + from maxminddb.types import ( + DatabaseSource, + Record, + RecordDict, + StrOrBytesPath, + ) _IPV4_MAX_NUM = 2**32 +# The path types, which the C extension also accepts. +_PATH_TYPES = (str, bytes, os.PathLike) +# The database types that the path modes pass to open(). +_PATH_OR_FD_TYPES = (*_PATH_TYPES, int) class Reader: @@ -45,24 +55,33 @@ class Reader: _metadata: Metadata _record_size: int _ipv4_start: int + _search_tree_size: int + _data_start: int def __init__( self, - database: str | bytes | int | PathLike[str] | PathLike[bytes] | IO[bytes], + database: DatabaseSource, mode: int = MODE_AUTO, ) -> None: """Reader for the MaxMind DB file format. Arguments: database: A path to a valid MaxMind DB file such as a GeoIP database - file, or a file descriptor in the case of MODE_FD. + file, or a binary file object for MODE_FD or MODE_AUTO. + MODE_AUTO, MODE_MMAP, MODE_FILE and MODE_MEMORY also + accept the file descriptor of a regular file. MODE_MEMORY + reads it from its current offset, the others from the + start. The reader closes it, even when the file is not a + valid database. mode: mode to open the database with. Valid mode are: * MODE_MMAP - read from memory map. * MODE_FILE - read database as standard file. * MODE_MEMORY - load database into memory. - * MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default. - * MODE_FD - the param passed via database is a file descriptor, not - a path. This mode implies MODE_MEMORY. + * MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Reads a + file object into memory, as MODE_FD does. + Default. + * MODE_FD - the param passed via database is a binary file + object, not a path. This mode implies MODE_MEMORY. """ filename = self._load_buffer(database, mode) @@ -94,9 +113,7 @@ def __init__( msg, ) - # The MaxMind DB spec fixes these keys and their value types. - fields: dict[str, Any] = metadata - self._metadata = Metadata(**fields) + self._metadata = Metadata(**_metadata_fields(metadata, filename)) self._record_size = self._metadata.record_size if self._record_size not in (24, 28, 32): msg = f"Unknown record size: {self._record_size}" @@ -121,6 +138,11 @@ def __init__( self._buffer, self._metadata.search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE, ) + # _resolve_data_pointer uses these on every lookup. + self._search_tree_size = self._metadata.search_tree_size + self._data_start = ( + self._search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE + ) self.closed = False ipv4_start = 0 @@ -190,27 +212,47 @@ def get_with_prefix_len( return self._resolve_data_pointer(pointer), prefix_len return None, prefix_len - def __iter__(self) -> Iterator: + def __iter__(self) -> Iterator[tuple[IPv4Network | IPv6Network, Record]]: return self._generate_children(0, 0, 0) - def _generate_children(self, node: int, depth: int, ip_acc: int) -> Iterator: - if ip_acc != 0 and node == self._ipv4_start: - # Skip nodes aliased to IPv4 + def _generate_children( + self, + node: int, + depth: int, + ip_acc: int, + ) -> Iterator[tuple[IPv4Network | IPv6Network, Record]]: + node_count = self._metadata.node_count + bits = 128 if self._metadata.ip_version == 6 else 32 + # Skip the IPv4 subtree when an IPv6 address other than ::/96 leads + # to it, as the C extension does. Inside the IPv4 subtree, or in an + # IPv4 tree, a record that points back to it is a cycle. + if ( + node == self._ipv4_start + and bits == 128 + and node < node_count + and ip_acc >> max(depth - 96, 0) != 0 + ): return - node_count = self._metadata.node_count if node > node_count: - bits = 128 if self._metadata.ip_version == 6 else 32 ip_acc <<= bits - depth - if ip_acc <= _IPV4_MAX_NUM and bits == 128: - depth -= 96 - yield ( - ipaddress.ip_network((ip_acc, depth)), - self._resolve_data_pointer( - node, - ), - ) + network: IPv4Network | IPv6Network + if bits == 32: + network = IPv4Network((ip_acc, depth)) + elif depth >= 96 and ip_acc < _IPV4_MAX_NUM: + # An IPv4 network in an IPv6 tree is at least /96, and its + # first 96 bits are zero. + network = IPv4Network((ip_acc, depth - 96)) + else: + network = IPv6Network((ip_acc, depth)) + yield (network, self._resolve_data_pointer(node)) elif node < node_count: + # A node at the full address depth has no valid children, and no + # record can point to the root. Only a corrupt tree, such as one + # with a cycle, has either. + if depth >= bits or (node == 0 and depth > 0): + msg = "The MaxMind DB file's search tree is corrupt" + raise InvalidDatabaseError(msg) left = self._read_node(node, 0) ip_acc <<= 1 depth += 1 @@ -265,9 +307,11 @@ def _read_node(self, node_number: int, index: int) -> int: raise InvalidDatabaseError(msg) def _resolve_data_pointer(self, pointer: int) -> Record: - resolved = pointer - self._metadata.node_count + self._metadata.search_tree_size + resolved = pointer - self._metadata.node_count + self._search_tree_size - if resolved >= self._buffer_size: + # 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: msg = "The MaxMind DB file's search tree is corrupt" raise InvalidDatabaseError(msg) @@ -276,44 +320,89 @@ def _resolve_data_pointer(self, pointer: int) -> Record: def _load_buffer( self, - database: str | bytes | int | PathLike[str] | PathLike[bytes] | IO[bytes], + database: DatabaseSource, mode: int = MODE_AUTO, - ) -> str: - filename: Any + ) -> object: + """Load the database and return a name for it in error messages.""" + if mode not in (MODE_AUTO, MODE_FD, MODE_FILE, MODE_MEMORY, MODE_MMAP): + msg = ( + f"Unsupported open mode ({mode}). Only MODE_AUTO, MODE_MMAP, " + "MODE_FILE, MODE_MEMORY and MODE_FD are supported by the pure " + "Python Reader" + ) + raise ValueError( + msg, + ) + + # bool is an int, but it is never a file descriptor. MODE_FD refuses + # it below, as any object without read(). + if isinstance(database, bool) and mode != MODE_FD: + msg = "Unsupported database type (bool). Pass a path or a file object." + raise TypeError(msg) + # open() also takes an object with __index__, such as numpy.int64, as + # a file descriptor. + if ( + mode != MODE_FD + and not isinstance(database, _PATH_OR_FD_TYPES) + and hasattr(type(database), "__index__") + ): + database = operator.index(database) # type: ignore[arg-type] + # A path wins over read() in the path modes, because some path objects + # also have a text read(). MODE_FD reads any object with read(). + if mode != MODE_FD and isinstance(database, _PATH_OR_FD_TYPES): + return self._load_path(database, mode) + # MODE_AUTO reads a file object into memory, as MODE_FD does. + if mode == MODE_FD or ( + mode == MODE_AUTO and callable(getattr(database, "read", None)) + ): + return self._load_file_object(database) + if mode == MODE_AUTO: + hint = "Pass a path or a binary file object." + else: + hint = "Use MODE_FD for a file object." + msg = f"Unsupported database type ({type(database).__name__}). {hint}" + raise TypeError(msg) + + def _load_path(self, database: StrOrBytesPath | int, mode: int) -> object: + """Memory-map or read a path or file descriptor.""" if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP: - with open(database, "rb") as db_file: # type: ignore[arg-type] + 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() - filename = database elif mode in (MODE_AUTO, MODE_FILE): - self._buffer = FileBuffer(database) # type: ignore[arg-type] + self._buffer = FileBuffer(database) self._buffer_size = self._buffer.size() - filename = database - elif mode == MODE_MEMORY: - with open(database, "rb") as db_file: # type: ignore[arg-type] + else: + with open(database, "rb") as db_file: buf = db_file.read() self._buffer = buf self._buffer_size = len(buf) - filename = database - elif mode == MODE_FD: - self._buffer = database.read() # type: ignore[union-attr] - self._buffer_size = len(self._buffer) # type: ignore[arg-type] - # io buffers are not guaranteed to have a name attribute - if hasattr(database, "name"): - filename = database.name # type: ignore[union-attr] - else: - filename = f"<{type(database)}>" - else: + return database + + def _load_file_object(self, database: DatabaseSource) -> object: + """Read a binary file object into memory.""" + read = getattr(database, "read", None) + if not callable(read): msg = ( - f"Unsupported open mode ({mode}). Only MODE_AUTO, MODE_FILE, " - "MODE_MEMORY and MODE_FD are supported by the pure Python " - "Reader" - ) - raise ValueError( - msg, + f"Unsupported database type for MODE_FD " + f"({type(database).__name__}). Pass a binary file object." ) - - return filename + raise TypeError(msg) + buf = read() + if isinstance(buf, bytearray): + # The decoder takes bytes. This copies the database once. + buf = bytes(buf) + if not isinstance(buf, bytes): + msg = f"The database file object returned {type(buf).__name__}, not bytes." + if isinstance(buf, str): + 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 + return f"<{type(database).__name__}>" def close(self) -> None: """Close the MaxMind DB file and returns the resources to the system. @@ -325,7 +414,7 @@ def close(self) -> None: self.closed = True - def __exit__(self, *_) -> None: # noqa: ANN002 + def __exit__(self, *_: object) -> None: self.close() def __enter__(self) -> Self: @@ -335,6 +424,58 @@ def __enter__(self) -> Self: return self +# The type of each metadata value. libmaxminddb also rejects a database with a +# missing key or a value of another type. It also checks the width and sign of +# each integer, which the decoder does not report. +_METADATA_TYPES: dict[str, type] = { + "binary_format_major_version": int, + "binary_format_minor_version": int, + "build_epoch": int, + "database_type": str, + "description": dict, + "ip_version": int, + "languages": list, + "node_count": int, + "record_size": int, +} + + +def _metadata_fields(metadata: RecordDict, filename: object) -> dict[str, Any]: + """Return the known metadata fields after a check of their types. + + A new minor version of the format can add keys. This ignores them. + """ + prefix = f"Error reading metadata in database file ({filename})." + fields: dict[str, Any] = {} + for key, value_type in _METADATA_TYPES.items(): + value = metadata.get(key) + # The exact type check rejects bool, a subclass of int. + valid = type(value) is value_type + 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()) + if not valid: + msg = f"{prefix} The {key} value is missing or has the wrong type." + raise InvalidDatabaseError(msg) + fields[key] = value + + # Range checks that libmaxminddb also makes. The reader decodes only the + # version 2 format, and ip_version drives the tree walk. libmaxminddb also + # rejects node_count 0, but this reader accepts an empty search tree. + if fields["binary_format_major_version"] != 2: + version = fields["binary_format_major_version"] + msg = f"{prefix} Unsupported binary format version {version}." + raise InvalidDatabaseError(msg) + if fields["ip_version"] not in (4, 6): + msg = f"{prefix} The ip_version is {fields['ip_version']}, not 4 or 6." + raise InvalidDatabaseError(msg) + if fields["build_epoch"] == 0: + msg = f"{prefix} The build_epoch is 0." + raise InvalidDatabaseError(msg) + return fields + + @dataclass(kw_only=True, frozen=True) class Metadata: """Metadata for the MaxMind DB reader.""" diff --git a/maxminddb/types.py b/maxminddb/types.py index dcb7b61e..4f06a5b0 100644 --- a/maxminddb/types.py +++ b/maxminddb/types.py @@ -1,10 +1,11 @@ -"""Types representing database records.""" +"""Types for database records and database arguments.""" from __future__ import annotations -from typing import TypeAlias +import os +from typing import Protocol, TypeAlias -Primitive: TypeAlias = str | bytes | bool | float | int +Primitive: TypeAlias = str | bytes | bytearray | bool | float | int RecordList: TypeAlias = list["Record"] """RecordList is a type for lists in a database record.""" @@ -13,3 +14,17 @@ """RecordDict is a type for dicts in a database record.""" Record: TypeAlias = Primitive | RecordList | RecordDict + +StrOrBytesPath: TypeAlias = str | bytes | os.PathLike[str] | os.PathLike[bytes] +"""StrOrBytesPath is a type for a path to a database file.""" + + +class SupportsRead(Protocol): + """SupportsRead is a type for a binary file object for MODE_FD or MODE_AUTO.""" + + def read(self) -> bytes: + """Return the remaining bytes.""" + + +DatabaseSource: TypeAlias = StrOrBytesPath | int | SupportsRead +"""DatabaseSource is a type for the database argument of a reader.""" diff --git a/tests/reader_test.py b/tests/reader_test.py index ce95714e..7c1902aa 100644 --- a/tests/reader_test.py +++ b/tests/reader_test.py @@ -1,16 +1,20 @@ from __future__ import annotations import contextlib +import dataclasses +import gzip import io import ipaddress import multiprocessing import os import pathlib +import struct import sys +import sysconfig import tempfile import threading import unittest -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, cast from unittest import mock import maxminddb @@ -29,11 +33,13 @@ MODE_MMAP, MODE_MMAP_EXT, ) +from maxminddb.decoder import Decoder if TYPE_CHECKING: from collections.abc import Iterator from maxminddb.reader import Reader + from maxminddb.types import DatabaseSource # Directory holding the shared MaxMind DB test fixtures. @@ -94,6 +100,87 @@ def address_space_in_use() -> int: resource.setrlimit(resource.RLIMIT_AS, (soft, hard)) +_METADATA_START_MARKER = b"\xab\xcd\xefMaxMind.com" +# libmaxminddb requires these unsigned integer types for metadata values. +_METADATA_UINT_TYPES = {"build_epoch": 9, "node_count": 6} +_UINT16_TYPE = 5 + + +def _database(records: tuple[int, ...], *, ip_version: int) -> 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". + """ + metadata = { + "binary_format_major_version": 2, + "binary_format_minor_version": 0, + "build_epoch": 1, + "database_type": "Test", + "description": {"en": "Test"}, + "ip_version": ip_version, + "languages": ["en"], + "node_count": len(records) // 2, + "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) + ) + + +def _database_with_metadata(**changes: object) -> bytes: + """Return the decoder test database with changed metadata. + + A value of None removes the key. + """ + data = pathlib.Path(f"{_TEST_DATA_DIR}/MaxMind-DB-test-decoder.mmdb").read_bytes() + start = data.rfind(_METADATA_START_MARKER) + len(_METADATA_START_MARKER) + (metadata, _) = Decoder(data, start).decode(start) + merged = {**cast("dict[str, object]", metadata), **changes} + edited = {k: v for k, v in merged.items() if v is not None} + return data[:start] + _encode_value(edited) + + +def _encode_value(value: object, key: str = "") -> bytes: + if isinstance(value, str): + encoded = value.encode() + return _encode_control(2, len(encoded)) + encoded + if isinstance(value, bool): + return _encode_control(14, int(value)) + if isinstance(value, float): + return _encode_control(3, 8) + struct.pack(">d", value) + if isinstance(value, int): + encoded = value.to_bytes((value.bit_length() + 7) // 8, "big") + type_num = _METADATA_UINT_TYPES.get(key, _UINT16_TYPE) + return _encode_control(type_num, len(encoded)) + encoded + if isinstance(value, list): + items = b"".join(_encode_value(v) for v in value) + return _encode_control(11, len(value)) + items + if isinstance(value, dict): + items = b"".join( + _encode_value(k) + _encode_value(v, k) for k, v in value.items() + ) + return _encode_control(7, len(value)) + items + msg = f"cannot encode {value!r}" + raise TypeError(msg) + + +def _encode_control(type_num: int, size: int) -> bytes: + # Sizes from 29 to 284 use one extra size byte. These tests need no more. + extended = b"" + if type_num > 7: + extended = bytes([type_num - 7]) + type_num = 0 + if size < 29: + return bytes([type_num << 5 | size]) + extended + return bytes([type_num << 5 | 29]) + extended + bytes([size - 29]) + + def get_reader_from_file_descriptor(filepath: str, mode: int) -> Reader: """Patches open_database() for class TestFDReader().""" if mode == MODE_FD: @@ -599,6 +686,141 @@ def test_search_tree_past_end_of_file(self) -> None: ): reader.get(self.ipf("1.1.1.1")) + def test_exhausted_iterator_stops_after_close(self) -> None: + reader = open_database( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb", + self.mode, + ) + iterator = iter(reader) + list(iterator) + reader.close() + self.assertEqual(next(iterator, "done"), "done") + + def test_record_that_points_to_the_root_is_rejected(self) -> None: + # Node 1's right record points back to the root. The left records point + # at data, so a walk through the root again would yield networks that + # the tree does not have, such as 192.0.0.0/3. + records = (18, 1, 18, 0) + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "root-record.mmdb" + path.write_bytes(_database(records, ip_version=4)) + with open_database(str(path), self.mode) as reader: + seen: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = [] + with self.assertRaisesRegex( + InvalidDatabaseError, + "search tree is corrupt", + ): + for network, _ in reader: + seen.append(network) + valid = { + ipaddress.ip_network("0.0.0.0/1"), + ipaddress.ip_network("128.0.0.0/2"), + } + self.assertLessEqual(set(seen), valid) + + def test_iterate_ipv6_networks_shorter_than_96_bits(self) -> None: + # One node whose two records point at the same data record, so the + # tree holds ::/1 and 8000::/1. + records = (17, 17) + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "short-ipv6.mmdb" + path.write_bytes(_database(records, ip_version=6)) + with open_database(str(path), self.mode) as reader: + self.assertEqual( + list(reader), + [ + (ipaddress.ip_network("::/1"), "net"), + (ipaddress.ip_network("8000::/1"), "net"), + ], + ) + + def test_record_that_points_into_the_separator_is_rejected(self) -> None: + # The left record, node_count + 1, points into the 16-byte separator + # between the search tree and the data section. libmaxminddb before + # 1.14 reports it as bad data. + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "separator.mmdb" + path.write_bytes(_database((2, 17), ip_version=4)) + with ( + open_database(str(path), self.mode) as reader, + self.assertRaisesRegex( + InvalidDatabaseError, + "search tree is corrupt|contains bad data", + ), + ): + reader.get(self.ipf("1.1.1.1")) + + def test_cyclic_search_tree_is_rejected(self) -> None: + data = bytearray( + pathlib.Path(f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb").read_bytes(), + ) + # Point the left record of node 1 back at node 1. + data[6:9] = b"\x00\x00\x01" + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "cyclic.mmdb" + path.write_bytes(data) + with open_database(str(path), self.mode) as reader: + iterator = iter(reader) + with self.assertRaisesRegex( + InvalidDatabaseError, + "search tree is corrupt", + ): + list(iterator) + # The iterator stops after an error, as a generator does. + self.assertEqual(next(iterator, "done"), "done") + + # A record that points back to the root of an IPv4 tree. + broken = f"{_TEST_DATA_DIR}/MaxMind-DB-test-broken-search-tree-24.mmdb" + with ( + open_database(broken, self.mode) as reader, + self.assertRaisesRegex(InvalidDatabaseError, "search tree is corrupt"), + ): + list(reader) + + # A record in the IPv4 subtree of an IPv6 tree that points back to the + # IPv4 start node, 96. + mixed = bytearray( + pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-mixed-24.mmdb" + ).read_bytes(), + ) + mixed[240 * 6 : 240 * 6 + 3] = (96).to_bytes(3, "big") + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "ipv4-cycle.mmdb" + path.write_bytes(mixed) + with ( + open_database(str(path), self.mode) as reader, + self.assertRaisesRegex(InvalidDatabaseError, "search tree is corrupt"), + ): + list(reader) + + def test_invalid_metadata_is_rejected(self) -> None: + cases: dict[str, dict[str, object]] = { + "missing languages": {"languages": None}, + "missing description": {"description": None}, + "string node_count": {"node_count": "1"}, + "double node_count": {"node_count": 1.5}, + "boolean record_size": {"record_size": True}, + "string languages": {"languages": "en"}, + "integer in languages": {"languages": [1]}, + "integer in description": {"description": {"en": 1}}, + "integer key in description": {"description": {1: "en"}}, + "integer database_type": {"database_type": 5}, + "ip_version 5": {"ip_version": 5}, + "binary_format_major_version 3": {"binary_format_major_version": 3}, + "build_epoch 0": {"build_epoch": 0}, + } + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "invalid-metadata.mmdb" + for name, changes in cases.items(): + with self.subTest(name): + path.write_bytes(_database_with_metadata(**changes)) + with ( + self.assertRaises(InvalidDatabaseError), + 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", @@ -849,6 +1071,11 @@ def _check_metadata( self.assertGreater(metadata.node_count, 36) self.assertEqual(metadata.record_size, record_size) + self.assertEqual(metadata.node_byte_size, record_size // 4) + self.assertEqual( + metadata.search_tree_size, + metadata.node_count * record_size // 4, + ) def _check_ip_v4(self, reader: Reader, file_name: str) -> None: for i in range(6): @@ -929,6 +1156,34 @@ 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_unknown_metadata_key_is_ignored(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = pathlib.Path(directory) / "unknown-key.mmdb" + path.write_bytes(_database_with_metadata(unknown_key="value")) + with maxminddb.extension.Reader(path) as reader: + metadata = reader.metadata() + self.assertEqual(metadata.database_type, "MaxMind DB Decoder Test") + self.assertFalse(hasattr(metadata, "unknown_key")) + + def test_file_object_is_refused(self) -> None: + with self.assertRaisesRegex(TypeError, "requires a path"): + open_database(io.BytesIO(b""), MODE_MMAP_EXT) + @unittest.skipIf( not has_maxminddb_extension() and not os.environ.get("MM_FORCE_EXT_TESTS"), @@ -942,6 +1197,148 @@ class TestExtensionReaderWithIPObjects(BaseTestReader): reader_class = maxminddb.extension.Reader +class TestModule(unittest.TestCase): + def test_mode_is_exported(self) -> None: + self.assertIs(maxminddb.Mode, maxminddb.const.Mode) + + +@unittest.skipIf( + not has_maxminddb_extension() and not os.environ.get("MM_FORCE_EXT_TESTS"), + "No C extension module found. Skipping tests", +) +class TestExtensionObjects(unittest.TestCase): + """Objects in states that crashed the extension.""" + + def test_uninitialized_metadata(self) -> None: + metadata_class = maxminddb.extension.Metadata + metadata = metadata_class.__new__(metadata_class) + self.assertIsNone(metadata.languages) + with self.assertRaisesRegex(AttributeError, "record_size is not set"): + _ = metadata.node_byte_size + with self.assertRaisesRegex(AttributeError, "node_count is not set"): + _ = metadata.search_tree_size + del metadata + + def test_metadata_missing_argument(self) -> None: + with self.assertRaisesRegex(TypeError, "missing required argument"): + maxminddb.extension.Metadata(binary_format_major_version=2) # type: ignore[call-arg] + + def test_metadata_unknown_argument(self) -> None: + with self.assertRaisesRegex(TypeError, "keyword argument"): + maxminddb.extension.Metadata( # type: ignore[call-arg] + binary_format_major_version=2, + binary_format_minor_version=0, + build_epoch=1, + database_type="db", + description={}, + ip_version=4, + languages=[], + node_count=1, + record_size=24, + unknown=1, + ) + + def test_iterate_uninitialized_reader(self) -> None: + reader_class = maxminddb.extension.Reader + reader = reader_class.__new__(reader_class) + with self.assertRaisesRegex(ValueError, "closed MaxMind DB"): + iter(reader) + + def test_enter_uninitialized_reader(self) -> None: + reader_class = maxminddb.extension.Reader + reader = reader_class.__new__(reader_class) + with self.assertRaisesRegex(ValueError, "closed MaxMind DB"): + reader.__enter__() + + def test_reinitialize_reader_is_refused(self) -> None: + path = f"{_TEST_DATA_DIR}/MaxMind-DB-test-decoder.mmdb" + with maxminddb.extension.Reader(path) as reader: + with self.assertRaisesRegex(ValueError, "reinitialize"): + reader.__init__(path) # type: ignore[misc] + self.assertIsNotNone(reader.get("::1.1.1.0")) + + # Re-init on a closed reader would leave this iterator pointing at a + # freed database. + closed = maxminddb.extension.Reader(path) + iterator = iter(closed) + next(iterator) + closed.close() + with self.assertRaisesRegex(ValueError, "reinitialize"): + closed.__init__(path) # type: ignore[misc] + + def test_reinitialize_metadata_is_refused(self) -> None: + fields: dict[str, Any] = { + "binary_format_major_version": 2, + "binary_format_minor_version": 0, + "build_epoch": 1, + "database_type": "db", + "description": {}, + "ip_version": 4, + "languages": [], + "node_count": 1, + "record_size": 24, + } + metadata = maxminddb.extension.Metadata(**fields) + with self.assertRaisesRegex(ValueError, "reinitialize"): + metadata.__init__(**{**fields, "record_size": 28}) # type: ignore[misc] + self.assertEqual(metadata.record_size, 24) + + @unittest.skipUnless( + hasattr(sys, "getrefcount") and not sysconfig.get_config_var("Py_GIL_DISABLED"), + "needs CPython reference counts on a build with the GIL", + ) + def test_freed_objects_release_their_type(self) -> None: + path = f"{_TEST_DATA_DIR}/MaxMind-DB-test-decoder.mmdb" + with maxminddb.extension.Reader(path) as reader: + classes = [type(reader), type(reader.metadata()), type(iter(reader))] + before = [sys.getrefcount(c) for c in classes] + for _ in range(10): + with maxminddb.extension.Reader(path) as reader: + reader.metadata() + iter(reader) + self.assertEqual([sys.getrefcount(c) for c in classes], before) + + @unittest.skipIf( + getattr(sys, "_is_gil_enabled", lambda: True)(), + "needs free-threaded Python", + ) + def test_threads_can_share_an_iterator(self) -> None: + path = f"{_TEST_DATA_DIR}/GeoIP2-City-Test.mmdb" + with maxminddb.extension.Reader(path) as reader: + expected = sum(1 for _ in reader) + + def count( + iterator: Iterator[object], + barrier: threading.Barrier, + counts: list[int], + ) -> None: + barrier.wait() + counts.append(sum(1 for _ in iterator)) + + # A race corrupts the heap only some of the time, so repeat. + for _ in range(5): + counts: list[int] = [] + args = (iter(reader), threading.Barrier(8), counts) + threads = [threading.Thread(target=count, args=args) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + self.assertEqual(sum(counts), expected) + + def test_iterator_type_is_not_instantiable(self) -> None: + reader = maxminddb.extension.Reader( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-decoder.mmdb", + ) + iterator_class = type(iter(reader)) + # The message differs across Python versions, so check only the type. + with self.assertRaises(TypeError): + iterator_class() + with self.assertRaises(TypeError): + iterator_class.__new__(iterator_class) + reader.close() + + class TestAutoReader(BaseTestReader): mode = MODE_AUTO @@ -987,6 +1384,180 @@ def setUp(self) -> None: class TestReaderInitialization(unittest.TestCase): + def test_metadata_types_match_metadata_fields(self) -> None: + self.assertEqual( + list(maxminddb.reader._METADATA_TYPES), # noqa: SLF001 + [field.name for field in dataclasses.fields(maxminddb.reader.Metadata)], + ) + + def test_auto_mode_reads_file_object_from_its_position(self) -> None: + data = pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb", + ).read_bytes() + with tempfile.TemporaryFile() as file_object: + file_object.write(b"header" + data) + file_object.seek(len(b"header")) + with maxminddb.open_database(file_object, MODE_AUTO) as reader: + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + + def test_auto_mode_reads_wrapped_and_piped_streams(self) -> None: + path = f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" + data = pathlib.Path(path).read_bytes() + + with tempfile.TemporaryDirectory() as directory: + gz_path = pathlib.Path(directory) / "db.mmdb.gz" + with gzip.open(gz_path, "wb") as compressed: + compressed.write(data) + with open(gz_path, "rb") as backing: + stream = io.BufferedReader(gzip.GzipFile(fileobj=backing)) + with maxminddb.open_database(stream, MODE_AUTO) as reader: + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + # The reader does not close the caller's stream. + self.assertFalse(stream.closed) + stream.close() + + # The fixture is far smaller than the pipe buffer, so the write to the + # pipe does not block. + read_fd, write_fd = os.pipe() + os.write(write_fd, data) + os.close(write_fd) + pipe = os.fdopen(read_fd, "rb") + try: + with maxminddb.open_database(pipe, MODE_AUTO) as reader: + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + self.assertFalse(pipe.closed) + finally: + pipe.close() + + def test_auto_mode_accepts_any_database_type(self) -> None: + path = f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" + data = pathlib.Path(path).read_bytes() + + class PathWithTextRead: + """A path object with a text read(), as py.path.local has.""" + + def __fspath__(self) -> str: + return path + + def read(self) -> str: + return "not the database" + + with open(path, "rb") as file_object: + sources: list[tuple[str, DatabaseSource]] = [ + ("path", path), + ("file object", file_object), + ("BytesIO", io.BytesIO(data)), + # A path wins over read(). + ("path with a text read()", PathWithTextRead()), + ] + for name, database in sources: + with ( + self.subTest(name), + maxminddb.open_database(database, MODE_AUTO) as reader, + ): + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + + # The pure Python reader takes ownership of a descriptor and closes it. + with maxminddb.reader.Reader(os.open(path, os.O_RDONLY), MODE_AUTO) as reader: + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + if has_maxminddb_extension(): + # open_database refuses a descriptor with the extension, as before. + descriptor = os.open(path, os.O_RDONLY) + self.addCleanup(os.close, descriptor) + with self.assertRaisesRegex(TypeError, r"\(int given\)"): + maxminddb.open_database(descriptor, MODE_AUTO) + + def test_database_type_must_match_mode(self) -> None: + path = f"{_TEST_DATA_DIR}/MaxMind-DB-test-decoder.mmdb" + reader_class = maxminddb.reader.Reader + with self.assertRaisesRegex(TypeError, r"MODE_FD \(str\)"): + reader_class(path, MODE_FD) + with open(path, "rb") as database: + for mode in (MODE_FILE, MODE_MEMORY, MODE_MMAP): + with ( + self.subTest(mode=mode), + self.assertRaisesRegex(TypeError, "Use MODE_FD"), + ): + reader_class(database, mode) + with self.assertRaisesRegex(ValueError, "Unsupported open mode"): + reader_class(database, MODE_MMAP_EXT) + + for bad in (None, object()): + with self.assertRaisesRegex(TypeError, "Unsupported database type"): + reader_class(bad, MODE_AUTO) # type: ignore[arg-type] + + for mode in (MODE_AUTO, MODE_FD): + with ( + self.subTest(mode=mode), + open(path, encoding="latin-1") as text_file, + self.assertRaisesRegex(TypeError, "binary mode"), + ): + reader_class(text_file, mode) # type: ignore[arg-type] + + # A bool is an int, but not a file descriptor. + with self.assertRaisesRegex(TypeError, r"\(bool\)"): + reader_class(False, MODE_AUTO) # noqa: FBT003 + + def test_path_modes_accept_a_descriptor_with_index(self) -> None: + class Descriptor: + """A file descriptor object, as numpy.int64 is.""" + + def __init__(self, fd: int) -> None: + self.fd = fd + + def __index__(self) -> int: + return self.fd + + path = f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" + # The reader takes ownership of the descriptor and closes it. + descriptor = Descriptor(os.open(path, os.O_RDONLY)) + with maxminddb.reader.Reader(descriptor, MODE_MMAP) as reader: # type: ignore[arg-type] + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + + if has_maxminddb_extension(): + # With the extension, MODE_AUTO refuses it, as it refuses an int. + fd = os.open(path, os.O_RDONLY) + self.addCleanup(os.close, fd) + with self.assertRaisesRegex(TypeError, r"\(Descriptor given\)"): + maxminddb.open_database(Descriptor(fd)) # type: ignore[arg-type] + # A bool gets the bool error, not advice to use MODE_MMAP. + with self.assertRaisesRegex(TypeError, r"\(bool\)\. Pass a path"): + maxminddb.open_database(False) # noqa: FBT003 + + def test_fd_mode_reads_any_binary_reader(self) -> None: + data = pathlib.Path( + f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb", + ).read_bytes() + + class FileWithPath(io.BytesIO): + def __fspath__(self) -> str: + return "does-not-exist.mmdb" + + class BytearrayReader: + def read(self) -> bytes: + return bytearray(data) # type: ignore[return-value] + + for database in (FileWithPath(data), BytearrayReader()): + with ( + self.subTest(type(database).__name__), + maxminddb.reader.Reader(database, MODE_FD) as reader, + ): + self.assertEqual(reader.get("1.1.1.1"), {"ip": "1.1.1.1"}) + + class NoneReader: + def read(self) -> None: + return None + + with self.assertRaisesRegex(TypeError, r"returned NoneType, not bytes\.$"): + maxminddb.reader.Reader(NoneReader(), MODE_FD) # type: ignore[arg-type] + + def test_unknown_metadata_key_is_ignored(self) -> None: + data = _database_with_metadata(unknown_key="value") + with maxminddb.reader.Reader(io.BytesIO(data), MODE_FD) as reader: + metadata = reader.metadata() + self.assertEqual(metadata.database_type, "MaxMind DB Decoder Test") + self.assertFalse(hasattr(metadata, "unknown_key")) + def test_empty_search_tree_is_accepted(self) -> None: data = pathlib.Path( f"{_TEST_DATA_DIR}/MaxMind-DB-test-ipv4-24.mmdb" @@ -1024,11 +1595,18 @@ def test_invalid_tree_metadata_is_rejected_on_open(self) -> None: 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"), + ( + _METADATA_START_MARKER + b"\x40", + InvalidDatabaseError, + "Error reading metadata", + ), + ( + _METADATA_START_MARKER + b"\xe0", + InvalidDatabaseError, + "missing or has the wrong type", + ), ( pathlib.Path( f"{_TEST_DATA_DIR}/MaxMind-DB-test-metadata-payload-limit.mmdb" diff --git a/tests/typing_test.py b/tests/typing_test.py new file mode 100644 index 00000000..353d52c0 --- /dev/null +++ b/tests/typing_test.py @@ -0,0 +1,51 @@ +# mypy: warn-unused-ignores +"""Static type checks for the public API. + +mypy checks this file in the lint environment. The code does not run. Each +type: ignore marks a call that must fail the type check. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import gzip + import io + from ipaddress import IPv4Network, IPv6Network + + from typing_extensions import assert_type + + import maxminddb + import maxminddb.extension + from maxminddb.types import Primitive, Record + + reader = maxminddb.open_database("GeoIP2-City.mmdb") + assert_type(reader.get("1.1.1.1"), Record | None) + assert_type(reader.get_with_prefix_len("1.1.1.1"), tuple[Record | None, int]) + for network, record in reader: + assert_type(network, IPv4Network | IPv6Network) + assert_type(record, Record) + assert_type(reader.metadata().search_tree_size, int) + + # Record includes the bytearray that the C extension returns for the + # bytes type. + value: Record = bytearray(b"\x00") + + # A TypeVar in either alias would give Record members of type Any. + def check_not_generic( + primitive: Primitive[str], # type: ignore[type-arg] + record: Record[str], # type: ignore[type-arg] + ) -> None: + pass + + def check_mode_fd(gzip_file: gzip.GzipFile, text_file: io.TextIOWrapper) -> None: + maxminddb.open_database(gzip_file, maxminddb.Mode.FD) + maxminddb.open_database(text_file, maxminddb.Mode.FD) # type: ignore[arg-type] + + extension_reader = maxminddb.extension.Reader("GeoIP2-City.mmdb") + for network, record in extension_reader: + assert_type(network, IPv4Network | IPv6Network) + assert_type(record, Record) + assert_type(extension_reader.metadata().node_byte_size, int) + maxminddb.extension.Reader(3) # type: ignore[arg-type]