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
3 changed files with 28 additions and 12 deletions

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 {