/**

* Copyright 2018 Jussi Judin

*

* Licensed under the Apache License, Version 2.0 (the "License");

* you may not use this file except in compliance with the License.

* You may obtain a copy of the License at

*

* http://www.apache.org/licenses/LICENSE-2.0

*

* Unless required by applicable law or agreed to in writing, software

* distributed under the License is distributed on an "AS IS" BASIS,

* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

* See the License for the specific language governing permissions and

* limitations under the License.

*/

#include <errno.h>

#include <stdio.h>

#include <signal.h>

#include <unistd.h>

#include <sys/shm.h>

#include <stdlib.h>

#include <stdbool.h>

#include <string.h>

#include <sys/wait.h>

#include <javafl_JavaAfl.h>

// Use afl's config.h for constants.

#ifdef HAVE_AFL_CONFIG_H

#include <config.h>

#else // #ifndef HAVE_AFL_CONFIG_H

#ifdef HAVE_AFL_MAP_SIZE_H

#include <afl-map-size.h>

#endif // #ifdef HAVE_AFL_MAP_SIZE_H

// These constants must be kept in sync with afl-fuzz:

#ifndef MAP_SIZE

#ifndef MAP_SIZE_POW2

#define MAP_SIZE_POW2 16;

#endif // #ifndef MAP_SIZE_POW2

static const size_t MAP_SIZE = 1 << MAP_SIZE_POW2;

#endif // #ifndef MAP_SIZE

static const char SHM_ENV_VAR[] = "__AFL_SHM_ID";

static const int FORKSRV_FD = 198;

#endif // #ifndef HAVE_AFL_CONFIG_H

// These are global helper variables to avoid discovering the same

// information again and again.

static void* g_afl_area = (void*)-1;

static void* g_zero_area = NULL;

static jfieldID g_map_field_id = NULL;

static bool g_is_persistent = false;

static bool g_initialized = false;

static void init_map_field(JNIEnv *env, jclass cls)

{

jfieldID map_field_id = (*env)->GetStaticFieldID(env, cls, "map", "[B");

if (map_field_id == NULL) {

fprintf(stderr, "No map field found from JavaAfl class!\n");

abort();

}

g_map_field_id = map_field_id;

}

static jint get_prev_location(JNIEnv *env, jclass cls)

{

jfieldID prev_location_field_id = (*env)->GetStaticFieldID(env, cls, "prev_location", "I");

return (*env)->GetStaticIntField(env, cls, prev_location_field_id);

}

static jobject get_map_field(JNIEnv *env, jclass cls)

{

return (*env)->GetStaticObjectField(env, cls, g_map_field_id);

}

JNIEXPORT jint JNICALL Java_javafl_JavaAfl__1get_1map_1size

(JNIEnv * env, jclass cls)

{

return MAP_SIZE;

}

JNIEXPORT void JNICALL Java_javafl_JavaAfl__1init_1impl

(JNIEnv * env, jclass cls, jboolean is_persistent)

{

if (g_initialized) {

fprintf(

stderr,

"Tried to initialize java-afl twice! "

"If you are using deferred or persistent mode, remember to "

"annotate your main() function with @javafl.CustomInit!\n"

);

abort();

}

bool use_forkserver = true;

{

int result = write(FORKSRV_FD + 1, "\x00\x00\x00\x00", 4);

if (result == -1) {

if (errno == EBADF) {

use_forkserver = false;

} else {

perror("Failed to send data to fork server");

abort();

}

}

}

bool child_stopped = false;

union

{

char child_pid_buf[4];

pid_t child_pid;

} child_pid_data = { .child_pid = 0 };

while (use_forkserver) {

union

{

char child_killed_buf[4];

unsigned child_killed;

} child_killed_data;

// Wait for parent. It can also tell us that it has killed the

// child process::

if (read(FORKSRV_FD, child_killed_data.child_killed_buf, 4) != 4) {

perror("Failed to read child killed data");

abort();

}

// This handles the race condition where the child receives

// SIGSTOP first in the persistent mode and then is killed by

// the parent for timing out.

if (child_stopped && child_killed_data.child_killed) {

child_stopped = false;

if (waitpid(child_pid_data.child_pid, NULL, 0) == -1) {

perror("Waiting for the child process failed!");

abort();

}

}

if (child_stopped) {

// In persistent mode the child will send SIGSTOP to

// itself after it has written map data to the shared

// memory. This makes it run for another round in

// persistent mode.

kill(child_pid_data.child_pid, SIGCONT);

child_stopped = false;

} else {

child_pid_data.child_pid = fork();

if (!child_pid_data.child_pid) {

// Child will directly jump to shared memory handling

// related code.

break;

}

}

// Parent will repeatedly write status information back to

// afl-fuzz process.

write(FORKSRV_FD + 1, child_pid_data.child_pid_buf, 4);

int wstatus;

int options = 0;

if (is_persistent) {

options = WUNTRACED;

}

waitpid(child_pid_data.child_pid, &wstatus, options);

child_stopped = WIFSTOPPED(wstatus);

write(FORKSRV_FD + 1, &wstatus, sizeof(wstatus));

}

if (use_forkserver) {

close(FORKSRV_FD);

close(FORKSRV_FD + 1);

}

g_is_persistent = is_persistent;

g_initialized = true;

const char* afl_shm_id = getenv(SHM_ENV_VAR);

if (afl_shm_id == NULL) {

return;

}

// This area of zeros is here only to be able to zero the map

// memory on Java side fast when in persistent fuzzing mode.

g_zero_area = calloc(1, MAP_SIZE);

jint start_location = get_prev_location(env, cls);

// Have at least something in the map so that afl-fuzz or

// afl-showmap don't give confusing hard to debug messages.

((char*)g_zero_area)[start_location] = 1;

g_afl_area = shmat(atoi(afl_shm_id), NULL, 0);

if (g_afl_area == (void*)-1) {

perror("No shared memory area!");

abort();

}

init_map_field(env, cls);

// It's possible that Java side instrumentation has already

// written something to the map. Reset it so that especially

// persistent mode gets a clean slate to start. Deferred mode

// also should benefit from the map being less full in the

// beginning.

(*env)->SetByteArrayRegion(

env, get_map_field(env, cls), 0, MAP_SIZE, g_zero_area);

}

/**

* Copies map data generated in Java side to the shared memory and at

* the same time zeroes it.

*/

static void send_map(JNIEnv * env, jclass cls)

{

if (g_afl_area != (void*)-1) {

jobject map_field = get_map_field(env, cls);

(*env)->GetByteArrayRegion(env, map_field, 0, MAP_SIZE, g_afl_area);

(*env)->SetByteArrayRegion(env, map_field, 0, MAP_SIZE, g_zero_area);

}

}

JNIEXPORT void JNICALL Java_javafl_JavaAfl__1after_1main

(JNIEnv * env, jclass cls)

{

// Do nothing if we're not running inside something that can read

// maps. This makes Java exit handlers to run properly.

if (g_afl_area == (void*)-1) {

return;

}

// In persistent mode javafl.JavaAfl.loop() does the final map update for

// us. Doing map updates after the main loop leads into

// instability, as there are map updates in the code after that.

// TODO rethink this approach, as there is quite a lot of ifs

// taking care of the persistent mode.

if (!g_is_persistent) {

send_map(env, cls);

}

// JVM likely waits for specific threads exit that do not exist

// anymore in the forked child processes. This exits from the

// program before any exit handlers get to run.

_Exit(0);

}

JNIEXPORT void JNICALL Java_javafl_JavaAfl__1handle_1uncaught_1exception

(JNIEnv * env, jclass cls)

{

if (g_afl_area == (void*)-1) {

return;

}

send_map(env, cls);

abort();

}

JNIEXPORT void JNICALL Java_javafl_JavaAfl__1send_1map

(JNIEnv * env, jclass cls)

{

send_map(env, cls);

kill(getpid(), SIGSTOP);

}