Sign in Sign up
kretrod/lodestone Public
Branches
master
645 lines (599 loc) · 23.5 KB Raw
//! The object layer: turn raw target memory into Java classes, objects and
//! fields, using only what VMStructs told us.
//!
//! Chain of reasoning for a single field read, e.g. `Minecraft.player.xo`:
//!
//!   ClassLoaderDataGraph::_head          (static address, from VMStructs)
//!     -> ClassLoaderData::_klasses       -> Klass* linked list via _next_link
//!        -> Klass::_name                 -> Symbol* -> "net/minecraft/..."
//!     -> InstanceKlass::_fieldinfo_stream -> UNSIGNED5 records -> field offsets
//!     -> Klass::_java_mirror             -> the Class object, where statics live
//!        -> read static oop              -> the Minecraft instance
//!           -> read instance oop field   -> the player
//!              -> read double field      -> the coordinate
//!
//! Metadata (Klass, Symbol, the field stream) lives in Metaspace and never
//! moves. Only oops move, and only at a GC — so every resolve starts from a
//! Klass and re-walks, rather than caching an object address.

use std::collections::HashMap;

use crate::vm::VmDb;
use crate::win::{Module, Process, Result};

/// Every offset and encoding parameter we need, read once at attach.
pub struct Layout {
    pub klass_name: u64,
    pub klass_mirror: u64,
    pub klass_next_link: u64,
    pub klass_super: u64,
    pub klass_layout_helper: u64,
    pub ik_fieldinfo: u64,
    pub ik_constants: u64,
    pub cld_klasses: u64,
    pub cld_next: u64,
    pub cldg_head: u64,
    pub sym_length: u64,
    pub sym_body: u64,
    pub arr_length: u64,
    pub arr_u1_data: u64,
    pub cp_header: u64,
    /// Base of HotSpot's built-in symbol table: injected fields name themselves
    /// with indices into this, not into the class's constant pool.
    pub vm_symbols: u64,
    pub oop_klass: u64,
    pub oop_base: u64,
    pub oop_shift: i32,
    pub klass_base: u64,
    pub klass_shift: i32,
    /// Offset of the length field in an array oop, and where elements start.
    pub array_length: u64,
    pub array_data: u64,
}

#[derive(Debug, Clone)]
pub struct Field {
    pub name: String,
    pub sig: String,
    pub offset: u64,
    pub access: u16,
    pub injected: bool,
}

impl Field {
    pub fn is_static(&self) -> bool {
        self.access & 0x0008 != 0
    }
    pub fn kind(&self) -> u8 {
        self.sig.as_bytes().first().copied().unwrap_or(b'?')
    }
}

/// A decoded Java value.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum JVal {
    Bool(bool),
    Byte(i8),
    Char(u16),
    Short(i16),
    Int(i32),
    Long(i64),
    Float(f32),
    Double(f64),
    /// Reference: absolute (decompressed) address of the target oop, 0 = null.
    Obj(u64),
}

impl std::fmt::Display for JVal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            JVal::Bool(v) => write!(f, "{v}"),
            JVal::Byte(v) => write!(f, "{v}"),
            JVal::Char(v) => write!(f, "{:?}", char::from_u32(*v as u32).unwrap_or('?')),
            JVal::Short(v) => write!(f, "{v}"),
            JVal::Int(v) => write!(f, "{v}"),
            JVal::Long(v) => write!(f, "{v}"),
            JVal::Float(v) => write!(f, "{v}"),
            JVal::Double(v) => write!(f, "{v}"),
            JVal::Obj(0) => write!(f, "null"),
            JVal::Obj(a) => write!(f, "@{a:#x}"),
        }
    }
}

pub struct Jvm {
    pub p: Process,
    pub db: VmDb,
    pub l: Layout,
    /// class name (with '/' separators) -> Klass* in Metaspace.
    classes: HashMap<String, u64>,
    /// Klass* -> declared fields, decoded once.
    field_cache: HashMap<u64, Vec<Field>>,
}

impl Jvm {
    pub fn attach(p: Process, jvm: &Module) -> Result<Self> {
        let db = VmDb::load(&p, jvm)?;
        let oop_base = p.u64(db.static_addr("CompressedOops", "_base")?)?;
        let oop_shift = p.i32(db.static_addr("CompressedOops", "_shift")?)?;
        let klass_base = p.u64(db.static_addr("CompressedKlassPointers", "_base")?)?;
        let klass_shift = p.i32(db.static_addr("CompressedKlassPointers", "_shift")?)?;
        let oop_klass = db.off("oopDesc", "_metadata._compressed_klass")?;
        // Arrays put their length in the gap right after the narrow klass, and
        // elements start at the (aligned) end of the header.
        let array_length = oop_klass + 4;
        let array_data = (array_length + 4 + 7) & !7;

        let l = Layout {
            klass_name: db.off("Klass", "_name")?,
            klass_mirror: db.off("Klass", "_java_mirror")?,
            klass_next_link: db.off("Klass", "_next_link")?,
            klass_super: db.off("Klass", "_super")?,
            klass_layout_helper: db.off("Klass", "_layout_helper")?,
            ik_fieldinfo: db.off("InstanceKlass", "_fieldinfo_stream")?,
            ik_constants: db.off("InstanceKlass", "_constants")?,
            cld_klasses: db.off("ClassLoaderData", "_klasses")?,
            cld_next: db.off("ClassLoaderData", "_next")?,
            cldg_head: db.static_addr("ClassLoaderDataGraph", "_head")?,
            sym_length: db.off("Symbol", "_length")?,
            sym_body: db.off("Symbol", "_body[0]")?,
            arr_length: db.off("Array<int>", "_length")?,
            arr_u1_data: db.off("Array<u1>", "_data")?,
            cp_header: db.type_size("ConstantPool")?,
            vm_symbols: db.static_addr("Symbol", "_vm_symbols[0]")?,
            oop_klass,
            oop_base,
            oop_shift,
            klass_base,
            klass_shift,
            array_length,
            array_data,
        };
        Ok(Self { p, db, l, classes: HashMap::new(), field_cache: HashMap::new() })
    }

    // ---- pointer encodings -------------------------------------------------

    pub fn decode_oop(&self, narrow: u32) -> u64 {
        if narrow == 0 {
            0
        } else {
            self.l.oop_base + ((narrow as u64) << self.l.oop_shift)
        }
    }

    pub fn encode_oop(&self, oop: u64) -> u32 {
        if oop == 0 {
            0
        } else {
            ((oop - self.l.oop_base) >> self.l.oop_shift) as u32
        }
    }

    pub fn decode_klass(&self, narrow: u32) -> u64 {
        if narrow == 0 {
            0
        } else {
            self.l.klass_base + ((narrow as u64) << self.l.klass_shift)
        }
    }

    /// The Klass of a live object.
    pub fn klass_of(&self, oop: u64) -> Result<u64> {
        Ok(self.decode_klass(self.p.u32(oop + self.l.oop_klass)?))
    }

    // ---- metadata ----------------------------------------------------------

    /// Symbol* -> text. Symbols are modified-UTF8; class/field names are ASCII
    /// in practice, so a lossy decode is honest here.
    pub fn symbol(&self, sym: u64) -> Result<String> {
        if sym == 0 {
            return Ok(String::new());
        }
        let len = self.p.u16(sym + self.l.sym_length)? as usize;
        if len > 4096 {
            return Err(format!("implausible symbol length {len} at {sym:#x}"));
        }
        let b = self.p.read_bytes(sym + self.l.sym_body, len)?;
        Ok(String::from_utf8_lossy(&b).into_owned())
    }

    pub fn klass_name(&self, klass: u64) -> Result<String> {
        self.symbol(self.p.u64(klass + self.l.klass_name)?)
    }

    /// The java.lang.Class object for a Klass — where its static fields live.
    /// `_java_mirror` is an OopHandle: a pointer to a fixed slot in OopStorage
    /// that holds a full-width oop. The slot never moves, the oop inside can.
    pub fn mirror(&self, klass: u64) -> Result<u64> {
        let handle = self.p.u64(klass + self.l.klass_mirror)?;
        if handle == 0 {
            return Ok(0);
        }
        self.p.u64(handle)
    }

    pub fn is_instance_klass(&self, klass: u64) -> Result<bool> {
        // layout_helper is a positive instance size for instance klasses and a
        // negative tagged word for array klasses.
        Ok(self.p.i32(klass + self.l.klass_layout_helper)? > 0)
    }

    /// Walk ClassLoaderDataGraph and index every loaded class by name.
    /// Classes keep loading as the game runs, so this can be re-run.
    pub fn load_classes(&mut self) -> Result<usize> {
        let mut map = HashMap::new();
        let mut cld = self.p.u64(self.l.cldg_head)?;
        let mut guard = 0;
        while cld != 0 && guard < 10_000 {
            guard += 1;
            let mut k = self.p.u64(cld + self.l.cld_klasses).unwrap_or(0);
            let mut kguard = 0;
            while k != 0 && kguard < 500_000 {
                kguard += 1;
                if let Ok(name) = self.klass_name(k) {
                    map.entry(name).or_insert(k);
                }
                k = match self.p.u64(k + self.l.klass_next_link) {
                    Ok(n) => n,
                    Err(_) => break,
                };
            }
            cld = self.p.u64(cld + self.l.cld_next).unwrap_or(0);
        }
        self.classes = map;
        Ok(self.classes.len())
    }

    pub fn classes(&self) -> &HashMap<String, u64> {
        &self.classes
    }

    /// Look up by internal name ("net/minecraft/client/Minecraft"); dots are
    /// accepted for convenience.
    pub fn class(&self, name: &str) -> Option<u64> {
        let n = name.replace('.', "/");
        self.classes.get(&n).copied()
    }

    // ---- field records -----------------------------------------------------

    /// ConstantPool slot -> Symbol*. CP entries are one word each, laid out
    /// immediately after the ConstantPool header.
    /// Symbol from HotSpot's own vmSymbols table (injected fields).
    fn vm_symbol(&self, index: u16) -> Result<String> {
        let sym = self.p.u64(self.l.vm_symbols + (index as u64) * 8)?;
        self.symbol(sym)
            .map_err(|e| format!("vmSymbols[{index}] -> {sym:#x}: {e}"))
    }

    fn cp_symbol(&self, cp: u64, index: u16) -> Result<String> {
        let slot = cp + self.l.cp_header + (index as u64) * 8;
        let sym = self.p.u64(slot)?;
        self.symbol(sym)
            .map_err(|e| format!("cp[{index}] slot {slot:#x} -> {sym:#x}: {e}"))
    }

    /// Decode `InstanceKlass::_fieldinfo_stream` — the UNSIGNED5-packed field
    /// records introduced in JDK 21. Format (from fieldInfo.hpp):
    ///
    ///   Stream := num_java_fields num_injected_fields Field[j+k] End
    ///   Field  := name sig offset access flags Optionals(flags)
    ///   Optionals := initval?[initialized] gsig?[generic] group?[contended]
    pub fn fields(&mut self, klass: u64) -> Result<Vec<Field>> {
        if let Some(v) = self.field_cache.get(&klass) {
            return Ok(v.clone());
        }
        let stream = self.p.u64(klass + self.l.ik_fieldinfo)?;
        let cp = self.p.u64(klass + self.l.ik_constants)?;
        if stream == 0 || cp == 0 {
            return Ok(Vec::new());
        }
        let len = self.p.i32(stream + self.l.arr_length)? as usize;
        if len > 1 << 20 {
            return Err(format!("implausible fieldinfo stream length {len}"));
        }
        let data = self.p.read_bytes(stream + self.l.arr_u1_data, len)?;

        let mut pos = 0usize;
        let java_fields = u5(&data, &mut pos)?;
        let injected_count = u5(&data, &mut pos)?;
        let total = java_fields as usize + injected_count as usize;
        let mut out = Vec::with_capacity(total);
        for _ in 0..total {
            let name_idx = u5(&data, &mut pos)? as u16;
            let sig_idx = u5(&data, &mut pos)? as u16;
            let offset = u5(&data, &mut pos)? as u64;
            let access = u5(&data, &mut pos)? as u16;
            let flags = u5(&data, &mut pos)?;
            // Optional trailing items, present only per their flag bit.
            if flags & (1 << 0) != 0 {
                u5(&data, &mut pos)?; // ConstantValue index
            }
            if flags & (1 << 2) != 0 {
                u5(&data, &mut pos)?; // generic signature index
            }
            if flags & (1 << 4) != 0 {
                u5(&data, &mut pos)?; // @Contended group
            }
            // Injected fields are VM bookkeeping (java.lang.Class's klass
            // pointer, String's hash, Thread's eetop…). They name themselves
            // out of vmSymbols because they have no constant pool entry.
            let injected = flags & (1 << 1) != 0;
            let (name, sig) = if injected {
                (self.vm_symbol(name_idx)?, self.vm_symbol(sig_idx)?)
            } else {
                (self.cp_symbol(cp, name_idx)?, self.cp_symbol(cp, sig_idx)?)
            };
            out.push(Field { name, sig, offset, access, injected });
        }
        self.field_cache.insert(klass, out.clone());
        Ok(out)
    }

    /// Debug view: the raw stream bytes plus every u5 value in order.
    pub fn fieldinfo_raw(&mut self, klass: u64) -> Result<(Vec<u8>, Vec<u32>)> {
        let stream = self.p.u64(klass + self.l.ik_fieldinfo)?;
        let len = self.p.i32(stream + self.l.arr_length)? as usize;
        let data = self.p.read_bytes(stream + self.l.arr_u1_data, len)?;
        let mut pos = 0usize;
        let mut vals = Vec::new();
        while pos < data.len() {
            match u5(&data, &mut pos) {
                Ok(v) => vals.push(v),
                Err(_) => break,
            }
        }
        Ok((data, vals))
    }

    /// Find a field on a class or any superclass.
    pub fn find_field(&mut self, klass: u64, name: &str) -> Result<Field> {
        let mut k = klass;
        let mut guard = 0;
        while k != 0 && guard < 64 {
            guard += 1;
            for f in self.fields(k)? {
                if f.name == name {
                    return Ok(f);
                }
            }
            k = self.p.u64(k + self.l.klass_super)?;
        }
        Err(format!("no field {name:?} on {} or its supers", self.klass_name(klass)?))
    }

    /// All fields including inherited ones, most-derived first.
    pub fn all_fields(&mut self, klass: u64) -> Result<Vec<(String, Field)>> {
        let mut out = Vec::new();
        let mut k = klass;
        let mut guard = 0;
        while k != 0 && guard < 64 {
            guard += 1;
            let owner = self.klass_name(k)?;
            for f in self.fields(k)? {
                out.push((owner.clone(), f));
            }
            k = self.p.u64(k + self.l.klass_super)?;
        }
        Ok(out)
    }

    // ---- values ------------------------------------------------------------

    /// Read a field value at `base + field.offset`, typed by its signature.
    /// For statics `base` is the mirror; for instance fields it is the oop.
    pub fn read_at(&self, base: u64, f: &Field) -> Result<JVal> {
        let a = base + f.offset;
        Ok(match f.kind() {
            b'Z' => JVal::Bool(self.p.u8(a)? != 0),
            b'B' => JVal::Byte(self.p.u8(a)? as i8),
            b'C' => JVal::Char(self.p.u16(a)?),
            b'S' => JVal::Short(self.p.u16(a)? as i16),
            b'I' => JVal::Int(self.p.i32(a)?),
            b'J' => JVal::Long(self.p.u64(a)? as i64),
            b'F' => JVal::Float(self.p.f32(a)?),
            b'D' => JVal::Double(self.p.f64(a)?),
            b'L' | b'[' => JVal::Obj(if self.l.oop_shift >= 0 {
                self.decode_oop(self.p.u32(a)?)
            } else {
                self.p.u64(a)?
            }),
            k => return Err(format!("unhandled signature {:?}", k as char)),
        })
    }

    pub fn write_at(&self, base: u64, f: &Field, v: JVal) -> Result<()> {
        let a = base + f.offset;
        match v {
            JVal::Bool(x) => self.p.write_u8(a, x as u8),
            JVal::Byte(x) => self.p.write_u8(a, x as u8),
            JVal::Char(x) => self.p.write_bytes(a, &x.to_le_bytes()),
            JVal::Short(x) => self.p.write_bytes(a, &x.to_le_bytes()),
            JVal::Int(x) => self.p.write_i32(a, x),
            JVal::Long(x) => self.p.write_u64(a, x as u64),
            JVal::Float(x) => self.p.write_f32(a, x),
            JVal::Double(x) => self.p.write_f64(a, x),
            // Writing a reference needs the GC's store barriers; primitives
            // don't. We deliberately never write oops.
            JVal::Obj(_) => Err("refusing to write an object reference".into()),
        }
    }

    /// Read a static field by name off a class.
    pub fn static_field(&mut self, klass: u64, name: &str) -> Result<JVal> {
        let f = self.find_field(klass, name)?;
        if !f.is_static() {
            return Err(format!("{name} is not static"));
        }
        let mirror = self.mirror(klass)?;
        if mirror == 0 {
            return Err("class has no mirror yet (not initialised?)".into());
        }
        self.read_at(mirror, &f)
    }

    /// Read an instance field by name off an object.
    pub fn get(&mut self, oop: u64, name: &str) -> Result<JVal> {
        let k = self.klass_of(oop)?;
        let f = self.find_field(k, name)?;
        self.read_at(oop, &f)
    }

    pub fn set(&mut self, oop: u64, name: &str, v: JVal) -> Result<()> {
        let k = self.klass_of(oop)?;
        let f = self.find_field(k, name)?;
        self.write_at(oop, &f, v)
    }

    /// Follow a chain of instance field names: `get_path(obj, ["player","xo"])`.
    pub fn get_path(&mut self, root: u64, path: &[&str]) -> Result<JVal> {
        let mut cur = JVal::Obj(root);
        for (i, name) in path.iter().enumerate() {
            let oop = match cur {
                JVal::Obj(0) => return Err(format!("null at {:?}", &path[..i])),
                JVal::Obj(a) => a,
                other => return Err(format!("{other} is not an object at {:?}", &path[..i])),
            };
            cur = self.get(oop, name)?;
        }
        Ok(cur)
    }

    /// Breadth-first search of the object graph for instances whose class name
    /// contains `needle`. The way to find the player when you do not yet know
    /// what the field is called.
    pub fn search(
        &mut self,
        root: u64,
        needle: &str,
        max_depth: usize,
        max_visit: usize,
    ) -> Result<Vec<(String, String)>> {
        let mut seen = std::collections::HashSet::new();
        let mut queue = std::collections::VecDeque::new();
        let mut hits = Vec::new();
        queue.push_back((root, String::new(), 0usize));
        seen.insert(root);
        while let Some((oop, path, depth)) = queue.pop_front() {
            if seen.len() > max_visit {
                break;
            }
            let klass = match self.klass_of(oop) {
                Ok(k) if k != 0 => k,
                _ => continue,
            };
            let cname = match self.klass_name(klass) {
                Ok(n) => n,
                Err(_) => continue,
            };
            if !path.is_empty() && cname.to_lowercase().contains(&needle.to_lowercase()) {
                hits.push((path.clone(), cname.clone()));
            }
            if depth >= max_depth {
                continue;
            }
            // Arrays: follow elements. Objects: follow reference fields.
            if cname.starts_with("[L") || cname.starts_with("[[") {
                for (i, el) in self.obj_array(oop, 256)?.into_iter().enumerate() {
                    if el != 0 && seen.insert(el) {
                        queue.push_back((el, format!("{path}[{i}]"), depth + 1));
                    }
                }
                continue;
            }
            let fields = match self.all_fields(klass) {
                Ok(f) => f,
                Err(_) => continue,
            };
            for (_, f) in fields {
                if f.is_static() || !matches!(f.kind(), b'L' | b'[') {
                    continue;
                }
                if let Ok(JVal::Obj(next)) = self.read_at(oop, &f) {
                    if next != 0 && seen.insert(next) {
                        let p = if path.is_empty() {
                            f.name.clone()
                        } else {
                            format!("{path}.{}", f.name)
                        };
                        queue.push_back((next, p, depth + 1));
                    }
                }
            }
        }
        Ok(hits)
    }

    pub fn class_name_of(&self, oop: u64) -> Result<String> {
        self.klass_name(self.klass_of(oop)?)
    }

    // ---- arrays and strings ------------------------------------------------

    pub fn array_len(&self, arr: u64) -> Result<i32> {
        self.p.i32(arr + self.l.array_length)
    }

    /// Element addresses of an object array.
    pub fn obj_array(&self, arr: u64, max: usize) -> Result<Vec<u64>> {
        if arr == 0 {
            return Ok(Vec::new());
        }
        let n = (self.array_len(arr)?.max(0) as usize).min(max);
        let raw = self.p.read_bytes(arr + self.l.array_data, n * 4)?;
        Ok((0..n)
            .map(|i| {
                let nw = u32::from_le_bytes(raw[i * 4..i * 4 + 4].try_into().unwrap());
                self.decode_oop(nw)
            })
            .collect())
    }

    /// Entries of a java.util.HashMap as (key, value) oops. Walks the bucket
    /// array and each collision chain — the same structure HashMap itself uses.
    pub fn map_entries(&mut self, map: u64) -> Result<Vec<(u64, u64)>> {
        if map == 0 {
            return Ok(Vec::new());
        }
        let table = match self.get(map, "table")? {
            JVal::Obj(a) => a,
            _ => return Ok(Vec::new()),
        };
        if table == 0 {
            return Ok(Vec::new());
        }
        let mut out = Vec::new();
        for bucket in self.obj_array(table, 4096)? {
            let mut node = bucket;
            let mut guard = 0;
            while node != 0 && guard < 1024 {
                guard += 1;
                let k = match self.get(node, "key") {
                    Ok(JVal::Obj(a)) => a,
                    _ => 0,
                };
                let v = match self.get(node, "value") {
                    Ok(JVal::Obj(a)) => a,
                    _ => 0,
                };
                out.push((k, v));
                node = match self.get(node, "next") {
                    Ok(JVal::Obj(a)) => a,
                    _ => 0,
                };
            }
        }
        Ok(out)
    }

    /// java.lang.String -> Rust String (LATIN1 and UTF16 coders).
    pub fn read_string(&mut self, oop: u64) -> Result<String> {
        if oop == 0 {
            return Ok(String::new());
        }
        let value = match self.get(oop, "value")? {
            JVal::Obj(a) => a,
            _ => return Err("String.value is not an array".into()),
        };
        if value == 0 {
            return Ok(String::new());
        }
        let coder = match self.get(oop, "coder")? {
            JVal::Byte(b) => b,
            _ => 0,
        };
        let n = self.array_len(value)?.max(0) as usize;
        let bytes = self.p.read_bytes(value + self.l.array_data, n.min(1 << 16))?;
        if coder == 0 {
            Ok(bytes.iter().map(|&b| b as char).collect())
        } else {
            let units: Vec<u16> = bytes
                .chunks_exact(2)
                .map(|c| u16::from_le_bytes([c[0], c[1]]))
                .collect();
            Ok(String::from_utf16_lossy(&units))
        }
    }
}

/// UNSIGNED5 decode, transcribed from `unsigned5.hpp`.
/// One excluded byte (X=1) so a 0 byte can terminate a stream; 191 "low" byte
/// values encode directly, the remaining 64 are continuation digits base 64.
fn u5(a: &[u8], pos: &mut usize) -> Result<u32> {
    const X: u32 = 1;
    const L: u32 = 191;
    const LG_H: u32 = 6;
    let b0 = *a.get(*pos).ok_or("fieldinfo stream truncated")? as u32;
    let mut sum = b0.wrapping_sub(X);
    if sum < L {
        *pos += 1;
        return Ok(sum);
    }
    let mut lg_h_i = LG_H;
    for i in 1..5usize {
        let bi = *a.get(*pos + i).ok_or("fieldinfo stream truncated")? as u32;
        sum = sum.wrapping_add((bi - X) << lg_h_i);
        if bi < X + L || i == 4 {
            *pos += i + 1;
            return Ok(sum);
        }
        lg_h_i += LG_H;
    }
    unreachable!()
}