mirror of
https://github.com/AFLplusplus/AFLplusplus.git
synced 2025-06-14 19:08:08 +00:00
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:
@ -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 {
|
||||
|
||||
|
Reference in New Issue
Block a user