Initial forward-edge CFI implementation (#3693)

* Initial forward-edge CFI implementation

Give the user the option to start all basic blocks that are targets
of indirect branches with the BTI instruction introduced by the
Branch Target Identification extension to the Arm instruction set
architecture.

Copyright (c) 2022, Arm Limited.

* Refactor `from_artifacts` to avoid second `make_executable` (#1)

This involves "parsing" twice but this is parsing just the header of an
ELF file so it's not a very intensive operation and should be ok to do
twice.

* Address the code review feedback

Copyright (c) 2022, Arm Limited.

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

cranelift/jit/src/backend.rs

@@ -1,6 +1,6 @@
 //! Defines `JITModule`.
 
-use crate::{compiled_blob::CompiledBlob, memory::Memory};
+use crate::{compiled_blob::CompiledBlob, memory::BranchProtection, memory::Memory};
 use cranelift_codegen::isa::TargetIsa;
 use cranelift_codegen::settings::Configurable;
 use cranelift_codegen::{self, ir, settings, MachReloc};
@@ -480,6 +480,12 @@ impl JITModule {
             );
         }
 
+        let branch_protection =
+            if cfg!(target_arch = "aarch64") && use_bti(&builder.isa.isa_flags()) {
+                BranchProtection::BTI
+            } else {
+                BranchProtection::None
+            };
         let mut module = Self {
             isa: builder.isa,
             hotswap_enabled: builder.hotswap_enabled,
@@ -487,9 +493,10 @@ impl JITModule {
             lookup_symbols: builder.lookup_symbols,
             libcall_names: builder.libcall_names,
             memory: MemoryHandle {
-                code: Memory::new(),
-                readonly: Memory::new(),
-                writable: Memory::new(),
+                code: Memory::new(branch_protection),
+                // Branch protection is not applicable to non-executable memory.
+                readonly: Memory::new(BranchProtection::None),
+                writable: Memory::new(BranchProtection::None),
             },
             declarations: ModuleDeclarations::default(),
             function_got_entries: SecondaryMap::new(),
@@ -959,3 +966,10 @@ fn lookup_with_dlsym(name: &str) -> Option<*const u8> {
         None
     }
 }
+
+fn use_bti(isa_flags: &Vec<settings::Value>) -> bool {
+    isa_flags
+        .iter()
+        .find(|&f| f.name == "use_bti")
+        .map_or(false, |f| f.as_bool().unwrap_or(false))
+}

cranelift/jit/src/memory.rs

@@ -104,6 +104,15 @@ impl Drop for PtrLen {
 
 // TODO: add a `Drop` impl for `cfg(target_os = "windows")`
 
+/// Type of branch protection to apply to executable memory.
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) enum BranchProtection {
+    /// No protection.
+    None,
+    /// Use the Branch Target Identification extension of the Arm architecture.
+    BTI,
+}
+
 /// JIT memory manager. This manages pages of suitably aligned and
 /// accessible memory. Memory will be leaked by default to have
 /// function pointers remain valid for the remainder of the
@@ -113,15 +122,17 @@ pub(crate) struct Memory {
     already_protected: usize,
     current: PtrLen,
     position: usize,
+    branch_protection: BranchProtection,
 }
 
 impl Memory {
-    pub(crate) fn new() -> Self {
+    pub(crate) fn new(branch_protection: BranchProtection) -> Self {
         Self {
             allocations: Vec::new(),
             already_protected: 0,
             current: PtrLen::new(),
             position: 0,
+            branch_protection,
         }
     }
 
@@ -157,14 +168,35 @@ impl Memory {
     pub(crate) fn set_readable_and_executable(&mut self) {
         self.finish_current();
 
+        let set_region_readable_and_executable = |ptr, len| {
+            if len != 0 {
+                if self.branch_protection == BranchProtection::BTI {
+                    #[cfg(all(target_arch = "aarch64", target_os = "linux"))]
+                    if std::arch::is_aarch64_feature_detected!("bti") {
+                        let prot = libc::PROT_EXEC | libc::PROT_READ | /* PROT_BTI */ 0x10;
+
+                        unsafe {
+                            if libc::mprotect(ptr as *mut libc::c_void, len, prot) < 0 {
+                                panic!("unable to make memory readable+executable");
+                            }
+                        }
+
+                        return;
+                    }
+                }
+
+                unsafe {
+                    region::protect(ptr, len, region::Protection::READ_EXECUTE)
+                        .expect("unable to make memory readable+executable");
+                }
+            }
+        };
+
         #[cfg(feature = "selinux-fix")]
         {
             for &PtrLen { ref map, ptr, len } in &self.allocations[self.already_protected..] {
-                if len != 0 && map.is_some() {
-                    unsafe {
-                        region::protect(ptr, len, region::Protection::READ_EXECUTE)
-                            .expect("unable to make memory readable+executable");
-                    }
+                if map.is_some() {
+                    set_region_readable_and_executable(ptr, len);
                 }
             }
         }
@@ -172,12 +204,7 @@ impl Memory {
         #[cfg(not(feature = "selinux-fix"))]
         {
             for &PtrLen { ptr, len } in &self.allocations[self.already_protected..] {
-                if len != 0 {
-                    unsafe {
-                        region::protect(ptr, len, region::Protection::READ_EXECUTE)
-                            .expect("unable to make memory readable+executable");
-                    }
-                }
+                set_region_readable_and_executable(ptr, len);
             }
         }