[Unity] Added ability to look for Classes inside the linked list of Transforms

src/game_engine/unity/scene.rs

@@ -2,6 +2,9 @@
 
 // References:
 // https://gist.githubusercontent.com/just-ero/92457b51baf85bd1e5b8c87de8c9835e/raw/8aa3e6b8da01fd03ff2ff0c03cbd018e522ef988/UnityScene.hpp
+// Offsets and logic for the GameObject functions taken from https://github.com/Micrologist/UnityInstanceDumper
+
+use core::{array, mem::MaybeUninit};
 
 use crate::{
     file_format::pe, future::retry, signature::Signature, string::ArrayCString, Address, Address32,
@@ -15,6 +18,7 @@ use crate::{
 /// the traditional class lookup in games with no useful static references.
 pub struct SceneManager {
     is_64_bit: bool,
+    is_il2cpp: bool,
     address: Address,
     offsets: &'static Offsets,
 }
@@ -30,6 +34,7 @@ impl SceneManager {
         let unity_player = process.get_module_range("UnityPlayer.dll").ok()?;
 
         let is_64_bit = pe::MachineType::read(process, unity_player.0)? == pe::MachineType::X86_64;
+        let is_il2cpp = process.get_module_address("GameAssembly.dll").is_ok();
 
         let address = if is_64_bit {
             let addr = SIG_64_BIT.scan_process_range(process, unity_player)? + 7;
@@ -48,6 +53,7 @@ impl SceneManager {
 
         Some(Self {
             is_64_bit,
+            is_il2cpp,
             address,
             offsets,
         })
@@ -73,6 +79,15 @@ impl SceneManager {
         self.read_pointer(process, addr + self.offsets.active_scene)
     }
 
+    /// `DontDestroyOnLoad` is a special Unity scene containing game objects
+    /// that must be preserved when switching between different scenes (eg. a
+    /// `scene1` starting some background music that continues when `scene2`
+    /// loads).
+    fn get_dont_destroy_on_load_scene_address(&self, process: &Process) -> Result<Address, Error> {
+        let addr = self.read_pointer(process, self.address)?;
+        Ok(addr + self.offsets.dont_destroy_on_load_scene)
+    }
+
     /// Returns the current scene index.
     ///
     /// The value returned is a [`i32`] because some games will show `-1` as their
@@ -121,14 +136,299 @@ impl SceneManager {
             })
             .filter(move |p| !fptr.is_null() && p.is_valid(process))
     }
+
+    /// Iterates over all root `Transform`s / [`GameObject`]s declared for the
+    /// specified scene.
+    ///
+    /// Each Unity scene normally has a linked list of `Transform`s (each one is
+    /// a [`GameObject`]). Each one can, recursively, have a child `Transform`
+    /// (and so on), and has a list of `Component`s, which are classes (eg.
+    /// `MonoBehaviour`) containing data we might want to retrieve for the auto
+    /// splitter logic.
+    fn root_game_objects<'a>(
+        &'a self,
+        process: &'a Process,
+        scene_address: Address,
+    ) -> Result<impl DoubleEndedIterator<Item = GameObject> + 'a, Error> {
+        let first_game_object =
+            self.read_pointer(process, scene_address + self.offsets.root_storage_container)?;
+
+        let number_of_root_game_objects = {
+            let mut index: usize = 0;
+            let mut temp_tr = first_game_object;
+
+            while temp_tr != scene_address + self.offsets.root_storage_container {
+                index += 1;
+                temp_tr = self.read_pointer(process, temp_tr)?;
+            }
+
+            index
+        };
+
+        let mut current_game_object = first_game_object;
+
+        Ok((0..number_of_root_game_objects).filter_map(move |n| {
+            let [first, _, third]: [Address; 3] = match self.is_64_bit {
+                true => process
+                    .read::<[Address64; 3]>(current_game_object)
+                    .ok()?
+                    .map(|a| a.into()),
+                false => process
+                    .read::<[Address32; 3]>(current_game_object)
+                    .ok()?
+                    .map(|a| a.into()),
+            };
+
+            let game_object = self
+                .read_pointer(process, third + self.offsets.game_object)
+                .ok()?;
+
+            // Load the next game object before looping, except at the last iteration of the loop
+            if n + 1 != number_of_root_game_objects {
+                current_game_object = first;
+            }
+
+            Some(GameObject {
+                address: game_object,
+            })
+        }))
+    }
+
+    /// Tries to find the specified root [`GameObject`] from the currently
+    /// active Unity scene.
+    pub fn get_root_game_object(&self, process: &Process, name: &str) -> Result<GameObject, Error> {
+        self.root_game_objects(process, self.get_current_scene_address(process)?)?
+            .find(|obj| {
+                obj.get_name::<128>(process, self)
+                    .unwrap_or_default()
+                    .matches(name)
+            })
+            .ok_or(Error {})
+    }
+
+    /// Tries to find the specified root [`GameObject`] from the
+    /// `DontDestroyOnLoad` Unity scene.
+    pub fn get_game_object_from_dont_destroy_on_load(
+        &self,
+        process: &Process,
+        name: &str,
+    ) -> Result<GameObject, Error> {
+        self.root_game_objects(
+            process,
+            self.get_dont_destroy_on_load_scene_address(process)?,
+        )?
+        .find(|obj| {
+            obj.get_name::<128>(process, self)
+                .unwrap_or_default()
+                .matches(name)
+        })
+        .ok_or(Error {})
+    }
+}
+
+/// A `GameObject` is a base class for all entities used in a Unity scene. All
+/// classes of interest useful for an auto splitter can be found starting from
+/// the addresses of the root `GameObject`s linked in each scene.
+pub struct GameObject {
+    address: Address,
+}
+
+impl GameObject {
+    /// Tries to return the name of the current `GameObject`.
+    pub fn get_name<const N: usize>(
+        &self,
+        process: &Process,
+        scene_manager: &SceneManager,
+    ) -> Result<ArrayCString<N>, Error> {
+        let name_ptr = scene_manager.read_pointer(
+            process,
+            self.address + scene_manager.offsets.game_object_name,
+        )?;
+        process.read(name_ptr)
+    }
+
+    /// Iterates over the classes referred to in the current `GameObject`.
+    pub fn classes<'a>(
+        &'a self,
+        process: &'a Process,
+        scene_manager: &'a SceneManager,
+    ) -> Result<impl Iterator<Item = Address> + 'a, Error> {
+        let number_of_components = process
+            .read::<u32>(self.address + scene_manager.offsets.number_of_object_components)?
+            as usize;
+
+        if number_of_components == 0 {
+            return Err(Error {});
+        }
+
+        let main_object = scene_manager
+            .read_pointer(process, self.address + scene_manager.offsets.game_object)?;
+
+        const ARRAY_SIZE: usize = 128;
+
+        let components: [Address; ARRAY_SIZE] = if scene_manager.is_64_bit {
+            let mut buf = [MaybeUninit::<[Address64; 2]>::uninit(); ARRAY_SIZE];
+            let slice =
+                process.read_into_uninit_slice(main_object, &mut buf[..number_of_components])?;
+
+            let mut iter = slice.iter_mut();
+            array::from_fn(|_| {
+                iter.next()
+                    .map(|&mut [_, second]| second.into())
+                    .unwrap_or_default()
+            })
+        } else {
+            let mut buf = [MaybeUninit::<[Address32; 2]>::uninit(); ARRAY_SIZE];
+            let slice =
+                process.read_into_uninit_slice(main_object, &mut buf[..number_of_components])?;
+
+            let mut iter = slice.iter_mut();
+            array::from_fn(|_| {
+                iter.next()
+                    .map(|&mut [_, second]| second.into())
+                    .unwrap_or_default()
+            })
+        };
+
+        Ok((0..number_of_components).filter_map(move |m| {
+            scene_manager
+                .read_pointer(process, components[m] + scene_manager.offsets.klass)
+                .ok()
+        }))
+    }
+
+    /// Tries to find the base address of a class in the current `GameObject`.
+    pub fn get_class(
+        &self,
+        process: &Process,
+        scene_manager: &SceneManager,
+        name: &str,
+    ) -> Result<Address, Error> {
+        self.classes(process, scene_manager)?
+            .find(|&c| {
+                let Ok(vtable) = scene_manager.read_pointer(process, c) else {
+                    return false;
+                };
+
+                let name_ptr = {
+                    match scene_manager.is_il2cpp {
+                        true => {
+                            let Ok(name_ptr) = scene_manager.read_pointer(
+                                process,
+                                vtable + 2 * if scene_manager.is_64_bit { 8 } else { 4 },
+                            ) else {
+                                return false;
+                            };
+
+                            name_ptr
+                        }
+                        false => {
+                            let Ok(vtable) = scene_manager.read_pointer(process, vtable) else {
+                                return false;
+                            };
+
+                            let Ok(name_ptr) = scene_manager
+                                .read_pointer(process, vtable + scene_manager.offsets.klass_name)
+                            else {
+                                return false;
+                            };
+
+                            name_ptr
+                        }
+                    }
+                };
+
+                process
+                    .read::<ArrayCString<128>>(name_ptr)
+                    .is_ok_and(|class_name| class_name.matches(name))
+            })
+            .ok_or(Error {})
+    }
+
+    /// Iterates over children `GameObject`s referred by the current one
+    pub fn children<'a>(
+        &'a self,
+        process: &'a Process,
+        scene_manager: &'a SceneManager,
+    ) -> Result<impl Iterator<Item = Self> + 'a, Error> {
+        let main_object = scene_manager
+            .read_pointer(process, self.address + scene_manager.offsets.game_object)?;
+
+        let transform = scene_manager.read_pointer(
+            process,
+            main_object + if scene_manager.is_64_bit { 8 } else { 4 },
+        )?;
+
+        let child_count =
+            process.read::<u32>(transform + scene_manager.offsets.children_count)? as usize;
+
+        if child_count == 0 {
+            return Err(Error {});
+        }
+
+        let child_pointer = scene_manager
+            .read_pointer(process, transform + scene_manager.offsets.children_pointer)?;
+
+        // Define an empty array and fill it later with the addresses of all child classes found for the current GameObject.
+        // Reading the whole array of pointers is (slightly) faster than reading each address in a loop
+        const ARRAY_SIZE: usize = 128;
+
+        let children: [Address; ARRAY_SIZE] = if scene_manager.is_64_bit {
+            let mut buf = [MaybeUninit::<Address64>::uninit(); ARRAY_SIZE];
+            let slice = process.read_into_uninit_slice(child_pointer, &mut buf[..child_count])?;
+
+            let mut iter = slice.iter_mut();
+            array::from_fn(|_| iter.next().copied().map(Into::into).unwrap_or_default())
+        } else {
+            let mut buf = [MaybeUninit::<Address32>::uninit(); ARRAY_SIZE];
+            let slice = process.read_into_uninit_slice(child_pointer, &mut buf[..child_count])?;
+
+            let mut iter = slice.iter_mut();
+            array::from_fn(|_| iter.next().copied().map(Into::into).unwrap_or_default())
+        };
+
+        Ok((0..child_count).filter_map(move |f| {
+            let game_object = scene_manager
+                .read_pointer(process, children[f] + scene_manager.offsets.game_object)
+                .ok()?;
+
+            Some(Self {
+                address: game_object,
+            })
+        }))
+    }
+
+    /// Tries to find a child `GameObject` from the current one.
+    pub fn get_child(
+        &self,
+        process: &Process,
+        scene_manager: &SceneManager,
+        name: &str,
+    ) -> Result<Self, Error> {
+        self.children(process, scene_manager)?
+            .find(|p| {
+                p.get_name::<128>(process, scene_manager)
+                    .is_ok_and(|obj_name| obj_name.matches(name))
+            })
+            .ok_or(Error {})
+    }
 }
 
 struct Offsets {
     scene_count: u8,
     loaded_scenes: u8,
     active_scene: u8,
+    dont_destroy_on_load_scene: u8,
     asset_path: u8,
     build_index: u8,
+    root_storage_container: u8,
+    game_object: u8,
+    game_object_name: u8,
+    number_of_object_components: u8,
+    klass: u8,
+    klass_name: u8,
+    children_count: u8,
+    children_pointer: u8,
 }
 
 impl Offsets {
@@ -138,15 +438,33 @@ impl Offsets {
                 scene_count: 0x18,
                 loaded_scenes: 0x28,
                 active_scene: 0x48,
+                dont_destroy_on_load_scene: 0x70,
                 asset_path: 0x10,
                 build_index: 0x98,
+                root_storage_container: 0xB0,
+                game_object: 0x30,
+                game_object_name: 0x60,
+                number_of_object_components: 0x40,
+                klass: 0x28,
+                klass_name: 0x48,
+                children_count: 0x80,
+                children_pointer: 0x70,
             },
             false => &Self {
                 scene_count: 0xC,
                 loaded_scenes: 0x18,
                 active_scene: 0x28,
+                dont_destroy_on_load_scene: 0x40,
                 asset_path: 0xC,
                 build_index: 0x70,
+                root_storage_container: 0x88,
+                game_object: 0x1C,
+                game_object_name: 0x3C,
+                number_of_object_components: 0x24,
+                klass: 0x18,
+                klass_name: 0x2C,
+                children_count: 0x58,
+                children_pointer: 0x50,
             },
         }
     }