Add OS-specific implementations; Use three 0x90 instead of 3-byte NOP…

…s in case the jmp is not 2 bytes long

src/arch/x86.rs

@@ -22,22 +22,34 @@ pub fn arch_jump_entry_instruction(
     }
 }
 
-/// With given branch as likely branch, initialize the instruction here as a 5-byte NOP instruction
 #[doc(hidden)]
 #[macro_export]
-macro_rules! arch_static_key_init_nop_with_given_branch_likely {
-    ($key:path, $branch:expr) => {'my_label: {
-        ::core::arch::asm!(
+macro_rules! arch_static_key_init_nop_asm_template {
+    () => {
+        ::core::concat!(
             r#"
             2:
             .byte 0x3e,0x8d,0x74,0x26,0x00
-            .pushsection __static_keys, "awR"
+            .pushsection "#,
+            $crate::os_static_key_sec_name!(),
+            r#", "awR"
             .balign 4
             .long 2b - .
             .long {0} - .
             .long {1} + {2} - .
             .popsection
-            "#,
+            "#
+        )
+    };
+}
+
+/// With given branch as likely branch, initialize the instruction here as a 5-byte NOP instruction
+#[doc(hidden)]
+#[macro_export]
+macro_rules! arch_static_key_init_nop_with_given_branch_likely {
+    ($key:path, $branch:expr) => {'my_label: {
+        ::core::arch::asm!(
+            $crate::arch_static_key_init_nop_asm_template!(),
             label {
                 break 'my_label !$branch;
             },
@@ -50,24 +62,36 @@ macro_rules! arch_static_key_init_nop_with_given_branch_likely {
     }};
 }
 
-// The `0x8d,0x76,0x00` is a 3-byte NOP, which is to make sure the `jmp {0}` is at least 5 bytes long.
-/// With given branch as likely branch, initialize the instruction here as JMP instruction
+// The `0x90,0x90,0x90` are three NOPs, which is to make sure the `jmp {0}` is at least 5 bytes long.
 #[doc(hidden)]
 #[macro_export]
-macro_rules! arch_static_key_init_jmp_with_given_branch_likely {
-    ($key:path, $branch:expr) => {'my_label: {
-        ::core::arch::asm!(
+macro_rules! arch_static_key_init_jmp_asm_template {
+    () => {
+        ::core::concat!(
             r#"
-            2: 
+            2:
                 jmp {0}
-            .byte 0x8d,0x76,0x00
-            .pushsection __static_keys, "awR"
+            .byte 0x90,0x90,0x90
+            .pushsection "#,
+            $crate::os_static_key_sec_name!(),
+            r#", "awR"
             .balign 4
             .long 2b - .
             .long {0} - .
             .long {1} + {2} - .
             .popsection
-            "#,
+            "#
+        )
+    };
+}
+
+/// With given branch as likely branch, initialize the instruction here as JMP instruction
+#[doc(hidden)]
+#[macro_export]
+macro_rules! arch_static_key_init_jmp_with_given_branch_likely {
+    ($key:path, $branch:expr) => {'my_label: {
+        ::core::arch::asm!(
+            $crate::arch_static_key_init_jmp_asm_template!(),
             label {
                 break 'my_label !$branch;
             },

src/arch/x86_64.rs

@@ -22,22 +22,34 @@ pub fn arch_jump_entry_instruction(
     }
 }
 
-/// With given branch as likely branch, initialize the instruction here as a 5-byte NOP instruction
 #[doc(hidden)]
 #[macro_export]
-macro_rules! arch_static_key_init_nop_with_given_branch_likely {
-    ($key:path, $branch:expr) => {'my_label: {
-        ::core::arch::asm!(
+macro_rules! arch_static_key_init_nop_asm_template {
+    () => {
+        ::core::concat!(
             r#"
             2:
             .byte 0x0f,0x1f,0x44,0x00,0x00
-            .pushsection __static_keys, "awR"
+            .pushsection "#,
+            $crate::os_static_key_sec_name!(),
+            r#", "awR"
             .balign 8
             .quad 2b - .
             .quad {0} - .
             .quad {1} + {2} - .
             .popsection
-            "#,
+            "#
+        )
+    };
+}
+
+/// With given branch as likely branch, initialize the instruction here as a 5-byte NOP instruction
+#[doc(hidden)]
+#[macro_export]
+macro_rules! arch_static_key_init_nop_with_given_branch_likely {
+    ($key:path, $branch:expr) => {'my_label: {
+        ::core::arch::asm!(
+            $crate::arch_static_key_init_nop_asm_template!(),
             label {
                 break 'my_label !$branch;
             },
@@ -50,24 +62,36 @@ macro_rules! arch_static_key_init_nop_with_given_branch_likely {
     }};
 }
 
-// The `0x0f,0x1f,0x00` is a 3-byte NOP, which is to make sure the `jmp {0}` is at least 5 bytes long.
-/// With given branch as likely branch, initialize the instruction here as JMP instruction
+// The `0x90,0x90,0x90` are three NOPs, which is to make sure the `jmp {0}` is at least 5 bytes long.
 #[doc(hidden)]
 #[macro_export]
-macro_rules! arch_static_key_init_jmp_with_given_branch_likely {
-    ($key:path, $branch:expr) => {'my_label: {
-        ::core::arch::asm!(
+macro_rules! arch_static_key_init_jmp_asm_template {
+    () => {
+        ::core::concat!(
             r#"
-            2: 
+            2:
                 jmp {0}
-            .byte 0x0f,0x1f,0x00
-            .pushsection __static_keys, "awR"
+            .byte 0x90,0x90,0x90
+            .pushsection "#,
+            $crate::os_static_key_sec_name!(),
+            r#", "awR"
             .balign 8
             .quad 2b - .
             .quad {0} - .
             .quad {1} + {2} - .
             .popsection
-            "#,
+            "#
+        )
+    };
+}
+
+/// With given branch as likely branch, initialize the instruction here as JMP instruction
+#[doc(hidden)]
+#[macro_export]
+macro_rules! arch_static_key_init_jmp_with_given_branch_likely {
+    ($key:path, $branch:expr) => {'my_label: {
+        ::core::arch::asm!(
+            $crate::arch_static_key_init_jmp_asm_template!(),
             label {
                 break 'my_label !$branch;
             },

src/lib.rs

@@ -8,6 +8,7 @@ extern crate std;
 
 mod arch;
 pub mod code_manipulate;
+mod os;
 
 use code_manipulate::CodeManipulator;
 
@@ -70,18 +71,6 @@ impl JumpEntry {
     }
 }
 
-// See https://sourceware.org/binutils/docs/ld/Input-Section-Example.html, modern linkers
-// will generate these two symbols indicating the start and end address of __static_keys
-// section. Note that the end address is excluded.
-extern "Rust" {
-    /// Address of this static is the start address of __static_keys section
-    #[link_name = "__start___static_keys"]
-    static mut JUMP_ENTRY_START: JumpEntry;
-    /// Address of this static is the end address of __styatic_keys section (excluded)
-    #[link_name = "__stop___static_keys"]
-    static mut JUMP_ENTRY_STOP: JumpEntry;
-}
-
 /// Static key to hold data about current status and which jump entries are associated with this key.
 ///
 /// For now, it is not encouraged to modify static key in a multi-thread application (which I don't think
@@ -173,8 +162,8 @@ impl<M: CodeManipulator, const S: bool> NoStdStaticKey<M, S> {
         if static_branch_unlikely!(DUMMY_STATIC_KEY) {
             return;
         }
-        let jump_entry_start_addr = core::ptr::addr_of_mut!(JUMP_ENTRY_START);
-        let jump_entry_stop_addr = core::ptr::addr_of_mut!(JUMP_ENTRY_STOP);
+        let jump_entry_start_addr = core::ptr::addr_of_mut!(os::JUMP_ENTRY_START);
+        let jump_entry_stop_addr = core::ptr::addr_of_mut!(os::JUMP_ENTRY_STOP);
         let jump_entry_len =
             unsafe { jump_entry_stop_addr.offset_from(jump_entry_start_addr) as usize };
         let jump_entries =
@@ -212,8 +201,8 @@ impl<M: CodeManipulator, const S: bool> NoStdStaticKey<M, S> {
 /// Count of jump entries in __static_keys section. Note that
 /// there will be several dummy jump entries inside this section.
 pub fn jump_entries_count() {
-    let jump_entry_start_addr = core::ptr::addr_of_mut!(JUMP_ENTRY_START);
-    let jump_entry_stop_addr = core::ptr::addr_of_mut!(JUMP_ENTRY_STOP);
+    let jump_entry_start_addr = core::ptr::addr_of_mut!(os::JUMP_ENTRY_START);
+    let jump_entry_stop_addr = core::ptr::addr_of_mut!(os::JUMP_ENTRY_STOP);
     unsafe { jump_entry_stop_addr.offset_from(jump_entry_start_addr) as usize };
 }
 
@@ -301,7 +290,7 @@ unsafe fn static_key_update<M: CodeManipulator, const S: bool>(
         return;
     }
     key.enabled = enabled;
-    let jump_entry_stop_addr = core::ptr::addr_of!(JUMP_ENTRY_STOP);
+    let jump_entry_stop_addr = core::ptr::addr_of!(os::JUMP_ENTRY_STOP);
     let mut jump_entry_addr = key.entries();
     if jump_entry_addr.is_null() {
         // This static key is never used
@@ -340,7 +329,10 @@ unsafe fn jump_entry_update<M: CodeManipulator>(jump_entry: &JumpEntry, enabled:
     };
     let code_bytes = arch::arch_jump_entry_instruction(jump_label_type, jump_entry);
 
-    let manipulator = M::mark_code_region_writable(jump_entry.code_addr() as *const _, 5);
+    let manipulator = M::mark_code_region_writable(
+        jump_entry.code_addr() as *const _,
+        arch::ARCH_JUMP_INS_LENGTH,
+    );
     core::ptr::copy_nonoverlapping(
         code_bytes.as_ptr(),
         jump_entry.code_addr() as usize as *mut u8,

src/os/linux.rs

@@ -0,0 +1,24 @@
+//! Linux-specific implementations
+
+use crate::JumpEntry;
+
+/// Section name of
+#[doc(hidden)]
+#[macro_export]
+macro_rules! os_static_key_sec_name {
+    () => {
+        "__static_keys"
+    };
+}
+
+// See https://sourceware.org/binutils/docs/ld/Input-Section-Example.html, modern linkers
+// will generate these two symbols indicating the start and end address of __static_keys
+// section. Note that the end address is excluded.
+extern "Rust" {
+    /// Address of this static is the start address of __static_keys section
+    #[link_name = "__start___static_keys"]
+    pub static mut JUMP_ENTRY_START: JumpEntry;
+    /// Address of this static is the end address of __static_keys section (excluded)
+    #[link_name = "__stop___static_keys"]
+    pub static mut JUMP_ENTRY_STOP: JumpEntry;
+}

src/os/mod.rs

@@ -0,0 +1,6 @@
+//! OS-specific implementations
+
+#[cfg(target_os = "linux")]
+mod linux;
+#[cfg(target_os = "linux")]
+pub use linux::*;