From 8d706a53c24bfe06a684468ab2799b0fcfb806a0 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 15:53:46 +0000
Subject: [PATCH 01/28] Add runtime patch and restore calls on trace
 transition.

---
 ykcapi/src/lib.rs           |  8 ----
 ykrt/src/lib.rs             |  3 +-
 ykrt/src/mt.rs              | 13 ++++++
 ykrt/src/trace/swt/mod.rs   |  1 +
 ykrt/src/trace/swt/patch.rs | 79 +++++++++++++++++++++++++++++++++++++
 5 files changed, 95 insertions(+), 9 deletions(-)
 create mode 100644 ykrt/src/trace/swt/patch.rs

diff --git a/ykcapi/src/lib.rs b/ykcapi/src/lib.rs
index 68eedfb9c..f573af8b7 100644
--- a/ykcapi/src/lib.rs
+++ b/ykcapi/src/lib.rs
@@ -14,8 +14,6 @@ use std::{
     ptr,
     sync::Arc,
 };
-#[cfg(tracer_swt)]
-use ykrt::trace_basicblock;
 use ykrt::{HotThreshold, Location, MT};
 
 #[no_mangle]
@@ -97,9 +95,3 @@ pub extern "C" fn yk_location_new() -> Location {
 pub extern "C" fn yk_location_drop(loc: Location) {
     drop(loc)
 }
-
-#[cfg(tracer_swt)]
-#[no_mangle]
-pub extern "C" fn yk_trace_basicblock(function_index: usize, block_index: usize) {
-    trace_basicblock(function_index, block_index);
-}
diff --git a/ykrt/src/lib.rs b/ykrt/src/lib.rs
index 15fda86a5..5e5a3ae6c 100644
--- a/ykrt/src/lib.rs
+++ b/ykrt/src/lib.rs
@@ -35,7 +35,8 @@ pub fn print_jit_state(state: &str) {
 }
 
 #[cfg(tracer_swt)]
-pub fn trace_basicblock(function_index: usize, block_index: usize) {
+#[no_mangle]
+pub extern "C" fn yk_trace_basicblock(function_index: usize, block_index: usize) {
     if mt::is_tracing() {
         trace::trace_basicblock(function_index, block_index)
     }
diff --git a/ykrt/src/mt.rs b/ykrt/src/mt.rs
index 8c049017e..48a3a09d7 100644
--- a/ykrt/src/mt.rs
+++ b/ykrt/src/mt.rs
@@ -17,6 +17,9 @@ use std::{
     thread,
 };
 
+#[cfg(tracer_swt)]
+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;
@@ -263,6 +266,11 @@ impl MT {
             TransitionControlPoint::StartTracing(hl) => {
                 #[cfg(feature = "yk_jitstate_debug")]
                 print_jit_state("start-tracing");
+                #[cfg(tracer_swt)]
+                unsafe {
+                    restore_trace_function();
+                }
+
                 let tracer = {
                     let lk = self.tracer.lock();
                     Arc::clone(&*lk)
@@ -279,6 +287,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) = THREAD_MTTHREAD.with(|mtt| {
diff --git a/ykrt/src/trace/swt/mod.rs b/ykrt/src/trace/swt/mod.rs
index d55a82ceb..9cd82f84c 100644
--- a/ykrt/src/trace/swt/mod.rs
+++ b/ykrt/src/trace/swt/mod.rs
@@ -6,6 +6,7 @@ use std::sync::Once;
 use std::{cell::RefCell, collections::HashMap, error::Error, ffi::CString, sync::Arc};
 
 mod iterator;
+pub(crate) mod patch;
 use iterator::SWTraceIterator;
 
 static FUNC_NAMES_INIT: Once = Once::new();
diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
new file mode 100644
index 000000000..8ca3ec54e
--- /dev/null
+++ b/ykrt/src/trace/swt/patch.rs
@@ -0,0 +1,79 @@
+use libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE};
+use std::mem;
+use std::{ffi::c_void, sync::Once};
+
+#[cfg(tracer_swt)]
+use crate::yk_trace_basicblock;
+
+#[cfg(tracer_swt)]
+static ORIGINAL_INSTRUCTIONS_INIT: Once = Once::new();
+#[cfg(tracer_swt)]
+static mut ORIGINAL_INSTRUCTIONS: [u8; 1] = [0; 1];
+#[cfg(tracer_swt)]
+static mut PATCH_INSTRUCTIONS: [u8; 1] = [0xC3];
+
+#[cfg(tracer_swt)]
+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);
+}
+
+#[cfg(tracer_swt)]
+unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
+    let page_size = sysconf(libc::_SC_PAGESIZE) as usize;
+
+    let func_address = ((function_ptr as usize) & !(page_size - 1)) as *mut c_void;
+    let page_size_aligned = (((function_ptr as usize) + mem::size_of_val(&function_ptr))
+        - (func_address as usize)) as usize;
+
+    // Set function memory region to be writable
+    let result = mprotect(
+        func_address,
+        page_size_aligned,
+        PROT_READ | PROT_WRITE | PROT_EXEC,
+    );
+    if result != 0 {
+        panic!("Failed to change memory protection to be writable");
+    }
+
+    // Set function memory region back to be non-writable
+    std::ptr::copy_nonoverlapping(code, function_ptr as *mut u8, size);
+
+    let result = mprotect(func_address, page_size_aligned, PROT_READ | PROT_EXEC);
+    if result != 0 {
+        panic!("Failed to change memory protection to not writable");
+    }
+}
+
+#[cfg(tracer_swt)]
+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,
+        );
+    });
+
+    patch_function(yk_trace_basicblock as usize, PATCH_INSTRUCTIONS.as_ptr(), 1);
+}
+
+#[cfg(tracer_swt)]
+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,
+    );
+}

From e1fd7741507062b1c9c20fab990b7a045f0a169e Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:08:51 +0000
Subject: [PATCH 02/28] Add docstring in patch module

---
 ykrt/src/trace/swt/patch.rs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 8ca3ec54e..701400aaf 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -12,6 +12,14 @@ static mut ORIGINAL_INSTRUCTIONS: [u8; 1] = [0; 1];
 #[cfg(tracer_swt)]
 static mut PATCH_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.
+///
 #[cfg(tracer_swt)]
 unsafe fn save_original_instructions(
     function_ptr: usize,
@@ -22,6 +30,14 @@ unsafe fn save_original_instructions(
     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`.
+///
 #[cfg(tracer_swt)]
 unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let page_size = sysconf(libc::_SC_PAGESIZE) as usize;

From ee563103037c4d013b3378455b2efe9f1ae73cb0 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:09:18 +0000
Subject: [PATCH 03/28] Add message to panic on empty trace in swt.

---
 ykrt/src/trace/swt/mod.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ykrt/src/trace/swt/mod.rs b/ykrt/src/trace/swt/mod.rs
index 9cd82f84c..ad05c8e23 100644
--- a/ykrt/src/trace/swt/mod.rs
+++ b/ykrt/src/trace/swt/mod.rs
@@ -112,7 +112,7 @@ impl TraceRecorder for SWTTraceRecorder {
             Err(TraceRecorderError::TraceTooLong)
         } else if aot_blocks.is_empty() {
             // FIXME: who should handle an empty trace?
-            panic!();
+            panic!("swt encountered an empty trace!");
         } else {
             Ok(Box::new(SWTraceIterator::new(aot_blocks)))
         }

From 461448f271e26bf4a1c9af19143efd6233d3f27c Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:12:41 +0000
Subject: [PATCH 04/28] Add patch calls to side-traces.

---
 ykrt/src/mt.rs | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/ykrt/src/mt.rs b/ykrt/src/mt.rs
index 48a3a09d7..a71f03888 100644
--- a/ykrt/src/mt.rs
+++ b/ykrt/src/mt.rs
@@ -338,6 +338,10 @@ impl MT {
                         self.stats.timing_state(TimingState::None);
                         #[cfg(feature = "yk_jitstate_debug")]
                         print_jit_state("stop-side-tracing");
+                        #[cfg(tracer_swt)]
+                        unsafe {
+                            patch_trace_function();
+                        }
                         self.queue_compile_job(
                             (utrace, promotions.into_boxed_slice()),
                             hl,
@@ -557,6 +561,10 @@ impl MT {
             TransitionGuardFailure::StartSideTracing(hl) => {
                 #[cfg(feature = "yk_jitstate_debug")]
                 print_jit_state("start-side-tracing");
+                #[cfg(tracer_swt)]
+                unsafe {
+                    restore_trace_function();
+                }
                 let tracer = {
                     let lk = self.tracer.lock();
                     Arc::clone(&*lk)

From e23727872eee175a55b676b6ff2fb22136ee2886 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:21:42 +0000
Subject: [PATCH 05/28] Add panic if runtime patch is called on non-x86
 platform.

---
 ykrt/src/trace/swt/patch.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 701400aaf..23409f3c0 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -65,8 +65,17 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     }
 }
 
+fn is_x86() -> bool {
+    std::env::consts::ARCH == "x86" || std::env::consts::ARCH == "x86_64"
+}
+
+/// This function is used to patch the `yk_trace_basicblock`
+/// function with a single `ret` (0xC3) instruction.
 #[cfg(tracer_swt)]
 pub(crate) unsafe fn patch_trace_function() {
+    if !is_x86() {
+        panic!("Only x86_64 architecture is supported for runtime patching!");
+    }
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
             yk_trace_basicblock as usize,
@@ -78,8 +87,13 @@ pub(crate) unsafe fn patch_trace_function() {
     patch_function(yk_trace_basicblock as usize, PATCH_INSTRUCTIONS.as_ptr(), 1);
 }
 
+/// This function is used to restore the original behavior of a
+/// previously patched `yk_trace_basicblock` function.
 #[cfg(tracer_swt)]
 pub(crate) unsafe fn restore_trace_function() {
+    if !is_x86() {
+        panic!("Only x86_64 architecture is supported for runtime patching!");
+    }
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
             yk_trace_basicblock as usize,

From 0699dbf93c3a8981e2a70711de09a313b93fb38c Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:27:30 +0000
Subject: [PATCH 06/28] Add test to swt patch.

---
 ykrt/src/trace/swt/patch.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 23409f3c0..ad65a328c 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -107,3 +107,21 @@ pub(crate) unsafe fn restore_trace_function() {
         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);
+            patch_function(test_function as usize, PATCH_INSTRUCTIONS.as_ptr(), 1);
+            assert_eq!(test_function(), 0);
+        }
+    }
+}

From 0db3e4f61ef484dfa3f66e20bf49daa1bcee4935 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:30:52 +0000
Subject: [PATCH 07/28] Add restore assertion to patch test.

---
 ykrt/src/trace/swt/patch.rs | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index ad65a328c..20f358a3a 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -120,8 +120,19 @@ mod patch_tests {
     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_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);
         }
     }
 }

From 6994b593d54d9c5171c8f609745b0b510600a05d Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Wed, 27 Mar 2024 16:31:18 +0000
Subject: [PATCH 08/28] Add assert to function restore.

---
 ykrt/src/trace/swt/patch.rs | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 20f358a3a..f80565e38 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -127,11 +127,7 @@ mod patch_tests {
             );
             patch_function(test_function as usize, PATCH_INSTRUCTIONS.as_ptr(), 1);
             assert_eq!(test_function(), 0);
-            patch_function(
-                test_function as usize,
-                ORIGINAL_INSTRUCTIONS.as_ptr(),
-                1,
-            );
+            patch_function(test_function as usize, ORIGINAL_INSTRUCTIONS.as_ptr(), 1);
             assert_eq!(test_function(), 42);
         }
     }

From 0e88a29e6ccb330d899d869631035d66300a8bd4 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 12:37:52 +0000
Subject: [PATCH 09/28] Add x86 compile-time guard

---
 ykrt/src/trace/swt/patch.rs | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index f80565e38..bd8b9fe6c 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -65,17 +65,11 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     }
 }
 
-fn is_x86() -> bool {
-    std::env::consts::ARCH == "x86" || std::env::consts::ARCH == "x86_64"
-}
-
 /// This function is used to patch the `yk_trace_basicblock`
 /// function with a single `ret` (0xC3) instruction.
 #[cfg(tracer_swt)]
+#[cfg(target_arch = "x86_64")]
 pub(crate) unsafe fn patch_trace_function() {
-    if !is_x86() {
-        panic!("Only x86_64 architecture is supported for runtime patching!");
-    }
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
             yk_trace_basicblock as usize,
@@ -90,10 +84,8 @@ pub(crate) unsafe fn patch_trace_function() {
 /// This function is used to restore the original behavior of a
 /// previously patched `yk_trace_basicblock` function.
 #[cfg(tracer_swt)]
+#[cfg(target_arch = "x86_64")]
 pub(crate) unsafe fn restore_trace_function() {
-    if !is_x86() {
-        panic!("Only x86_64 architecture is supported for runtime patching!");
-    }
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
             yk_trace_basicblock as usize,

From 86e404b3b7c0c9f633b9c49f748bc65811d2186e Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 12:38:33 +0000
Subject: [PATCH 10/28] Add compile-time guard.

---
 ykrt/src/mt.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ykrt/src/mt.rs b/ykrt/src/mt.rs
index a71f03888..c36669c78 100644
--- a/ykrt/src/mt.rs
+++ b/ykrt/src/mt.rs
@@ -18,6 +18,7 @@ use std::{
 };
 
 #[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};
@@ -267,6 +268,7 @@ impl MT {
                 #[cfg(feature = "yk_jitstate_debug")]
                 print_jit_state("start-tracing");
                 #[cfg(tracer_swt)]
+                #[cfg(target_arch = "x86_64")]
                 unsafe {
                     restore_trace_function();
                 }
@@ -562,6 +564,7 @@ impl MT {
                 #[cfg(feature = "yk_jitstate_debug")]
                 print_jit_state("start-side-tracing");
                 #[cfg(tracer_swt)]
+                #[cfg(target_arch = "x86_64")]
                 unsafe {
                     restore_trace_function();
                 }

From 3725c7189fb0159b52a412da595f760370cf2ff4 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 12:41:55 +0000
Subject: [PATCH 11/28] Remove PROT_EXEC on mprotect.

---
 ykrt/src/trace/swt/patch.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index bd8b9fe6c..d953fa809 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -1,4 +1,4 @@
-use libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE};
+use libc::{mprotect, size_t, sysconf, PROT_READ, PROT_WRITE};
 use std::mem;
 use std::{ffi::c_void, sync::Once};
 
@@ -50,7 +50,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let result = mprotect(
         func_address,
         page_size_aligned,
-        PROT_READ | PROT_WRITE | PROT_EXEC,
+        PROT_READ | PROT_WRITE,
     );
     if result != 0 {
         panic!("Failed to change memory protection to be writable");

From 799644e64d38f56c9447bcac350e86ce0b4e6182 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 12:43:28 +0000
Subject: [PATCH 12/28] Revert "Remove PROT_EXEC on mprotect."

This reverts commit 3725c7189fb0159b52a412da595f760370cf2ff4.
---
 ykrt/src/trace/swt/patch.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index d953fa809..bd8b9fe6c 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -1,4 +1,4 @@
-use libc::{mprotect, size_t, sysconf, PROT_READ, PROT_WRITE};
+use libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE};
 use std::mem;
 use std::{ffi::c_void, sync::Once};
 
@@ -50,7 +50,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let result = mprotect(
         func_address,
         page_size_aligned,
-        PROT_READ | PROT_WRITE,
+        PROT_READ | PROT_WRITE | PROT_EXEC,
     );
     if result != 0 {
         panic!("Failed to change memory protection to be writable");

From d572cabf3deeef89567a60ca671ac93d08d6743f Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 12:53:02 +0000
Subject: [PATCH 13/28] Add documentation.

---
 ykrt/src/trace/swt/patch.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index bd8b9fe6c..9fa929a60 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -6,10 +6,13 @@ use std::{ffi::c_void, sync::Once};
 use crate::yk_trace_basicblock;
 
 #[cfg(tracer_swt)]
+// This is used to ensure that the original instructions are only saved once.
 static ORIGINAL_INSTRUCTIONS_INIT: Once = Once::new();
 #[cfg(tracer_swt)]
+// Original instructions of the function that is patched with `PATCH_INSTRUCTIONS`.
 static mut ORIGINAL_INSTRUCTIONS: [u8; 1] = [0; 1];
 #[cfg(tracer_swt)]
+// 0xC3 is a `ret` instruction on x86_64.
 static mut PATCH_INSTRUCTIONS: [u8; 1] = [0xC3];
 
 /// This function is used to save the original instructions of a function to .

From 0bd8f3617e95057590fafe45a72eecf6cd96a922 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 13:08:13 +0000
Subject: [PATCH 14/28] Add module docstring.

---
 ykrt/src/trace/swt/patch.rs | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 9fa929a60..9f2e82e2a 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -1,3 +1,25 @@
+//! 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 libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE};
 use std::mem;
 use std::{ffi::c_void, sync::Once};

From b2aee3dd12b8144c7e0f1cf3ab7c8be19029950b Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Fri, 29 Mar 2024 17:53:47 +0000
Subject: [PATCH 15/28] Encapsulate x86 instructions.

---
 ykrt/src/mt.rs              |  2 --
 ykrt/src/trace/swt/patch.rs | 16 +++++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/ykrt/src/mt.rs b/ykrt/src/mt.rs
index e3b890339..adb329c05 100644
--- a/ykrt/src/mt.rs
+++ b/ykrt/src/mt.rs
@@ -261,7 +261,6 @@ impl MT {
                 #[cfg(feature = "yk_jitstate_debug")]
                 print_jit_state("start-tracing");
                 #[cfg(tracer_swt)]
-                #[cfg(target_arch = "x86_64")]
                 unsafe {
                     restore_trace_function();
                 }
@@ -559,7 +558,6 @@ impl MT {
                 #[cfg(feature = "yk_jitstate_debug")]
                 print_jit_state("start-side-tracing");
                 #[cfg(tracer_swt)]
-                #[cfg(target_arch = "x86_64")]
                 unsafe {
                     restore_trace_function();
                 }
diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index ad5457004..5b857a9eb 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -31,11 +31,11 @@ use crate::trace::swt::yk_trace_basicblock;
 // This is used to ensure that the original instructions are only saved once.
 static ORIGINAL_INSTRUCTIONS_INIT: Once = Once::new();
 #[cfg(tracer_swt)]
-// Original instructions of the function that is patched with `PATCH_INSTRUCTIONS`.
+// Original instructions of the function that is patched with `PATCH_X86_INSTRUCTIONS`.
 static mut ORIGINAL_INSTRUCTIONS: [u8; 1] = [0; 1];
 #[cfg(tracer_swt)]
 // 0xC3 is a `ret` instruction on x86_64.
-static mut PATCH_INSTRUCTIONS: [u8; 1] = [0xC3];
+static mut PATCH_X86_INSTRUCTIONS: [u8; 1] = [0xC3];
 
 /// This function is used to save the original instructions of a function to .
 ///
@@ -93,7 +93,6 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
 /// This function is used to patch the `yk_trace_basicblock`
 /// function with a single `ret` (0xC3) instruction.
 #[cfg(tracer_swt)]
-#[cfg(target_arch = "x86_64")]
 pub(crate) unsafe fn patch_trace_function() {
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
@@ -102,14 +101,17 @@ pub(crate) unsafe fn patch_trace_function() {
             1,
         );
     });
-
-    patch_function(yk_trace_basicblock as usize, PATCH_INSTRUCTIONS.as_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.
 #[cfg(tracer_swt)]
-#[cfg(target_arch = "x86_64")]
 pub(crate) unsafe fn restore_trace_function() {
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
@@ -142,7 +144,7 @@ mod patch_tests {
                 ORIGINAL_INSTRUCTIONS.as_mut_ptr(),
                 1,
             );
-            patch_function(test_function as usize, PATCH_INSTRUCTIONS.as_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);

From 40c0696e76ca685f7334716461d13816e852ac3e Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Tue, 9 Apr 2024 16:32:01 +0100
Subject: [PATCH 16/28] Remove PROT_EXEC from mprotect.

---
 ykrt/src/trace/swt/patch.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 5b857a9eb..079e607b3 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -75,7 +75,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let result = mprotect(
         func_address,
         page_size_aligned,
-        PROT_READ | PROT_WRITE | PROT_EXEC,
+        PROT_READ | PROT_WRITE,
     );
     if result != 0 {
         panic!("Failed to change memory protection to be writable");

From 44011539f55d86055d5f1a9d69452069d90397d4 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Tue, 9 Apr 2024 16:35:15 +0100
Subject: [PATCH 17/28] Format.

---
 ykrt/src/trace/swt/patch.rs | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 079e607b3..67abd5b03 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -72,11 +72,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
         - (func_address as usize)) as usize;
 
     // Set function memory region to be writable
-    let result = mprotect(
-        func_address,
-        page_size_aligned,
-        PROT_READ | PROT_WRITE,
-    );
+    let result = mprotect(func_address, page_size_aligned, PROT_READ | PROT_WRITE);
     if result != 0 {
         panic!("Failed to change memory protection to be writable");
     }

From 5ead8930eb6a4375069233f5aec35775b8c740f2 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 11 Apr 2024 18:55:59 +0100
Subject: [PATCH 18/28] Remove not needed cfgs.

---
 ykrt/src/trace/swt/patch.rs | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 67abd5b03..ebcf84832 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -24,16 +24,14 @@ use libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE};
 use std::mem;
 use std::{ffi::c_void, sync::Once};
 
-#[cfg(tracer_swt)]
 use crate::trace::swt::yk_trace_basicblock;
 
-#[cfg(tracer_swt)]
 // This is used to ensure that the original instructions are only saved once.
 static ORIGINAL_INSTRUCTIONS_INIT: Once = Once::new();
-#[cfg(tracer_swt)]
+
 // Original instructions of the function that is patched with `PATCH_X86_INSTRUCTIONS`.
 static mut ORIGINAL_INSTRUCTIONS: [u8; 1] = [0; 1];
-#[cfg(tracer_swt)]
+
 // 0xC3 is a `ret` instruction on x86_64.
 static mut PATCH_X86_INSTRUCTIONS: [u8; 1] = [0xC3];
 
@@ -45,7 +43,6 @@ static mut PATCH_X86_INSTRUCTIONS: [u8; 1] = [0xC3];
 /// * `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.
 ///
-#[cfg(tracer_swt)]
 unsafe fn save_original_instructions(
     function_ptr: usize,
     instructions: *mut u8,
@@ -63,7 +60,6 @@ unsafe fn save_original_instructions(
 /// * `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`.
 ///
-#[cfg(tracer_swt)]
 unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let page_size = sysconf(libc::_SC_PAGESIZE) as usize;
 
@@ -88,7 +84,6 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
 
 /// This function is used to patch the `yk_trace_basicblock`
 /// function with a single `ret` (0xC3) instruction.
-#[cfg(tracer_swt)]
 pub(crate) unsafe fn patch_trace_function() {
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(
@@ -107,7 +102,6 @@ pub(crate) unsafe fn patch_trace_function() {
 
 /// This function is used to restore the original behavior of a
 /// previously patched `yk_trace_basicblock` function.
-#[cfg(tracer_swt)]
 pub(crate) unsafe fn restore_trace_function() {
     ORIGINAL_INSTRUCTIONS_INIT.call_once(|| {
         save_original_instructions(

From b6fcbd48083439b6511dc549131c899b8b28486a Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 11 Apr 2024 18:56:53 +0100
Subject: [PATCH 19/28] Add x86_64 guard.

---
 ykrt/src/trace/swt/mod.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ykrt/src/trace/swt/mod.rs b/ykrt/src/trace/swt/mod.rs
index 2e51b740d..c2d9ad80a 100644
--- a/ykrt/src/trace/swt/mod.rs
+++ b/ykrt/src/trace/swt/mod.rs
@@ -15,6 +15,7 @@ use std::{
     sync::{Arc, LazyLock},
 };
 
+#[cfg(target_arch = "x86_64")]
 pub(crate) mod patch;
 
 #[derive(Debug, Eq, PartialEq, Clone)]

From 6f4bdc3a3a5d4ac15a130d7e3a877a8f9af062c3 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sat, 13 Apr 2024 11:45:35 +0100
Subject: [PATCH 20/28] Set mprotect as read+write on patch.

---
 ykrt/src/trace/swt/patch.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index e6d5be2fd..ebcf84832 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -68,7 +68,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
         - (func_address as usize)) as usize;
 
     // Set function memory region to be writable
-    let result = mprotect(func_address, page_size_aligned, PROT_READ | PROT_WRITE | PROT_EXEC);
+    let result = mprotect(func_address, page_size_aligned, PROT_READ | PROT_WRITE);
     if result != 0 {
         panic!("Failed to change memory protection to be writable");
     }

From 8e00241271862bd3afdccf0f4c308217938f5098 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sun, 14 Apr 2024 16:13:00 +0100
Subject: [PATCH 21/28] Use Layout::from_size_align for aligned page
 calculation.

---
 ykrt/src/trace/swt/patch.rs | 51 +++++++++++++++++++++----------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index ebcf84832..732a04ea7 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -20,11 +20,10 @@
 //! execution flow of functions at runtime. Improper use can lead to
 //! undefined behaviour, memory corruption, or crashes.
 
-use libc::{mprotect, size_t, sysconf, PROT_EXEC, PROT_READ, PROT_WRITE};
-use std::mem;
-use std::{ffi::c_void, sync::Once};
-
 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();
@@ -61,24 +60,32 @@ unsafe fn save_original_instructions(
 /// * `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(libc::_SC_PAGESIZE) as usize;
-
-    let func_address = ((function_ptr as usize) & !(page_size - 1)) as *mut c_void;
-    let page_size_aligned = (((function_ptr as usize) + mem::size_of_val(&function_ptr))
-        - (func_address as usize)) as usize;
-
-    // Set function memory region to be writable
-    let result = mprotect(func_address, page_size_aligned, PROT_READ | PROT_WRITE);
-    if result != 0 {
-        panic!("Failed to change memory protection to be writable");
-    }
-
-    // Set function memory region back to be non-writable
-    std::ptr::copy_nonoverlapping(code, function_ptr as *mut u8, size);
-
-    let result = mprotect(func_address, page_size_aligned, PROT_READ | PROT_EXEC);
-    if result != 0 {
-        panic!("Failed to change memory protection to not writable");
+    let page_size = sysconf(_SC_PAGESIZE) as usize;
+
+    let page_address = (function_ptr & !(page_size - 1)) as *mut c_void;
+    let start_offset = function_ptr - (page_address as usize) + size;
+
+    match Layout::from_size_align(start_offset, page_size) {
+        Ok(layout) => {
+            // Set function memory page as writable
+            let result = mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE);
+            if result != 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
+            let result = mprotect(page_address, layout.size(), PROT_READ | PROT_EXEC);
+            if result != 0 {
+                panic!("Failed to change memory protection back to executable");
+            }
+        }
+        Err(e) => {
+            panic!(
+                "Failed to create layout for fuction instuction patch: {:?}",
+                e
+            );
+        }
     }
 }
 

From ada68bbcfbc3f25fb4fd71c9444a497a9b62c240 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 18 Apr 2024 18:54:30 +0100
Subject: [PATCH 22/28] Change match to unwrap().

---
 ykrt/src/trace/swt/patch.rs | 37 ++++++++++++++++---------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 732a04ea7..7fd03f47d 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -65,27 +65,22 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let page_address = (function_ptr & !(page_size - 1)) as *mut c_void;
     let start_offset = function_ptr - (page_address as usize) + size;
 
-    match Layout::from_size_align(start_offset, page_size) {
-        Ok(layout) => {
-            // Set function memory page as writable
-            let result = mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE);
-            if result != 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
-            let result = mprotect(page_address, layout.size(), PROT_READ | PROT_EXEC);
-            if result != 0 {
-                panic!("Failed to change memory protection back to executable");
-            }
-        }
-        Err(e) => {
-            panic!(
-                "Failed to create layout for fuction instuction patch: {:?}",
-                e
-            );
-        }
+    // 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)
+        .expect("Failed to create layout for function memory page");
+
+    // Set function memory page as writable
+    let result = mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE);
+    if result != 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
+    let result = mprotect(page_address, layout.size(), PROT_READ | PROT_EXEC);
+    if result != 0 {
+        panic!("Failed to change memory protection back to executable");
     }
 }
 

From 3e7d023a35d6c0b8e5c9954ba4f446dea07a4857 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 18 Apr 2024 19:02:57 +0100
Subject: [PATCH 23/28] Recover and panic on mprotect.

---
 ykrt/src/trace/swt/patch.rs | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 7fd03f47d..2119f11bd 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -70,16 +70,13 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let layout = Layout::from_size_align(start_offset, page_size)
         .expect("Failed to create layout for function memory page");
 
-    // Set function memory page as writable
-    let result = mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE);
-    if result != 0 {
-        panic!("Failed to change memory protection to be writable");
-    }
+    // Set function memory page as writable.
+    // Ignoring mprotect call failure.
+    mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE);
     // Copy the new code over
     std::ptr::copy_nonoverlapping(code, function_ptr as *mut u8, size);
     // Set function memory page as readable
-    let result = mprotect(page_address, layout.size(), PROT_READ | PROT_EXEC);
-    if result != 0 {
+    if mprotect(page_address, layout.size(), PROT_READ | PROT_EXEC) != 0 {
         panic!("Failed to change memory protection back to executable");
     }
 }

From 18b7fee60b1d193c14539a03002937c7788e10f6 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 18 Apr 2024 19:12:39 +0100
Subject: [PATCH 24/28] Remove +size from start_offset calculation.

---
 ykrt/src/trace/swt/patch.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 2119f11bd..68e6eb5c3 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -63,7 +63,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let page_size = sysconf(_SC_PAGESIZE) as usize;
 
     let page_address = (function_ptr & !(page_size - 1)) as *mut c_void;
-    let start_offset = function_ptr - (page_address as usize) + size;
+    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.

From e1be48cfe17c1d07dbe077472e51fcba2056b220 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Thu, 18 Apr 2024 19:16:20 +0100
Subject: [PATCH 25/28] Add comments.

---
 ykrt/src/trace/swt/patch.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 68e6eb5c3..450d34fbc 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -61,8 +61,9 @@ unsafe fn save_original_instructions(
 ///
 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

From 0428dc1696e940437ea26afd96954097e5316878 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sat, 20 Apr 2024 13:23:31 +0100
Subject: [PATCH 26/28] Use unwrap over expect.

---
 ykrt/src/trace/swt/patch.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 450d34fbc..69d83ee6e 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -68,8 +68,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
 
     // 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)
-        .expect("Failed to create layout for function memory page");
+    let layout = Layout::from_size_align(start_offset, page_size).unwrap();
 
     // Set function memory page as writable.
     // Ignoring mprotect call failure.

From 6be7bef8f5968aec6e44f792422d446f306c4c7a Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Sat, 20 Apr 2024 13:31:13 +0100
Subject: [PATCH 27/28] Early return if mprotect failes to set memmory page as
 writable.

---
 ykrt/src/trace/swt/patch.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index 69d83ee6e..ac63a8616 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -71,8 +71,9 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
     let layout = Layout::from_size_align(start_offset, page_size).unwrap();
 
     // Set function memory page as writable.
-    // Ignoring mprotect call failure.
-    mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE);
+    if mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE) != 0 {
+        return;
+    }
     // Copy the new code over
     std::ptr::copy_nonoverlapping(code, function_ptr as *mut u8, size);
     // Set function memory page as readable

From 6f0abb36df22b19df4880e289a96538387e9e392 Mon Sep 17 00:00:00 2001
From: Pavel Durov <pavel@p3ld3v.dev>
Date: Tue, 23 Apr 2024 20:40:23 +0100
Subject: [PATCH 28/28] Add panic on mprotect failure.

---
 ykrt/src/trace/swt/patch.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ykrt/src/trace/swt/patch.rs b/ykrt/src/trace/swt/patch.rs
index ac63a8616..6468fa7e0 100644
--- a/ykrt/src/trace/swt/patch.rs
+++ b/ykrt/src/trace/swt/patch.rs
@@ -72,7 +72,7 @@ unsafe fn patch_function(function_ptr: usize, code: *const u8, size: size_t) {
 
     // Set function memory page as writable.
     if mprotect(page_address, layout.size(), PROT_READ | PROT_WRITE) != 0 {
-        return;
+        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);