Creating a JVMTI Agent with Rust

I think there is little demand for this, but I will write about how to create a Java JVMTI agent using Rust.

What is JVMTI?

JVMTI stands for JVM Tool Interface.
By specifying it at startup or dynamically attaching the JVMTI agent after startup, you can monitor or modify the behavior of Java.

How to Create It

Prerequisites

$ cat /etc/os-release
NAME="Arch Linux"
...
$ cargo --version
cargo 1.51.0 (43b129a20 2021-03-16)
$ cargo install --list|grep cargo-edit
cargo-edit v0.7.0:
$ pacman -Q jdk-openjdk jre-openjdk
jdk-openjdk 15.0.2.u7-1
jre-openjdk 15.0.2.u7-1

Preparing the Crate

First, create the crate.

$ cargo new --lib jvmti-study
     Created library `jvmti-study` package
$ cd jvmti-study

Configure the crate for a dynamic library.

Cargo.toml

...
[lib]
crate-type = ["cdylib"]
...

bindgen

Create a sys-module using bindgen.
First, add bindgen to the build dependencies.

$ cargo add --build bindgen
    Updating 'https://github.com/rust-lang/crates.io-index' index
      Adding bindgen v0.57.0 to build-dependencies

Create a build.rs script.

build.rs

use std::env;
use std::path;

fn main() {
    const LIB: &str = "/usr/lib/jvm/default-runtime/lib/server";
    const INCLUDE: &str = "/usr/lib/jvm/default/include";
    const INCLUDE_LINUX: &str = "/usr/lib/jvm/default/include/linux";

    println!("cargo:rustc-link-lib=jvm");
    println!("cargo:rustc-link-search=native={}", LIB);

    let bindings = bindgen::builder()
        .header_contents("bindings.h", "#include <jvmti.h>")
        .clang_arg(format!("-I{}", INCLUDE))
        .clang_arg(format!("-I{}", INCLUDE_LINUX))
        .derive_debug(true)
        .derive_default(true)
        .parse_callbacks(Box::new(bindgen::CargoCallbacks))
        .generate()
        .expect("failed to generate bindgen.");

    let out_path = path::PathBuf::from(env::var("OUT_DIR").unwrap());
    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("failed to write bindings.rs.");
}

Embed the generated code into the source code.

src/lib.rs

#[allow(non_upper_case_globals)]
#[allow(non_camel_case_types)]
#[allow(non_snake_case)]
#[allow(unused)]
mod sys {
    include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
}

Implementing a Minimal Agent

The entry points from the JVM are as follows.
You only need to implement the necessary parts.

• Agent_OnLoad
  • This is the entry point for an agent that runs at startup.
  • Since the JVM is not yet fully initialized, you cannot access JNI functions at this point.
• Agent_OnAttach
  • This is the entry point for an agent that is dynamically attached while the JVM is running.
• Agent_OnUnload
  • This is the entry point when the JVMTI agent finishes.
  • Since you cannot access the jvmtiEnv created during OnLoad / OnAttach, it is usually easier to implement cleanup processing in the VmDeath event if needed.

src/lib.rs

...
#[no_mangle]
pub extern "C" fn Agent_OnLoad(
    _vm: *mut sys::JavaVM,
    _options: *const std::os::raw::c_char,
    _reserved: *const std::ffi::c_void,
) -> sys::jint {
    println!("Hello, JVMTI!");
    0
}

#[no_mangle]
pub extern "C" fn Agent_OnAttach(
    _vm: *mut sys::JavaVM,
    _options: *const std::os::raw::c_char,
    _reserved: *const std::ffi::c_void,
) -> sys::jint {
    println!("JVMTI Attached.");
    0
}

#[no_mangle]
pub extern "C" fn Agent_OnUnload(_vm: *mut sys::JavaVM) {
    println!("Good bye.");
}

Compile

$ cargo build
...
$ ls target/debug/libjvmti_study.so
target/debug/libjvmti_study.so

Verification of Operation

Java program for verification:

Test.java

public class Test {
    public static void main(String...args) throws Exception {
        System.out.println("Hello, world!");
        Thread.sleep(Long.MAX_VALUE);
    }
}

$ javac Test.java

First, verify the operation of the Java program itself:

$ java Test
Hello, world!
^C

Start the agent when launching Java:

$ java -agentpath:target/debug/libjvmti_study.so Test
Hello, JVMTI!
Hello, world!
^CGood bye.

Attach dynamically:

$ java Test &
[1] 84459
$ jcmd 84459 JVMTI.agent_load target/debug/libjvmti_study.so
84459:
JVMTI Attached.
return code: 0
$ kill 84459
Good bye.

Obtaining jvmtiEnv

The jvmtiEnv is a structure that stores the JVMTI functions, similar to JNIEnv in JNI.
You can obtain it using the GetEnv function, just like JNIEnv in JNI.
By setting the version passed to GetEnv to JVMTI_VERSION instead of a JNI version, you can obtain the jvmtiEnv.

src/lib.rs

#[no_mangle]
pub extern "C" fn Agent_OnLoad(
    vm: *mut sys::JavaVM,
    _options: *const std::os::raw::c_char,
    _reserved: *const std::ffi::c_void,
) -> sys::jint {
    println!("Hello, JVMTI!");

    let mut jvmti_env = std::ptr::null_mut();
    let jvmti_env = unsafe {
        let get_env = (**vm).GetEnv.unwrap();
        if get_env(vm, &mut jvmti_env, sys::JVMTI_VERSION as i32) != sys::JNI_OK as i32 {
            eprintln!("failed to get jvmtiEnv");
            return -1;
        }
        jvmti_env as *mut sys::jvmtiEnv
    };

    0
}

Hooking Events

You can hook various events fired by the JVM.
You can hook events using SetEventCallbacks and SetEventNotificationMode.

src/lib.rs

unsafe extern "C" fn on_vm_init(_jvmti_env: *mut sys::jvmtiEnv, _jni_env: *mut sys::JNIEnv, _thread: sys::jthread) {
    println!("on vm init.");
}

unsafe extern "C" fn on_vm_death(_jvmti_env: *mut sys::jvmtiEnv, _jni_env: *mut sys::JNIEnv) {
    println!("on vm death");
}
...
#[no_mangle]
pub extern "C" fn Agent_OnLoad(
    vm: *mut sys::JavaVM,
    options: *const std::os::raw::c_char,
    _reserved: *const std::ffi::c_void,
) -> sys::jint {
...
    let callbacks = sys::jvmtiEventCallbacks {
        VMInit: Some(on_vm_init),
        VMDeath: Some(on_vm_death),
        ..Default::default()
    };
    unsafe {
        let set_event_callback = (**jvmti_env).SetEventCallbacks.unwrap();
        if set_event_callback(jvmti_env, &callbacks, std::mem::size_of::<sys::jvmtiEventCallbacks>() as i32) != sys::JNI_OK {
            eprintln!("failed to set event callbacks.");
            return -1;
        }
    };
    unsafe {
        let set_event_notification_mode = (**jvmti_env).SetEventNotificationMode.unwrap();
        if set_event_notification_mode(jvmti_env, sys::jvmtiEventMode_JVMTI_ENABLE, sys::jvmtiEvent_JVMTI_EVENT_VM_INIT, std::ptr::null_mut()) != sys::JNI_OK {
            eprintln!("failed to set event notification mode.");
            return -1;
        }
        if set_event_notification_mode(jvmti_env, sys::jvmtiEventMode_JVMTI_ENABLE, sys::jvmtiEvent_JVMTI_EVENT_VM_DEATH, std::ptr::null_mut()) != sys::JNI_OK {
            eprintln!("failed to set event notification mode.");
            return -1;
        }
    };
...
}

$ java -agentpath:target/debug/libjvmti_study.so Test
Hello, JVMTI!
on vm init.
Hello, world!
^Con vm death
Good bye.

Environment per JVMTI Agent

You can set and retrieve environment data for each JVMTI agent using the following functions. The JVM manages a pointer to the environment for each agent.

• SetEnvironmentLocalStorage
• GetEnvironmentLocalStorage

src/lib.rs

unsafe extern "C" fn on_vm_init(jvmti_env: *mut sys::jvmtiEnv, _jni_env: *mut sys::JNIEnv, _thread: sys::jthread) {
    println!("on vm init.");

    let get_environment_local_storage = (**jvmti_env).GetEnvironmentLocalStorage.unwrap();
    let mut env = std::ptr::null_mut();
    if get_environment_local_storage(jvmti_env, &mut env) != sys::JNI_OK {
        eprintln!("failed to set environment local storage.");
        return
    }
    let env = std::sync::Arc::<MyEnv>::from_raw(env as _);
    println!("env: {:?}", env);
}
...
#[no_mangle]
pub extern "C" fn Agent_OnLoad(
    vm: *mut sys::JavaVM,
    options: *const std::os::raw::c_char,
    _reserved: *const std::ffi::c_void,
) -> sys::jint {
...
    let options = (!options.is_null()).then(|| unsafe {
        std::ffi::CStr::from_ptr(options).to_string_lossy().to_string()
    });
    let env = std::sync::Arc::new(MyEnv { options, });
    let env = std::sync::Arc::into_raw(env);

    unsafe {
        let set_environment_local_storage = (**jvmti_env).SetEnvironmentLocalStorage.unwrap();
        if set_environment_local_storage(jvmti_env, env as _) != sys::JNI_OK {
            eprintln!("failed to set environment local storage.");
            return -1;
        }
    }
...
}

Conclusion

With this much capability, it seems like you could create something quite interesting.