Refactoring `fuzz_py` API
MODIFY
include/afl-fuzz.h
include/afl-fuzz.h
@@ -627,7 +627,9 @@ int init_py_module(u8*);
void finalize_py_module();
void init_py(unsigned int seed);
-void fuzz_py(char*, size_t, char*, size_t, char**, size_t*);
+/* TODO: unify fuzz interface for custom mutator and Python mutator */
+size_t fuzz_py(u8*, size_t, u8*, size_t, unsigned int);
+void fuzz_py_original(char*, size_t, char*, size_t, char**, size_t*);
size_t pre_save_py(u8* data, size_t size, u8** new_data);
u32 init_trim_py(u8*, size_t);
u32 post_trim_py(u8);

MODIFY
src/afl-fuzz-mutators.c
src/afl-fuzz-mutators.c
@@ -281,7 +281,6 @@ void load_custom_mutator_py(const char* module_name) {
mutator->name = module_name;
ACTF("Loading Python mutator library from '%s'...", module_name);
- /* TODO: unify "init" and "fuzz" */
if (py_functions[PY_FUNC_INIT])
mutator->afl_custom_init = init_py;

MODIFY
src/afl-fuzz-one.c
src/afl-fuzz-one.c
@@ -1647,7 +1647,7 @@ python_stage:
ck_read(fd, new_buf, target->len, target->fname);
close(fd);
- fuzz_py(out_buf, len, new_buf, target->len, &retbuf, &retlen);
+ fuzz_py_original(out_buf, len, new_buf, target->len, &retbuf, &retlen);
ck_free(new_buf);

MODIFY
src/afl-fuzz-python.c
src/afl-fuzz-python.c
@@ -32,6 +32,8 @@ int init_py_module(u8* module_name) {
if (!module_name) return 1;
+ Py_Initialize();
+
#if PY_MAJOR_VERSION >= 3
PyObject* py_name = PyUnicode_FromString(module_name);
#else
@@ -58,7 +60,12 @@ int init_py_module(u8* module_name) {
if (!py_functions[py_idx] || !PyCallable_Check(py_functions[py_idx])) {
- if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
+ if (py_idx == PY_FUNC_PRE_SAVE) {
+
+ // Implenting the pre_save API is optional for now
+ if (PyErr_Occurred()) PyErr_Print();
+
+ } else if (py_idx >= PY_FUNC_INIT_TRIM && py_idx <= PY_FUNC_TRIM) {
// Implementing the trim API is optional for now
if (PyErr_Occurred()) PyErr_Print();
@@ -152,7 +159,8 @@ void init_py(unsigned int seed) {
}
}
-void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
+void fuzz_py_original(char* buf, size_t buflen,
+ char* add_buf, size_t add_buflen,
char** ret, size_t* retlen) {
if (py_module != NULL) {
@@ -204,6 +212,72 @@ void fuzz_py(char* buf, size_t buflen, char* add_buf, size_t add_buflen,
}
+size_t fuzz_py(u8* data, size_t size, u8* mutated_out, size_t max_size,
+ unsigned int seed) {
+
+ size_t out_size;
+ PyObject *py_args, *py_value;
+ py_args = PyTuple_New(3);
+
+ py_value = PyByteArray_FromStringAndSize(data, size);
+ if (!py_value) {
+
+ Py_DECREF(py_args);
+ FATAL("Failed to convert arguments");
+
+ }
+
+ PyTuple_SetItem(py_args, 0, py_value);
+
+#if PY_MAJOR_VERSION >= 3
+ py_value = PyLong_FromLong(max_size);
+#else
+ py_value = PyInt_FromLong(max_size);
+#endif
+ if (!py_value) {
+
+ Py_DECREF(py_args);
+ FATAL("Failed to convert arguments");
+
+ }
+
+ PyTuple_SetItem(py_args, 1, py_value);
+
+#if PY_MAJOR_VERSION >= 3
+ py_value = PyLong_FromLong(seed);
+#else
+ py_value = PyInt_FromLong(seed);
+#endif
+ if (!py_value) {
+
+ Py_DECREF(py_args);
+ FATAL("Failed to convert arguments");
+
+ }
+
+ PyTuple_SetItem(py_args, 2, py_value);
+
+ py_value = PyObject_CallObject(py_functions[PY_FUNC_FUZZ], py_args);
+
+ Py_DECREF(py_args);
+
+ if (py_value != NULL) {
+
+ out_size = PyByteArray_Size(py_value);
+ memcpy(mutated_out, PyByteArray_AsString(py_value), out_size);
+ Py_DECREF(py_value);
+
+ return out_size;
+
+ } else {
+
+ PyErr_Print();
+ FATAL("Call failed");
+
+ }
+
+}
+
size_t pre_save_py(u8* data, size_t size, u8** new_data) {
size_t new_size;

MODIFY
src/afl-fuzz.c
src/afl-fuzz.c
@@ -862,7 +862,6 @@ int main(int argc, char** argv, char** envp) {
check_cpu_governor();
setup_post();
- setup_custom_mutator();
setup_shm(dumb_mode);
if (!in_bitmap) memset(virgin_bits, 255, MAP_SIZE);
@@ -873,6 +872,8 @@ int main(int argc, char** argv, char** envp) {
setup_dirs_fds();
+ setup_custom_mutator();
+
setup_cmdline_file(argv + optind);
read_testcases();

Refactoring `fuzz_py` API