diff --git a/docs/usage.md b/docs/usage.md index a4ec38d..261ced4 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -176,9 +176,16 @@ with open('hs.db', 'wb') as f: f.write(serialized) # Deserializing (loading from bytes): -db = hyperscan.loadb(serialized) +db = hyperscan.loadb(serialized, mode=hyperscan.HS_MODE_BLOCK) +db.scan(b'foobar', match_event_handler=callback) ``` +Use the same scanning mode as the serialized database. By default, `loadb` +allocates scratch space, so the restored database can scan immediately. +Pass `alloc_scratch=False` to manage scratch allocation yourself; assign +`db.scratch = hyperscan.Scratch(db)` or pass an appropriate scratch object +when scanning. + ## Chimera Mode ```python diff --git a/src/hyperscan/__init__.pyi b/src/hyperscan/__init__.pyi index 177f26e..1ff4e18 100644 --- a/src/hyperscan/__init__.pyi +++ b/src/hyperscan/__init__.pyi @@ -82,12 +82,14 @@ def dumpb(database: "Database") -> bytes: """ -def loadb(buf: ByteString, mode: int) -> "Database": +def loadb(buf: ByteString, mode: int, alloc_scratch: bool = True) -> "Database": """Deserializes a Hyperscan database. Args: buf (bytes): A serialized Hyperscan database. mode (int): The expected mode of the database. + alloc_scratch (bool, optional): Allocate scratch space for scanning. + Defaults to True. If False, supply scratch space before scanning. Returns: :class:`Database`: The deserialized database instance. diff --git a/src/hyperscan/extension.c b/src/hyperscan/extension.c index e072d10..455916b 100644 --- a/src/hyperscan/extension.c +++ b/src/hyperscan/extension.c @@ -1465,24 +1465,40 @@ static PyObject *loadb(PyObject *self, PyObject *args, PyObject *kwds) char *buf; PyObject *obuf = Py_None; - PyObject *odb; - odb = PyObject_CallFunctionObjArgs((PyObject *)&DatabaseType, NULL); - Database *db = (Database *)odb; + unsigned int mode; + int alloc_scratch = 1; - static char *kwlist[] = {"buf", "mode", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "OI", kwlist, &obuf, &db->mode)) + static char *kwlist[] = {"buf", "mode", "alloc_scratch", NULL}; + if (!PyArg_ParseTupleAndKeywords( + args, kwds, "OI|p", kwlist, &obuf, &mode, &alloc_scratch)) HS_LOCK_RETURN_NULL(); if (!PyBytes_Check(obuf)) { PyErr_SetString(PyExc_TypeError, "buf must be a bytestring"); HS_LOCK_RETURN_NULL(); } + PyObject *odb = PyObject_CallFunctionObjArgs((PyObject *)&DatabaseType, NULL); + if (odb == NULL) + HS_LOCK_RETURN_NULL(); + Database *db = (Database *)odb; + db->mode = mode; + Py_ssize_t length = PyBytes_Size(obuf); buf = PyBytes_AsString(obuf); hs_error_t err = hs_deserialize_database(buf, length, &db->hs_db); - HANDLE_HYPERSCAN_ERR(err, NULL); - if (PyErr_Occurred()) - HS_LOCK_RETURN_NULL(); + if (err != HS_SUCCESS) { + Py_DECREF(odb); + HANDLE_HYPERSCAN_ERR(err, NULL); + } + if (alloc_scratch) { + PyObject *scratch = + PyObject_CallFunction((PyObject *)&ScratchType, "O", odb); + if (scratch == NULL) { + Py_DECREF(odb); + HS_LOCK_RETURN_NULL(); + } + db->scratch = scratch; + } HS_LOCK_RETURN(odb); } @@ -1499,11 +1515,14 @@ static PyMethodDef HyperscanMethods[] = { {"loadb", (PyCFunction)loadb, METH_VARARGS | METH_KEYWORDS, - "loadb(buf, mode)\n" + "loadb(buf, mode, alloc_scratch=True)\n" " Deserializes a Hyperscan database.\n\n" " Args:\n" - " buf (:obj:`bytearray`): A serialized Hyperscan database.\n" - " mode (int): The expected mode of the database.\n\n" + " buf (bytes): A serialized Hyperscan database.\n" + " mode (int): The expected mode of the database.\n" + " alloc_scratch (bool, optional): Allocate scratch space for\n" + " scanning. Defaults to True. If False, supply scratch\n" + " space before scanning.\n\n" " Returns:\n" " :class:`Database`: The deserialized database instance.\n\n"}, {NULL}}; diff --git a/tests/test_hyperscan.py b/tests/test_hyperscan.py index fcfac2d..37da00a 100644 --- a/tests/test_hyperscan.py +++ b/tests/test_hyperscan.py @@ -386,7 +386,8 @@ def test_database_deserialize_scan(db_fixture_name, request, mocker): original_db: hyperscan.Database = request.getfixturevalue(db_fixture_name) serialized_db = hyperscan.dumpb(original_db) db = hyperscan.loadb(serialized_db, original_db.mode) - db.scratch = hyperscan.Scratch(db) + assert isinstance(db.scratch, hyperscan.Scratch) + assert db.scratch.database is db if db_fixture_name == "database_stream": with db.stream(match_event_handler=callback) as stream: stream.scan(b"foobar") @@ -395,6 +396,46 @@ def test_database_deserialize_scan(db_fixture_name, request, mocker): if db_fixture_name == "database_vector": buf = [buf] db.scan(buf, match_event_handler=callback) + callback.assert_called() + + +@pytest.mark.parametrize( + "db_fixture_name", + ["database_stream", "database_block", "database_vector"], +) +@pytest.mark.parametrize("alloc_scratch", [False, True]) +def test_database_deserialize_scratch(db_fixture_name, alloc_scratch, request, mocker): + original_db = request.getfixturevalue(db_fixture_name) + db = hyperscan.loadb( + hyperscan.dumpb(original_db), original_db.mode, alloc_scratch=alloc_scratch + ) + if alloc_scratch: + assert isinstance(db.scratch, hyperscan.Scratch) + else: + assert db.scratch is None + db.scratch = hyperscan.Scratch(db) + + callback = mocker.Mock(return_value=None) + if db_fixture_name == "database_stream": + with db.stream(match_event_handler=callback) as stream: + stream.scan(b"foobar") + else: + data = [b"foo", b"bar"] if db_fixture_name == "database_vector" else b"foobar" + db.scan(data, match_event_handler=callback) + callback.assert_called() + + +@pytest.mark.parametrize( + "data, error", + [ + (None, TypeError), + (b"", hyperscan.InvalidError), + (b"not a database", hyperscan.InvalidError), + ], +) +def test_database_deserialize_invalid(data, error): + with pytest.raises(error): + hyperscan.loadb(data, hyperscan.HS_MODE_BLOCK) def test_database_exception_in_callback(database_block, mocker):