crates/client/src/jni.rs at c385e97d220ed804fd1e5c5eeaf27b48d0d82d40

kretrod/lodestone Public
Kretrod Dear ImGui Enigma menu, Linux support, silent-aim fix, crystal and mace c385e97 · 6 hours ago History
700 lines (646 loc) · 22.5 KB Raw
//! A thin JNI layer.
//!
//! Inside the process we no longer have to decode HotSpot's heap by hand: the
//! JVM exposes a proper API. `JNI_GetCreatedJavaVMs` hands us the VM that is
//! already running, and because our hook executes on the game's render thread —
//! itself a Java thread — `GetEnv` gives us a usable environment with no
//! attaching required.
//!
//! Two rules keep this safe:
//!   * every frame runs inside a local reference frame, so references we make
//!     while drawing are released together instead of leaking the heap away;
//!   * every call that can throw is followed by an exception check, because a
//!     pending exception makes the *next* JNI call undefined behaviour.

#![allow(non_snake_case)]

use std::ffi::{c_void, CString};

use jni_sys::{
    jboolean, jclass, jdouble, jfieldID, jfloat, jint, jmethodID, jobject, jvalue, JNIEnv, JavaVM,
    JNI_OK, JNI_VERSION_1_6,
};
#[cfg(windows)]
use windows_sys::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress};

type GetCreatedJavaVMs =
    unsafe extern "system" fn(*mut *mut JavaVM, jint, *mut jint) -> jint;

/// Look up `JNI_GetCreatedJavaVMs` in the JVM we are living inside.
#[cfg(windows)]
unsafe fn get_created_java_vms() -> Option<GetCreatedJavaVMs> {
    // SAFETY: jvm.dll is loaded (we are inside the JVM); the export is looked
    // up by name and null-checked before being called.
    let module = GetModuleHandleA(c"jvm.dll".as_ptr() as *const u8);
    if module.is_null() {
        return None;
    }
    let f = GetProcAddress(module, c"JNI_GetCreatedJavaVMs".as_ptr() as *const u8)?;
    Some(std::mem::transmute(f))
}

/// The same, through the dynamic linker. libjvm.so is already mapped into this
/// process and exports the symbol globally, so the default scope finds it; the
/// explicit open is a fallback for a JVM loaded into a private namespace.
#[cfg(unix)]
unsafe fn get_created_java_vms() -> Option<GetCreatedJavaVMs> {
    let name = c"JNI_GetCreatedJavaVMs";
    let mut p = libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr());
    if p.is_null() {
        let h = libc::dlopen(c"libjvm.so".as_ptr(), libc::RTLD_NOW | libc::RTLD_NOLOAD);
        if !h.is_null() {
            p = libc::dlsym(h, name.as_ptr());
        }
    }
    if p.is_null() {
        return None;
    }
    Some(std::mem::transmute(p))
}

/// The JVM already running in this process, or None if there isn't one.
pub fn java_vm() -> Option<*mut JavaVM> {
    // SAFETY: the export is looked up by name and null-checked before the call,
    // and the out-parameters are plain stack slots.
    unsafe {
        let f = get_created_java_vms()?;
        let mut vm: *mut JavaVM = std::ptr::null_mut();
        let mut count: jint = 0;
        if f(&mut vm, 1, &mut count) != JNI_OK || count < 1 || vm.is_null() {
            return None;
        }
        Some(vm)
    }
}

#[derive(Clone, Copy)]
pub struct Jni {
    pub env: *mut JNIEnv,
}

impl Jni {
    /// An environment for the calling thread, if it is already a Java thread.
    /// The render thread is, which is exactly where our hook runs.
    pub fn current(vm: *mut JavaVM) -> Option<Jni> {
        // SAFETY: vm came from JNI_GetCreatedJavaVMs and its vtable is valid.
        unsafe {
            let mut env: *mut c_void = std::ptr::null_mut();
            let get_env = (**vm).GetEnv?;
            if get_env(vm, &mut env, JNI_VERSION_1_6 as jint) != JNI_OK || env.is_null() {
                return None;
            }
            Some(Jni { env: env as *mut JNIEnv })
        }
    }

    /// Reserve a batch of local references; dropping the guard frees them all.
    pub fn frame(&self, capacity: jint) -> Option<Frame> {
        // SAFETY: env is valid for this thread.
        unsafe {
            let push = (**self.env).PushLocalFrame?;
            if push(self.env, capacity) != JNI_OK {
                return None;
            }
            Some(Frame { env: self.env })
        }
    }

    /// Clear any pending exception. A live exception poisons every later call,
    /// so this runs after anything that can throw.
    pub fn clear_exception(&self) -> bool {
        // SAFETY: env is valid; both entry points are always present.
        unsafe {
            let check = match (**self.env).ExceptionCheck {
                Some(f) => f,
                None => return false,
            };
            if check(self.env) == 0 {
                return false;
            }
            if let Some(clear) = (**self.env).ExceptionClear {
                clear(self.env);
            }
            true
        }
    }

    /// Look up a class by binary name ("net/minecraft/client/Minecraft").
    ///
    /// `FindClass` resolves against the class loader of the current native
    /// frame. In the render thread that is LWJGL's, which under Fabric is the
    /// same loader the game's classes live in — but if that ever stops being
    /// true we fall back to asking the thread's context class loader directly.
    pub fn find_class(&self, name: &str) -> Option<jclass> {
        // SAFETY: name is NUL-terminated for the call; result is checked.
        unsafe {
            let c = CString::new(name).ok()?;
            let find = (**self.env).FindClass?;
            let cls = find(self.env, c.as_ptr());
            if !cls.is_null() {
                return Some(cls);
            }
            self.clear_exception();
            self.find_class_via_context_loader(name)
        }
    }

    unsafe fn find_class_via_context_loader(&self, name: &str) -> Option<jclass> {
        let find = (**self.env).FindClass?;
        let thread_cls = find(self.env, c"java/lang/Thread".as_ptr());
        if thread_cls.is_null() {
            self.clear_exception();
            return None;
        }
        let current = self.static_method(thread_cls, "currentThread", "()Ljava/lang/Thread;")?;
        let call_obj = (**self.env).CallStaticObjectMethodA?;
        let thread = call_obj(self.env, thread_cls, current, std::ptr::null());
        if self.clear_exception() || thread.is_null() {
            return None;
        }
        let get_loader =
            self.method(thread_cls, "getContextClassLoader", "()Ljava/lang/ClassLoader;")?;
        let call = (**self.env).CallObjectMethodA?;
        let loader = call(self.env, thread, get_loader, std::ptr::null());
        if self.clear_exception() || loader.is_null() {
            return None;
        }
        let loader_cls = find(self.env, c"java/lang/ClassLoader".as_ptr());
        if loader_cls.is_null() {
            self.clear_exception();
            return None;
        }
        let load = self.method(loader_cls, "loadClass", "(Ljava/lang/String;)Ljava/lang/Class;")?;
        // loadClass wants dots, not slashes.
        let dotted = self.new_string(&name.replace('/', "."))?;
        let args = [jvalue { l: dotted }];
        let cls = call(self.env, loader, load, args.as_ptr());
        if self.clear_exception() || cls.is_null() {
            return None;
        }
        Some(cls as jclass)
    }

    /// Define a class into a running class loader from raw bytecode.
    ///
    /// This is how a handler written in Java gets into a game that was never
    /// built to accept one: compile it here, hand the bytes to the loader that
    /// already holds the game's classes, and it can see and be seen by them.
    pub fn define_class(&self, name: &str, loader: jobject, bytes: &[u8]) -> Option<jclass> {
        // SAFETY: name is NUL-terminated, and the buffer outlives the call.
        unsafe {
            let c = CString::new(name).ok()?;
            let f = (**self.env).DefineClass?;
            let cls = f(
                self.env,
                c.as_ptr(),
                loader,
                bytes.as_ptr() as *const i8,
                bytes.len() as i32,
            );
            if self.clear_exception() || cls.is_null() {
                None
            } else {
                Some(cls)
            }
        }
    }

    pub fn new_string(&self, s: &str) -> Option<jobject> {
        // SAFETY: the UTF-8 buffer is NUL-terminated and outlives the call.
        unsafe {
            let c = CString::new(s).ok()?;
            let f = (**self.env).NewStringUTF?;
            let o = f(self.env, c.as_ptr());
            if self.clear_exception() || o.is_null() {
                None
            } else {
                Some(o)
            }
        }
    }

    /// Release one local reference early. Scanning thousands of entities in a
    /// frame would otherwise pile up references until the frame ends.
    pub fn delete_local(&self, obj: jobject) {
        if obj.is_null() {
            return;
        }
        // SAFETY: obj is a live local reference created on this thread.
        unsafe {
            if let Some(f) = (**self.env).DeleteLocalRef {
                f(self.env, obj);
            }
        }
    }

    /// Promote a reference so it survives past the current frame.
    pub fn global(&self, obj: jobject) -> Option<jobject> {
        // SAFETY: obj is a live local reference.
        unsafe {
            let f = (**self.env).NewGlobalRef?;
            let g = f(self.env, obj);
            if g.is_null() {
                None
            } else {
                Some(g)
            }
        }
    }

    /// Release a global reference. Pairs with `global`.
    pub fn delete_global(&self, obj: jobject) {
        if obj.is_null() {
            return;
        }
        // SAFETY: obj came from NewGlobalRef and is released exactly once.
        unsafe {
            if let Some(f) = (**self.env).DeleteGlobalRef {
                f(self.env, obj);
            }
        }
    }

    pub fn field(&self, cls: jclass, name: &str, sig: &str) -> Option<jfieldID> {
        // SAFETY: both strings are NUL-terminated; a miss throws, so clear it.
        unsafe {
            let n = CString::new(name).ok()?;
            let s = CString::new(sig).ok()?;
            let f = (**self.env).GetFieldID?;
            let id = f(self.env, cls, n.as_ptr(), s.as_ptr());
            if self.clear_exception() || id.is_null() {
                None
            } else {
                Some(id)
            }
        }
    }

    pub fn static_field(&self, cls: jclass, name: &str, sig: &str) -> Option<jfieldID> {
        // SAFETY: as above.
        unsafe {
            let n = CString::new(name).ok()?;
            let s = CString::new(sig).ok()?;
            let f = (**self.env).GetStaticFieldID?;
            let id = f(self.env, cls, n.as_ptr(), s.as_ptr());
            if self.clear_exception() || id.is_null() {
                None
            } else {
                Some(id)
            }
        }
    }

    pub fn method(&self, cls: jclass, name: &str, sig: &str) -> Option<jmethodID> {
        // SAFETY: as above.
        unsafe {
            let n = CString::new(name).ok()?;
            let s = CString::new(sig).ok()?;
            let f = (**self.env).GetMethodID?;
            let id = f(self.env, cls, n.as_ptr(), s.as_ptr());
            if self.clear_exception() || id.is_null() {
                None
            } else {
                Some(id)
            }
        }
    }

    pub fn static_method(&self, cls: jclass, name: &str, sig: &str) -> Option<jmethodID> {
        // SAFETY: as above.
        unsafe {
            let n = CString::new(name).ok()?;
            let s = CString::new(sig).ok()?;
            let f = (**self.env).GetStaticMethodID?;
            let id = f(self.env, cls, n.as_ptr(), s.as_ptr());
            if self.clear_exception() || id.is_null() {
                None
            } else {
                Some(id)
            }
        }
    }

    // ---- reads -------------------------------------------------------------

    pub fn obj_field(&self, obj: jobject, id: jfieldID) -> Option<jobject> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: obj is live and id was resolved against its class.
        unsafe {
            let f = (**self.env).GetObjectField?;
            let v = f(self.env, obj, id);
            if v.is_null() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn static_obj_field(&self, cls: jclass, id: jfieldID) -> Option<jobject> {
        // SAFETY: cls is live and id was resolved against it.
        unsafe {
            let f = (**self.env).GetStaticObjectField?;
            let v = f(self.env, cls, id);
            if v.is_null() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn bool_field(&self, obj: jobject, id: jfieldID) -> Option<bool> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: obj is live and id matches its class.
        unsafe { Some((**self.env).GetBooleanField?(self.env, obj, id) != 0) }
    }

    pub fn int_field(&self, obj: jobject, id: jfieldID) -> Option<i32> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe { Some((**self.env).GetIntField?(self.env, obj, id)) }
    }

    pub fn float_field(&self, obj: jobject, id: jfieldID) -> Option<f32> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe { Some((**self.env).GetFloatField?(self.env, obj, id)) }
    }

    pub fn double_field(&self, obj: jobject, id: jfieldID) -> Option<f64> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe { Some((**self.env).GetDoubleField?(self.env, obj, id)) }
    }

    // ---- writes ------------------------------------------------------------
    //
    // Unlike the external build, these go through the JVM, so the GC's write
    // barriers run and object fields are as safe to set as primitive ones.

    pub fn set_bool(&self, obj: jobject, id: jfieldID, v: bool) {
        if obj.is_null() {
            return;
        }
        // SAFETY: obj is live and id matches its class.
        unsafe {
            if let Some(f) = (**self.env).SetBooleanField {
                f(self.env, obj, id, v as jboolean);
            }
        }
    }

    pub fn set_int(&self, obj: jobject, id: jfieldID, v: i32) {
        if obj.is_null() {
            return;
        }
        // SAFETY: as above.
        unsafe {
            if let Some(f) = (**self.env).SetIntField {
                f(self.env, obj, id, v);
            }
        }
    }

    pub fn set_float(&self, obj: jobject, id: jfieldID, v: f32) {
        if obj.is_null() {
            return;
        }
        // SAFETY: as above.
        unsafe {
            if let Some(f) = (**self.env).SetFloatField {
                f(self.env, obj, id, v as jfloat);
            }
        }
    }

    pub fn set_double(&self, obj: jobject, id: jfieldID, v: f64) {
        if obj.is_null() {
            return;
        }
        // SAFETY: as above.
        unsafe {
            if let Some(f) = (**self.env).SetDoubleField {
                f(self.env, obj, id, v as jdouble);
            }
        }
    }

    pub fn set_static_bool(&self, cls: jclass, id: jfieldID, v: bool) {
        // SAFETY: cls declares the field and it is a boolean.
        unsafe {
            if let Some(f) = (**self.env).SetStaticBooleanField {
                f(self.env, cls, id, v as jboolean);
            }
        }
    }

    pub fn set_static_int(&self, cls: jclass, id: jfieldID, v: i32) {
        // SAFETY: cls declares the static int field.
        unsafe {
            if let Some(f) = (**self.env).SetStaticIntField {
                f(self.env, cls, id, v);
            }
        }
    }

    pub fn set_static_long(&self, cls: jclass, id: jfieldID, v: i64) {
        // SAFETY: cls declares the static long field.
        unsafe {
            if let Some(f) = (**self.env).SetStaticLongField {
                f(self.env, cls, id, v);
            }
        }
    }

    pub fn set_obj(&self, obj: jobject, id: jfieldID, v: jobject) {
        if obj.is_null() {
            return;
        }
        // SAFETY: as above; v may legitimately be null.
        unsafe {
            if let Some(f) = (**self.env).SetObjectField {
                f(self.env, obj, id, v);
            }
        }
    }

    // ---- calls -------------------------------------------------------------

    pub fn call_void(&self, obj: jobject, id: jmethodID, args: &[jvalue]) {
        if obj.is_null() {
            return;
        }
        // SAFETY: args matches the method's descriptor at every call site.
        unsafe {
            if let Some(f) = (**self.env).CallVoidMethodA {
                f(self.env, obj, id, args.as_ptr());
            }
            self.clear_exception();
        }
    }

    pub fn call_obj(&self, obj: jobject, id: jmethodID, args: &[jvalue]) -> Option<jobject> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe {
            let f = (**self.env).CallObjectMethodA?;
            let v = f(self.env, obj, id, args.as_ptr());
            if self.clear_exception() || v.is_null() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn call_static_obj(
        &self,
        cls: jclass,
        id: jmethodID,
        args: &[jvalue],
    ) -> Option<jobject> {
        // SAFETY: args matches the method's descriptor at every call site.
        unsafe {
            let f = (**self.env).CallStaticObjectMethodA?;
            let v = f(self.env, cls, id, args.as_ptr());
            if self.clear_exception() || v.is_null() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn call_static_int(&self, cls: jclass, id: jmethodID, args: &[jvalue]) -> Option<i32> {
        // SAFETY: args matches the method's descriptor at every call site.
        unsafe {
            let f = (**self.env).CallStaticIntMethodA?;
            let v = f(self.env, cls, id, args.as_ptr());
            if self.clear_exception() {
                None
            } else {
                Some(v)
            }
        }
    }

    /// Construct a Java object. Used for the block positions the level wants.
    pub fn new_object(&self, cls: jclass, ctor: jmethodID, args: &[jvalue]) -> Option<jobject> {
        // SAFETY: ctor belongs to cls and args matches its descriptor.
        unsafe {
            let f = (**self.env).NewObjectA?;
            let v = f(self.env, cls, ctor, args.as_ptr());
            if self.clear_exception() || v.is_null() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn call_bool(&self, obj: jobject, id: jmethodID, args: &[jvalue]) -> Option<bool> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe {
            let f = (**self.env).CallBooleanMethodA?;
            let v = f(self.env, obj, id, args.as_ptr());
            if self.clear_exception() {
                None
            } else {
                Some(v != 0)
            }
        }
    }

    /// Call a method without virtual dispatch — which is how a private method
    /// has to be invoked, since it has no vtable slot to look up.
    pub fn call_nonvirtual_bool(
        &self,
        obj: jobject,
        cls: jclass,
        id: jmethodID,
        args: &[jvalue],
    ) -> Option<bool> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: cls declares the method and args matches its descriptor.
        unsafe {
            let f = (**self.env).CallNonvirtualBooleanMethodA?;
            let v = f(self.env, obj, cls, id, args.as_ptr());
            if self.clear_exception() {
                None
            } else {
                Some(v != 0)
            }
        }
    }

    pub fn call_float(&self, obj: jobject, id: jmethodID, args: &[jvalue]) -> Option<f32> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe {
            let f = (**self.env).CallFloatMethodA?;
            let v = f(self.env, obj, id, args.as_ptr());
            if self.clear_exception() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn call_double(&self, obj: jobject, id: jmethodID, args: &[jvalue]) -> Option<f64> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe {
            let f = (**self.env).CallDoubleMethodA?;
            let v = f(self.env, obj, id, args.as_ptr());
            if self.clear_exception() {
                None
            } else {
                Some(v)
            }
        }
    }

    pub fn call_int(&self, obj: jobject, id: jmethodID, args: &[jvalue]) -> Option<i32> {
        if obj.is_null() {
            return None;
        }
        // SAFETY: as above.
        unsafe {
            let f = (**self.env).CallIntMethodA?;
            let v = f(self.env, obj, id, args.as_ptr());
            if self.clear_exception() {
                None
            } else {
                Some(v)
            }
        }
    }

    /// Reference identity, used to keep the local player out of scans.
    pub fn same_object(&self, a: jobject, b: jobject) -> bool {
        // SAFETY: both are live references (or null, which is allowed).
        unsafe {
            match (**self.env).IsSameObject {
                Some(f) => f(self.env, a, b) != 0,
                None => false,
            }
        }
    }

    /// `obj instanceof cls`
    pub fn is_instance(&self, obj: jobject, cls: jclass) -> bool {
        if obj.is_null() || cls.is_null() {
            return false;
        }
        // SAFETY: both references are live.
        unsafe {
            match (**self.env).IsInstanceOf {
                Some(f) => f(self.env, obj, cls) != 0,
                None => false,
            }
        }
    }

    pub fn rust_string(&self, s: jobject) -> Option<String> {
        if s.is_null() {
            return None;
        }
        // SAFETY: s is a java.lang.String; the buffer is released before return.
        unsafe {
            let get = (**self.env).GetStringUTFChars?;
            let release = (**self.env).ReleaseStringUTFChars?;
            let ptr = get(self.env, s, std::ptr::null_mut());
            if ptr.is_null() {
                return None;
            }
            let out = std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned();
            release(self.env, s, ptr);
            Some(out)
        }
    }
}

/// Releases every local reference made inside it.
pub struct Frame {
    env: *mut JNIEnv,
}

impl Drop for Frame {
    fn drop(&mut self) {
        // SAFETY: matched with the PushLocalFrame that produced this guard.
        unsafe {
            if let Some(pop) = (**self.env).PopLocalFrame {
                pop(self.env, std::ptr::null_mut());
            }
        }
    }
}