Use buffer protocol to retrieve result from python post_process (#605)

Saves an extra copy, gives post processing functions more flexibility
This commit is contained in:
Dustin Spicuzza 2020-11-18 14:29:17 -05:00 committed by GitHub
parent b260204b72
commit cd0a25be5e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 28 additions and 12 deletions

View File

@ -130,6 +130,9 @@ def introspection():
`post_process` function. This function is then transforming the data into the
format expected by the API before executing the target.
This can return any python object that implements the buffer protocol and
supports PyBUF_SIMPLE. These include bytes, bytearray, etc.
- `queue_new_entry` (optional):
This methods is called after adding a new test case to the queue.

View File

@ -326,8 +326,7 @@ typedef struct py_mutator {
u8 * fuzz_buf;
size_t fuzz_size;
u8 * post_process_buf;
size_t post_process_size;
Py_buffer post_process_buf;
u8 * trim_buf;
size_t trim_size;

View File

@ -134,6 +134,18 @@ static py_mutator_t *init_py_module(afl_state_t *afl, u8 *module_name) {
PyObject * py_module = py->py_module;
PyObject **py_functions = py->py_functions;
// initialize the post process buffer; ensures it's always valid
PyObject *unused_bytes = PyByteArray_FromStringAndSize("OHAI", 4);
if (!unused_bytes) { FATAL("allocation failed!"); }
if (PyObject_GetBuffer(unused_bytes, &py->post_process_buf, PyBUF_SIMPLE) ==
-1) {
FATAL("buffer initialization failed");
}
Py_DECREF(unused_bytes);
if (py_module != NULL) {
u8 py_notrim = 0, py_idx;
@ -313,7 +325,6 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
struct custom_mutator *mutator;
mutator = ck_alloc(sizeof(struct custom_mutator));
mutator->post_process_buf = NULL;
mutator->name = module_name;
ACTF("Loading Python mutator library from '%s'...", module_name);
@ -403,10 +414,13 @@ struct custom_mutator *load_custom_mutator_py(afl_state_t *afl,
size_t post_process_py(void *py_mutator, u8 *buf, size_t buf_size,
u8 **out_buf) {
size_t py_out_buf_size;
PyObject * py_args, *py_value;
py_mutator_t *py = (py_mutator_t *)py_mutator;
// buffer returned previously must be released; initialized during init
// so we don't need to do comparisons
PyBuffer_Release(&py->post_process_buf);
py_args = PyTuple_New(1);
py_value = PyByteArray_FromStringAndSize(buf, buf_size);
if (!py_value) {
@ -426,20 +440,20 @@ size_t post_process_py(void *py_mutator, u8 *buf, size_t buf_size,
if (py_value != NULL) {
py_out_buf_size = PyByteArray_Size(py_value);
if (PyObject_GetBuffer(py_value, &py->post_process_buf, PyBUF_SIMPLE) ==
-1) {
if (unlikely(!afl_realloc(BUF_PARAMS(post_process), py_out_buf_size))) {
PFATAL("alloc");
PyErr_Print();
FATAL(
"Python custom mutator: post_process call return value not a "
"bytes-like object");
}
memcpy(py->post_process_buf, PyByteArray_AsString(py_value),
py_out_buf_size);
Py_DECREF(py_value);
*out_buf = py->post_process_buf;
return py_out_buf_size;
*out_buf = (u8 *)py->post_process_buf.buf;
return py->post_process_buf.len;
} else {