diff --git a/ykrt/src/mt.rs b/ykrt/src/mt.rs index 9b3189c98..dd48656f6 100644 --- a/ykrt/src/mt.rs +++ b/ykrt/src/mt.rs @@ -17,6 +17,10 @@ use std::{ thread, }; +#[cfg(tracer_swt)] +#[cfg(target_arch = "x86_64")] +use crate::trace::swt::patch::{patch_trace_function, restore_trace_function}; + use parking_lot::{Condvar, Mutex, MutexGuard}; #[cfg(not(all(feature = "yk_testing", not(test))))] use parking_lot_core::SpinWait; @@ -269,6 +273,10 @@ impl MT { } TransitionControlPoint::StartTracing(hl) => { log_jit_state("start-tracing"); + #[cfg(tracer_swt)] + unsafe { + restore_trace_function(); + } let tracer = { let lk = self.tracer.lock(); Arc::clone(&*lk) @@ -285,6 +293,11 @@ impl MT { } } TransitionControlPoint::StopTracing => { + #[cfg(tracer_swt)] + unsafe { + patch_trace_function(); + } + // Assuming no bugs elsewhere, the `unwrap`s cannot fail, because `StartTracing` // will have put a `Some` in the `Rc`. let (hl, thread_tracer, promotions) = @@ -330,6 +343,10 @@ impl MT { Ok(utrace) => { self.stats.timing_state(TimingState::None); log_jit_state("stop-side-tracing"); + #[cfg(tracer_swt)] + unsafe { + patch_trace_function(); + } self.queue_compile_job( (utrace, promotions.into_boxed_slice()), hl, @@ -547,6 +564,10 @@ impl MT { TransitionGuardFailure::NoAction => todo!(), TransitionGuardFailure::StartSideTracing(hl) => { log_jit_state("start-side-tracing"); + #[cfg(tracer_swt)] + unsafe { + restore_trace_function(); + } let tracer = { let lk = self.tracer.lock(); Arc::clone(&*lk) diff --git a/ykrt/src/trace/swt/mod.rs b/ykrt/src/trace/swt/mod.rs index d3287b987..d1ceb6a3f 100644 --- a/ykrt/src/trace/swt/mod.rs +++ b/ykrt/src/trace/swt/mod.rs @@ -15,6 +15,9 @@ use std::{ sync::{Arc, LazyLock}, }; +#[cfg(target_arch = "x86_64")] +pub(crate) mod patch; + #[derive(Debug, Eq, PartialEq, Clone)] struct TracingBBlock { function_index: usize, @@ -105,7 +108,7 @@ impl TraceRecorder for SWTTraceRecorder { Err(TraceRecorderError::TraceTooLong) } else if bbs.is_empty() { // FIXME: who should handle an empty trace? - panic!(); + panic!("swt encountered an empty trace!"); } else { Ok(Box::new(SWTraceIterator::new(bbs))) } diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs new file mode 100644 index 000000000..6468fa7e0 --- /dev/null +++ b/ykrt/src/trace/swt/patch.rs @@ -0,0 +1,143 @@ +//! This module provides functionality for dynamically patching and +//! restoring `yk_trace_basicblock` function at runtime. It includes +//! mechanisms to save the original instructions of a target function, +//! patch the target function with new instructions, and restore the +//! original instructions. +//! +//! This is particularly useful in scenarios such as software tracing, +//! where the tracing function may be dynamically modified to have +//! single return instruction and later restored to their original state +//! for performance gain. +//! +//! The module relies on low-level operations like memory protection +//! changes to manipulate function code safely. It is designed to be +//! used with `tracer_swt` configurations, enabling conditional +//! compilation for tracing functionalities. +//! +//! # Warning +//! +//! This module performs low-level memory operations and modifies the +//! execution flow of functions at runtime. Improper use can lead to +//! undefined behaviour, memory corruption, or crashes. + +use crate::trace::swt::yk_trace_basicblock; +use libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE, _SC_PAGESIZE}; +use std::alloc::Layout; +use std::{ffi::c_void, sync::Once}; + +// This is used to ensure that the original instructions are only saved once. +static ORIGINAL_INSTRUCTIONS_INIT: Once = Once::new(); + +// Original instructions of the function that is patched with `PATCH_X86_INSTRUCTIONS`. +static mut ORIGINAL_INSTRUCTIONS: [u8; 1] = [0; 1]; + +// 0xC3 is a `ret` instruction on x86_64. +static mut PATCH_X86_INSTRUCTIONS: [u8; 1] = [0xC3]; + +/// This function is used to save the original instructions of a function to . +/// +/// # Arguments +/// +/// * `function_ptr` - A usize representing the memory address of the function. +/// * `instructions` - A mutable pointer to a u8 where the original instructions will be saved. +/// * `num_of_instructions` - A usize indicating the number of instructions to save. +/// +unsafe fn save_original_instructions( + function_ptr: usize, + instructions: *mut u8, + num_of_instructions: usize, +) { + let func_ptr: *const () = function_ptr as *const (); + std::ptr::copy_nonoverlapping(func_ptr as *const u8, instructions, num_of_instructions); +} + +/// This function is used to patch a function instructions at runtime. +/// +/// # Arguments +/// +/// * `function_ptr` - A usize representing the memory address of the function to be patched. +/// * `code` - A constant pointer to a u8 vector where the new instructions are located. +/// * `size` - A size_t indicating the number of bytes to copy from `code`. +/// +unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) { + let page_size = sysconf(_SC_PAGESIZE) as usize; + // Align the function_ptr address down to the nearest page boundary. + let page_address = (function_ptr & !(page_size - 1)) as *mut c_void; + // Calculate the offset of function_ptr from the page start. + let start_offset = function_ptr - (page_address as usize); + + // This unwrap should be safe since we are using a page that is + // based on function_ptr with a known location. + let layout = Layout::from_size_align(start_offset, page_size).unwrap(); + + // Set function memory page as writable. + if mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE) != 0 { + panic!("Failed to change memory protection to be writable"); + } + // Copy the new code over + std::ptr::copy_nonoverlapping(code, function_ptr as *mut u8, size); + // Set function memory page as readable + if mprotect(page_address, layout.size(), PROT_READ | PROT_EXEC) != 0 { + panic!("Failed to change memory protection back to executable"); + } +} + +/// This function is used to patch the `yk_trace_basicblock` +/// function with a single `ret` (0xC3) instruction. +pub(crate) unsafe fn patch_trace_function() { + ORIGINAL_INSTRUCTIONS_INIT.call_once(|| { + save_original_instructions( + yk_trace_basicblock as usize, + ORIGINAL_INSTRUCTIONS.as_mut_ptr(), + 1, + ); + }); + #[cfg(target_arch = "x86_64")] + patch_function( + yk_trace_basicblock as usize, + PATCH_X86_INSTRUCTIONS.as_ptr(), + 1, + ); +} + +/// This function is used to restore the original behavior of a +/// previously patched `yk_trace_basicblock` function. +pub(crate) unsafe fn restore_trace_function() { + ORIGINAL_INSTRUCTIONS_INIT.call_once(|| { + save_original_instructions( + yk_trace_basicblock as usize, + ORIGINAL_INSTRUCTIONS.as_mut_ptr(), + 1, + ); + }); + patch_function( + yk_trace_basicblock as usize, + ORIGINAL_INSTRUCTIONS.as_ptr(), + 1, + ); +} + +#[cfg(test)] +mod patch_tests { + use super::*; + + fn test_function() -> i32 { + return 42; + } + + #[test] + fn test_runtime_patch() { + unsafe { + assert_eq!(test_function(), 42); + save_original_instructions( + test_function as usize, + ORIGINAL_INSTRUCTIONS.as_mut_ptr(), + 1, + ); + patch_function(test_function as usize, PATCH_X86_INSTRUCTIONS.as_ptr(), 1); + assert_eq!(test_function(), 0); + patch_function(test_function as usize, ORIGINAL_INSTRUCTIONS.as_ptr(), 1); + assert_eq!(test_function(), 42); + } + } +}