Improved: Memory Efficiency of Assembly Hooks

docs/dev/design/assembly-hooks/overview.md

@@ -138,6 +138,8 @@ The following table below shows common hook lengths, for:
 
 !!! info "Extra: [Thread Safety on x86](../common.md#thread-safety-on-x86)"
 
+### Stub Memory Layout
+
 In order to support thread safety, while retaining maximum runtime performance, the buffers where the 
 original and hook code are contained have a very specific memory layout (shown below)
 
@@ -150,7 +152,7 @@ original and hook code are contained have a very specific memory layout (shown b
 Emplacing the jump to the hook function itself, and patching within the hook function should be atomic
 whenever it is possible on the platform.
 
-### Example
+#### Example
 
 If the *'Original Code'* was:
 
@@ -185,6 +187,28 @@ hook: ; Backup (Hook)
  b back_to_code
 ```
 
+### Heap Layout
+
+Each Assembly Hook contains a pointer to the heap stub (seen above) and a pointer to the heap.
+
+The heap contains all information required to perform operations on the stub.
+
+```text
+- AssemblyHookPackedProps
+ - Enabled Flag
+ - Offset of Hook Function (Also length of HookFunction/OriginalCode block)
+ - Offset of Original Code
+- [Hook Function / Original Code]
+```
+
+The data in the heap contains a short 'AssemblyHookPackedProps' struct, detailing the data that is required
+to make a temporary branch to the stub/hook function. After that is the either the `hook function` bytes or
+the `original code` bytes, depending on the state of the hook.
+
+The hook uses a 'swapping' system, where the `[Hook Function / Original Code]` block in the stub is swapped
+with the `[Hook Function / Original Code]` block in the heap. When one contains the code for `Hook Function`,
+the other contains the code for `Original Code`. This is memory efficient.
+
 ### Switching State
 
 !!! info "When transitioning between Enabled/Disabled state, we place a temporary branch at `entry`, this allows us to manipulate the remaining code safely."
@@ -212,8 +236,9 @@ This means a few functionalities must be supported here:
 
 Assembly hook info is packed by default to save on memory space. By default, the following limits apply:
 
-| Property | 4 Byte Instruction (e.g. ARM) | x86 | Unknown |
-| -------------------- | ----------------------------- | ------ | ------- |
-| Max Branch Length | 4 | 5 | 8 |
-| Max Orig Code Length | 16KiB | 4KiB | 128MiB |
-| Max Hook Code Length | 2MiB | 128KiB | 1GiB |
+| Property | 4 Byte Instruction (e.g. ARM64) | Other (e.g. x86) |
+| -------------------- | ------------------------------- | ---------------- |
+| Max Orig Code Length | 128KiB | 32KiB |
+| Max Hook Code Length | 128KiB | 32KiB |
+
+!!! note "These limits may increase in the future if additional functionality warrants extending metadata length."

projects/reloaded-hooks-portable/src/api/hooks/assembly/assembly_hook.rs

@@ -9,21 +9,34 @@ use crate::{
  settings::assembly_hook_settings::AssemblyHookSettings,
  traits::register_info::RegisterInfo,
  },
- helpers::atomic_write_masked::atomic_write_masked,
+ helpers::{
+ atomic_write_masked::atomic_write_masked, jit_jump_operation::create_jump_operation,
+ },
  internal::assembly_hook::create_assembly_hook,
 };
-
-use core::marker::PhantomData;
+use alloc::vec::Vec;
 use core::ptr::NonNull;
-
-#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-use super::assembly_hook_props_x86::*;
-
-#[cfg(target_arch = "aarch64")]
+use core::{marker::PhantomData, slice::from_raw_parts_mut};
+
+#[cfg(any(
+ target_arch = "aarch64",
+ target_arch = "arm",
+ target_arch = "mips",
+ target_arch = "powerpc",
+ target_arch = "riscv32",
+ target_arch = "riscv64"
+))]
 use super::assembly_hook_props_4byteins::*;
 
-#[cfg(not(any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")))]
-use super::assembly_hook_props_unknown::*;
+#[cfg(not(any(
+ target_arch = "aarch64",
+ target_arch = "arm",
+ target_arch = "mips",
+ target_arch = "powerpc",
+ target_arch = "riscv32",
+ target_arch = "riscv64"
+)))]
+use super::assembly_hook_props_other::*;
 
 /// Represents an assembly hook.
 #[repr(C)] // Not 'packed' because this is not in array and malloc in practice will align this.
@@ -142,19 +155,45 @@ where
  }
 
  /// Writes the hook to memory, either enabling or disabling it based on the provided parameters.
- fn write_hook(&self, branch_opcode: &[u8], code: &[u8], num_bytes: usize) {
- // Write the branch first, as per docs
+ unsafe fn swap_hook(&self, temp_branch_offset: usize) {
+ let props = self.props.as_ref();
+
+ // Backup current code from swap buffer.
+ let swap_buffer_real = props.get_swap_buffer();
+ let swap_buffer_copy = swap_buffer_real.to_vec();
+
+ // Copy current code into swap buffer
+ let buf_buffer_real =
+ from_raw_parts_mut(self.stub_address as *mut u8, props.get_swap_size());
+ swap_buffer_real.copy_from_slice(buf_buffer_real);
+
+ // JIT temp branch to hook/orig code.
+ let mut vec = Vec::<u8>::with_capacity(8);
+ _ = create_jump_operation::<TRegister, TJit, TBufferFactory, TBuffer>(
+ self.stub_address,
+ true,
+ self.stub_address + temp_branch_offset,
+ None,
+ &mut vec,
+ );
+ let branch_opcode = &vec;
+ let branch_bytes = branch_opcode.len();
+
+ // Write the temp branch first, as per docs
  // This also overwrites some extra code afterwards, but that's a-ok for now.
  unsafe {
- atomic_write_masked::<TBuffer>(self.stub_address, branch_opcode, num_bytes);
+ atomic_write_masked::<TBuffer>(self.stub_address, branch_opcode, branch_bytes);
  }
 
  // Now write the remaining code
- TBuffer::overwrite(self.stub_address + num_bytes, &code[num_bytes..]);
+ TBuffer::overwrite(
+ self.stub_address + branch_bytes,
+ &swap_buffer_copy[branch_bytes..],
+ );
 
  // And now re-insert the code we temp overwrote with the branch
  unsafe {
- atomic_write_masked::<TBuffer>(self.stub_address, code, num_bytes);
+ atomic_write_masked::<TBuffer>(self.stub_address, &swap_buffer_copy, branch_bytes);
  }
  }
 
@@ -164,13 +203,13 @@ where
  /// If the hook is disabled, this function will write the hook to memory.
  pub fn enable(&self) {
  unsafe {
- let props = self.props.as_ref();
- let num_bytes = props.get_branch_to_hook_len();
- self.write_hook(
-  props.get_branch_to_hook_slice(),
- props.get_enabled_code(),
-  num_bytes,
- );
+ let props = &mut (*self.props.as_ptr());
+ if props.is_enabled() {
+  return;
+ };
+
+ self.swap_hook(props.get_swap_size());
+ props.set_is_enabled(true);
  }
  }
 
@@ -180,13 +219,13 @@ where
  /// If the hook is enabled, this function will no-op the hook.
  pub fn disable(&self) {
  unsafe {
- let props = self.props.as_ref();
- let num_bytes = props.get_branch_to_orig_len();
- self.write_hook(
-  props.get_branch_to_orig_slice(),
- props.get_disabled_code(),
-  num_bytes,
- );
+ let props = &mut (*self.props.as_ptr());
+ if !props.is_enabled() {
+  return;
+ };
+
+ self.swap_hook(props.get_swap_size() + props.get_hook_fn_size());
+ props.set_is_enabled(false);
  }
  }
 

projects/reloaded-hooks-portable/src/api/hooks/assembly/assembly_hook_props_4byteins.rs

@@ -1,68 +1,48 @@
 use bitfield::bitfield;
 
 bitfield! {
- /// Defines the data layout of the Assembly Hook data.
- /// For architectures which use 4 byte instructions.
+ /// Defines the data layout of the Assembly Hook data for architectures
+ /// with fixed instruction sizes of 4 bytes.
  pub struct AssemblyHookPackedProps(u32);
  impl Debug;
 
  /// True if the hook is enabled, else false.
  pub is_enabled, set_is_enabled: 0;
 
- /// Length of the 'disabled code' array.
- u16, disabled_code_len, set_disabled_code_len_impl: 12, 1; // Max 16KiB.
+ reserved, _: 1;
 
- /// Length of the 'enabled code' array.
- u32, enabled_code_len, set_enabled_code_len_impl: 31, 13; // Max 2MiB.
+ /// Size of the 'swap' space where hook function and original code are swapped out.
+ /// Represented in number of 4-byte instructions.
+ u16, swap_size, set_swap_size_impl: 16, 2; // Max 32Ki instructions, 128KiB.
+
+ /// Size of the 'swap' space where hook function and original code are swapped out.
+ /// Represented in number of 4-byte instructions.
+ u16, hook_fn_size, set_hook_fn_size_impl: 31, 17; // Max 32Ki instructions, 128KiB.
 }
 
 impl AssemblyHookPackedProps {
- /// Gets the length of the 'enabled code' array.
- pub fn get_enabled_code_len(&self) -> usize {
- self.enabled_code_len() as usize * 4
- }
-
- /// Gets the length of the 'disabled code' array.
- pub fn get_disabled_code_len(&self) -> usize {
- self.disabled_code_len() as usize * 4
- }
-
- /// Sets the length of the 'enabled code' array with a minimum value of 4.
- pub fn set_enabled_code_len(&mut self, len: usize) {
- debug_assert!(
- len >= 4 && len % 4 == 0,
- "Length must be a multiple of 4 and at least 4"
- );
- self.set_enabled_code_len_impl((len / 4) as u32);
+ pub fn get_swap_size(&self) -> usize {
+ (self.swap_size() as usize) * 4 // Convert from instructions to bytes
  }
 
- /// Sets the length of the 'disabled code' array with a minimum value of 4.
- pub fn set_disabled_code_len(&mut self, len: usize) {
+ pub fn set_swap_size(&mut self, size: usize) {
  debug_assert!(
- len >= 4 && len % 4 == 0,
- "Length must be a multiple of 4 and at least 4"
+ size % 4 == 0 && size <= 128 * 1024,
+ "Swap size must be a multiple of 4 and at most 128KiB"
  );
- self.set_disabled_code_len_impl((len / 4) as u16);
- }
-
- /// Sets the 'branch to orig' length field based on the provided length.
- pub fn set_branch_to_hook_len(&mut self, _len: usize) {
- // no-op, for API compatibility
- }
-
- /// Gets the length of the 'branch to hook' array. Always 4 for AArch64.
- pub fn get_branch_to_hook_len(&self) -> usize {
- 4
+ self.set_swap_size_impl((size / 4) as u16); // Convert from bytes to instructions
  }
 
- /// Sets the 'branch to orig' length field based on the provided length.
- pub fn set_branch_to_orig_len(&mut self, _len: usize) {
- // no-op, for API compatibility
+ pub fn get_hook_fn_size(&self) -> usize {
+ (self.hook_fn_size() as usize) * 4 // Convert from instructions to bytes
  }
 
- /// Gets the length of the 'branch to orig' array. Always 4 for AArch64.
- pub fn get_branch_to_orig_len(&self) -> usize {
- 4
+ pub fn set_hook_fn_size(&mut self, size: usize) {
+ debug_assert!(
+ size % 4 == 0 && size <= 128 * 1024,
+ "Hook function size must be a multiple of 4 and at most 128KiB"
+ );
+ self.set_hook_fn_size_impl((size / 4) as u16); // Convert from bytes to instructions
  }
 }
 
@@ -71,34 +51,35 @@ mod tests {
  use super::*;
 
  #[test]
- fn test_enabled_and_disabled_code_lengths() {
+ fn test_swap_and_hook_fn_sizes() {
  let mut props = AssemblyHookPackedProps(0);
 
- // Test setting and getting enabled code length
- props.set_enabled_code_len(123 * 4); // Multiples of 4
- assert_eq!(props.get_enabled_code_len(), 123 * 4);
+ // Test setting and getting swap size
+ props.set_swap_size(1024); // 256 instructions
+ assert_eq!(props.get_swap_size(), 1024);
+
+ // Test setting and getting hook function size
+ props.set_hook_fn_size(2048); // 512 instructions
+ assert_eq!(props.get_hook_fn_size(), 2048);
 
- // Test setting and getting disabled code length
- props.set_disabled_code_len(456 * 4); // Multiples of 4
- assert_eq!(props.get_disabled_code_len(), 456 * 4);
+ // Test upper limits
+ props.set_swap_size(127 * 1024); // 32Ki instructions
+ props.set_hook_fn_size(127 * 1024); // 32Ki instructions
+ assert_eq!(props.get_swap_size(), 127 * 1024);
+ assert_eq!(props.get_hook_fn_size(), 127 * 1024);
  }
 
  #[test]
- #[should_panic(expected = "Length must be a multiple of 4 and at least 4")]
- fn test_invalid_code_length() {
+ #[should_panic(expected = "Swap size must be a multiple of 4 and at most 128KiB")]
+ fn test_swap_size_limit() {
  let mut props = AssemblyHookPackedProps(0);
- props.set_enabled_code_len(5); // Should panic, not a multiple of 4
+ props.set_swap_size(129 * 1024); // Should panic
  }
 
  #[test]
- fn test_branch_lengths() {
+ #[should_panic(expected = "Hook function size must be a multiple of 4 and at most 128KiB")]
+ fn test_hook_fn_size_limit() {
  let mut props = AssemblyHookPackedProps(0);
-
- // Setting and getting branch lengths, always 4 for AArch64
- props.set_branch_to_hook_len(4);
- assert_eq!(props.get_branch_to_hook_len(), 4);
-
- props.set_branch_to_orig_len(4);
- assert_eq!(props.get_branch_to_orig_len(), 4);
+ props.set_hook_fn_size(129 * 1024); // Should panic
  }
 }