agent.cpp

#define IIF(cond) IIF_ ## cond
#define IIF_false(trueBranch, falseBranch) falseBranch
#define IIF_true(trueBranch, falseBranch) trueBranch

#define WRAP_METHOD(methodName, isStatic, sig, returnType, parameters...) ({ \
    static const auto methodId = env->IIF(isStatic)(GetStaticMethodID, GetMethodID)(clazz, methodName, sig); \
    returnType(*wrapper)(JNIEnv*, jobject, ##parameters) = [](auto env, auto self, auto...params) { \
        if (isBlockingCall(jvmti)) { \
            reportBlockingCall(jvmti, env, methodId); \
        } else { \
            static const auto method = (returnType(*)(JNIEnv*, jobject, ...)) originalMethods[methodId]; \
            return method(env, self, params...); \
        } \
    }; \
    replacements[methodId] = (void *) wrapper; \
    if (originalMethods.find(methodId) != originalMethods.end()) { \
        overrides.push_back({ \
            .name = methodName, \
            .signature = sig, \
            .fnPtr = (void *) wrapper \
        }); \
    } \
})

static void JNICALL callbackNativeMethodBindEvent(jvmtiEnv *, JNIEnv*, jthread, jmethodID method, void* address, void** newAddress) {
    if (originalMethods.find(method) == originalMethods.end()) {
        originalMethods[method] = address;
    }

    auto replacementIterator = replacements.find(method);
    if (replacementIterator != replacements.end()) {
        *newAddress = replacementIterator++->second;
    }
}

void JNICALL callbackVMInitEvent(jvmtiEnv *, JNIEnv* env, jthread thread) {
    int javaVersion;
    jvmti->GetVersionNumber(&javaVersion);
    javaVersion = (javaVersion & JVMTI_VERSION_MASK_MAJOR) >> JVMTI_VERSION_SHIFT_MAJOR;
    // printf("VM (version: %d) initiated\n", javaVersion);

    {
        auto clazz = env->FindClass("java/lang/Thread");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("sleep", true, "(J)V", void, jlong);
        WRAP_METHOD("yield", true, "()V", void);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/lang/Object");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("wait", false, "(J)V", void, jlong);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/io/RandomAccessFile");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("read0", false, "()I", int);
        WRAP_METHOD("readBytes", false, "([BII)I", int, jbyteArray,jint,jint);
        WRAP_METHOD("write0", false, "(I)V", void, jint);
        WRAP_METHOD("writeBytes", false, "([BII)V", void, jbyteArray,jint,jint);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/net/PlainDatagramSocketImpl");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("connect0", false, "(Ljava/net/InetAddress;I)V", void, jobject,jint);
        WRAP_METHOD("peekData", false, "(Ljava/net/DatagramPacket;)I", int, jobject);
        WRAP_METHOD("send", false, "(Ljava/net/DatagramPacket;)V", void, jobject);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/net/PlainSocketImpl");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("socketAccept", false, "(Ljava/net/SocketImpl;)V", void, jobject,jint);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass(javaVersion >= 9 ? "java/lang/ProcessImpl" : "java/lang/UNIXProcess");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("forkAndExec", false, "(I[B[B[BI[BI[B[IZ)I", int, int,jbyteArray,jbyteArray,jbyteArray,int,jbyteArray,int,jbyteArray,jintArray,jboolean);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/net/SocketInputStream");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("socketRead0", false, "(Ljava/io/FileDescriptor;[BIII)I", int, jobject,jbyteArray,jint,jint,jint);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/net/SocketOutputStream");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("socketWrite0", false, "(Ljava/io/FileDescriptor;[BII)V", void, jobject,jbyteArray,jint,jint);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/io/FileInputStream");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("read0", false, "()I", int);
        WRAP_METHOD("readBytes", false, "([BII)I", int, jbyteArray,jint,jint);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass("java/io/FileOutputStream");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("write", false, "(IZ)V", int, jint,jboolean);
        WRAP_METHOD("writeBytes", false, "([BIIZ)V", int, jbyteArray,jint,jint,jboolean);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    {
        auto clazz = env->FindClass(javaVersion >= 9 ? "jdk/internal/misc/Unsafe" : "sun/misc/Unsafe");
        std::vector<JNINativeMethod> overrides;
        WRAP_METHOD("park", false, "(ZJ)V", void, jboolean,jlong);
        env->RegisterNatives(clazz, overrides.data(), overrides.size());
    }

    // printf("Native methods registered\n");
}