Merge branch 'Evian-Zhang:master' into master

Cargo.toml

@@ -17,12 +17,13 @@ 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"
 
 [target.'cfg(target_os = "windows")'.dependencies]
-windows = { version = "0.58", features = ["Win32_System_SystemInformation", "Win32_System_Memory"] }
+windows = { version = "0.58", features = ["Win32_System_SystemInformation", "Win32_System_Memory", "Win32_System_Threading", "Win32_System_Diagnostics_Debug"] }
 
 [dev-dependencies]
 trybuild = "1"

src/arch/x86.rs

@@ -43,16 +43,17 @@ macro_rules! arch_static_key_init_nop_asm_template {
     };
 }
 
-// The `0x90,0x90,0x90` are three NOPs, which is to make sure the `jmp {0}` is at least 5 bytes long.
+// Here we do not use `jmp {0}` because it may be compiled into a 3-byte jmp instead of 5 byte.
+// See https://stackoverflow.com/q/74771372/10005095
 #[doc(hidden)]
 #[macro_export]
 macro_rules! arch_static_key_init_jmp_asm_template {
     () => {
         ::core::concat!(
             r#"
             2:
-                jmp {0}
-            .byte 0x90,0x90,0x90
+            .byte 0xe9
+            .long ({0} - 4) - .
             .pushsection "#,
             $crate::os_static_key_sec_name_attr!(),
             r#"

src/arch/x86_64.rs

@@ -43,16 +43,17 @@ macro_rules! arch_static_key_init_nop_asm_template {
     };
 }
 
-// The `0x90,0x90,0x90` are three NOPs, which is to make sure the `jmp {0}` is at least 5 bytes long.
+// Here we do not use `jmp {0}` because it may be compiled into a 3-byte jmp instead of 5 byte.
+// See https://stackoverflow.com/q/74771372/10005095
 #[doc(hidden)]
 #[macro_export]
 macro_rules! arch_static_key_init_jmp_asm_template {
     () => {
         ::core::concat!(
             r#"
             2:
-                jmp {0}
-            .byte 0x90,0x90,0x90
+            .byte 0xe9
+            .long ({0} - 4) - .
             .pushsection "#,
             $crate::os_static_key_sec_name_attr!(),
             r#"

src/os/linux.rs

@@ -40,26 +40,54 @@ 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.");
         }
     }
 }

src/os/windows.rs

@@ -1,8 +1,10 @@
 //! Windows-specific implementations
 
 use windows::Win32::System::{
+    Diagnostics::Debug::FlushInstructionCache,
     Memory::{VirtualProtect, PAGE_EXECUTE_READWRITE, PAGE_PROTECTION_FLAGS},
     SystemInformation::{GetSystemInfo, SYSTEM_INFO},
+    Threading::GetCurrentProcess,
 };
 
 use crate::{code_manipulate::CodeManipulator, JumpEntry};
@@ -69,5 +71,9 @@ impl CodeManipulator for ArchCodeManipulator {
         if res.is_err() {
             panic!("Unable to restore code region to non-writable");
         }
+        let res = unsafe { FlushInstructionCache(GetCurrentProcess(), Some(addr), L) };
+        if res.is_err() {
+            panic!("Failed to flush instruction cache");
+        }
     }
 }