intercept_exceptions.c

#include <jvmti.h>
#include <jni.h>
#include <string.h>

static jvmtiEnv *jvmti = NULL;
static jclass outOfMemoryErrorClass = NULL;

void check_exact_exception_type(JNIEnv *jni_env, jobject exception) {
    jclass exceptionClass = (*jni_env)->GetObjectClass(jni_env, exception);
    if ((*jni_env)->IsSameObject(jni_env, exceptionClass, outOfMemoryErrorClass)) {
        printf("Exact type: OutOfMemoryError captured!\n");
    } else {
        printf("Other exception captured.\n");
    }
}

void JNICALL ExceptionCallback(jvmtiEnv *jvmti_env,
                               JNIEnv* jni_env,
                               jthread thread,
                               jmethodID method,
                               jlocation location,
                               jobject exception,
                               jmethodID catch_method,
                               jlocation catch_location) {
    char *methodName;
    (*jvmti_env)->GetMethodName(jvmti_env, method, &methodName, NULL, NULL);
    printf("Exception thrown in method: %s\n", methodName);
    check_exact_exception_type(jni_env, exception);
    (*jvmti_env)->Deallocate(jvmti_env, (unsigned char*)methodName);
}

void JNICALL ExceptionCatchCallback(jvmtiEnv *jvmti_env,
                                    JNIEnv* jni_env,
                                    jthread thread,
                                    jmethodID method,
                                    jlocation location,
                                    jobject exception) {
    char *methodName;
    (*jvmti_env)->GetMethodName(jvmti_env, method, &methodName, NULL, NULL);
    printf("Exception caught in method: %s\n", methodName);
    check_exact_exception_type(jni_env, exception);
    (*jvmti_env)->Deallocate(jvmti_env, (unsigned char*)methodName);
}

JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *vm, char *options, void *reserved) {
    jint res;
    jvmtiError err;

    // Get the JVMTI environment
    res = (*vm)->GetEnv(vm, (void **)&jvmti, JVMTI_VERSION_1_0);
    if (res != JNI_OK || jvmti == NULL) {
        return JNI_ERR;
    }

    // Enable the Exception event
    err = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, JVMTI_EVENT_EXCEPTION, NULL);
    if (err != JVMTI_ERROR_NONE) {
        return JNI_ERR;
    }

    // Enable the ExceptionCatch event
    err = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, JVMTI_EVENT_EXCEPTION_CATCH, NULL);
    if (err != JVMTI_ERROR_NONE) {
        return JNI_ERR;
    }

    // Register callbacks
    jvmtiEventCallbacks callbacks;
    memset(&callbacks, 0, sizeof(callbacks));
    callbacks.Exception = &ExceptionCallback;
    callbacks.ExceptionCatch = &ExceptionCatchCallback;
    err = (*jvmti)->SetEventCallbacks(jvmti, &callbacks, sizeof(callbacks));
    if (err != JVMTI_ERROR_NONE) {
        return JNI_ERR;
    }

    // Find the OutOfMemoryError class
    JNIEnv *jni_env;
    res = (*vm)->GetEnv(vm, (void **)&jni_env, JNI_VERSION_1_2);
    if (res != JNI_OK) {
        return JNI_ERR;
    }
    outOfMemoryErrorClass = (*jni_env)->FindClass(jni_env, "java/lang/OutOfMemoryError");
    if (outOfMemoryErrorClass == NULL) {
        return JNI_ERR;
    }
    outOfMemoryErrorClass = (*jni_env)->NewGlobalRef(jni_env, outOfMemoryErrorClass);

    return JNI_OK;
}

JNIEXPORT void JNICALL Agent_OnUnload(JavaVM *vm) {
    // Disable the Exception and ExceptionCatch events
    (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_DISABLE, JVMTI_EVENT_EXCEPTION, NULL);
    (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_DISABLE, JVMTI_EVENT_EXCEPTION_CATCH, NULL);

    // Clean up global references
    JNIEnv *jni_env;
    if ((*vm)->GetEnv(vm, (void **)&jni_env, JNI_VERSION_1_2) == JNI_OK) {
        if (outOfMemoryErrorClass != NULL) {
            (*jni_env)->DeleteGlobalRef(jni_env, outOfMemoryErrorClass);
        }
    }
}