From 022a7f9d5dca7b14d9b30f826776dce64a0a0ae7 Mon Sep 17 00:00:00 2001
From: Evian-Zhang <evianzhang1999@163.com>
Date: Fri, 9 Aug 2024 16:29:45 +0800
Subject: [PATCH] Use mmap and mremap to allow multi-thread static key
 modifying. Add clear cache for linux OS

---
 Cargo.toml      |  1 +
 src/os/linux.rs | 48 +++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index c9c9e97..e41e888 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,7 @@ maintenance = { status = "actively-developed" }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 libc = { version = "0.2", default-features = false }
+clear-cache = "0.1"
 
 [target.'cfg(target_os = "macos")'.dependencies]
 mach2 = "0.4"
diff --git a/src/os/linux.rs b/src/os/linux.rs
index 3dbe59c..a00cab9 100644
--- a/src/os/linux.rs
+++ b/src/os/linux.rs
@@ -40,26 +40,56 @@ impl CodeManipulator for ArchCodeManipulator {
         } else {
             page_size
         };
-        let res = unsafe {
-            libc::mprotect(
-                aligned_addr,
+
+        // Create a temp mmap, which will store updated content of corresponding pages
+        let mmaped_addr = unsafe {
+            libc::mmap(
+                core::ptr::null_mut(),
                 aligned_length,
-                libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
+                libc::PROT_READ | libc::PROT_WRITE,
+                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
+                -1,
+                0,
             )
         };
-        if res != 0 {
-            panic!("Unable to make code region writable");
+        if mmaped_addr == libc::MAP_FAILED {
+            panic!("Failed to create temp mappings");
+        }
+        unsafe {
+            let addr_in_mmap = mmaped_addr.offset(addr.offset_from(aligned_addr));
+            core::ptr::copy_nonoverlapping(aligned_addr, mmaped_addr, aligned_length);
+            core::ptr::copy_nonoverlapping(data.as_ptr(), addr_in_mmap.cast(), L);
         }
-        core::ptr::copy_nonoverlapping(data.as_ptr(), addr.cast(), L);
         let res = unsafe {
             libc::mprotect(
-                aligned_addr,
+                mmaped_addr,
                 aligned_length,
                 libc::PROT_READ | libc::PROT_EXEC,
             )
         };
         if res != 0 {
-            panic!("Unable to restore code region to non-writable");
+            panic!("Unable to make mmaped mapping executable.");
+        }
+        // Remap the created temp mmaping to replace old mapping
+        let res = unsafe {
+            libc::mremap(
+                mmaped_addr,
+                aligned_length,
+                aligned_length,
+                libc::MREMAP_MAYMOVE | libc::MREMAP_FIXED,
+                // Any previous mapping at the address range specified by new_address and new_size is unmapped.
+                // So, no memory leak
+                aligned_addr,
+            )
+        };
+        if res == libc::MAP_FAILED {
+            panic!("Failed to mremap.");
+        }
+        let res = unsafe {
+            clear_cache::clear_cache(addr, addr.add(L));
+        };
+        if !res {
+            panic!("Failed to clear cache.");
         }
     }
 }