cranelift: improve syscall error/oom handling in JIT module (#5173)

* cranelift: improve syscall error/oom handling in JIT module

The JIT module has several places where it `expect`s or `panic`s
on syscall or allocator errors. For example, `mmap` and `mprotect`
can fail if Linux `vm.max_map_count` is not high enough, and some
users may wish to handle this error rather than immediately
crashing.

This commit plumbs these errors upward as new `ModuleError`
types, so that callers of jit module functions like
`finalize_definitions` and `define_function` can handle them
(or just `unwrap()`, as desired).

* cranelift: Remove ModuleError::Syscall variant

Syscall errors can just be folded into the generic Backend error,
which is an anyhow::Error

* cranelift-jit: return io::ErrorKind::OutOfMemory for alloc failure

Just using `io::Error::last_os_error()` is not correct as global
allocator impls are not required to set errno

cranelift/filetests/src/function_runner.rs

@@ -244,7 +244,7 @@ impl TestFileCompiler {
         // Finalize the functions which we just defined, which resolves any
         // outstanding relocations (patching in addresses, now that they're
         // available).
-        self.module.finalize_definitions();
+        self.module.finalize_definitions()?;
 
         Ok(CompiledTestFile {
             module: Some(self.module),

cranelift/jit/examples/jit-minimal.rs

@@ -78,7 +78,7 @@ fn main() {
     module.clear_context(&mut ctx);
 
     // Perform linking.
-    module.finalize_definitions();
+    module.finalize_definitions().unwrap();
 
     // Get a raw pointer to the generated code.
     let code_b = module.get_finalized_function(func_b);

cranelift/jit/src/backend.rs

@@ -427,7 +427,9 @@ impl JITModule {
     ///
     /// Use `get_finalized_function` and `get_finalized_data` to obtain the final
     /// artifacts.
-    pub fn finalize_definitions(&mut self) {
+    ///
+    /// Returns ModuleError in case of allocation or syscall failure
+    pub fn finalize_definitions(&mut self) -> ModuleResult<()> {
         for func in std::mem::take(&mut self.functions_to_finalize) {
             let decl = self.declarations.get_function_decl(func);
             assert!(decl.linkage.is_definable());
@@ -455,12 +457,13 @@ impl JITModule {
         }
 
         // Now that we're done patching, prepare the memory for execution!
-        self.memory.readonly.set_readonly();
+        self.memory.readonly.set_readonly()?;
-        self.memory.code.set_readable_and_executable();
+        self.memory.code.set_readable_and_executable()?;
 
         for update in self.pending_got_updates.drain(..) {
             unsafe { update.entry.as_ref() }.store(update.ptr as *mut _, Ordering::SeqCst);
         }
+        Ok(())
     }
 
     /// Create a new `JITModule`.
@@ -688,7 +691,10 @@ impl Module for JITModule {
             .memory
             .code
             .allocate(size, align)
-            .expect("TODO: handle OOM etc.");
+            .map_err(|e| ModuleError::Allocation {
+                message: "unable to alloc function",
+                err: e,
+            })?;
 
         {
             let mem = unsafe { std::slice::from_raw_parts_mut(ptr, size) };
@@ -770,7 +776,10 @@ impl Module for JITModule {
             .memory
             .code
             .allocate(size, align)
-            .expect("TODO: handle OOM etc.");
+            .map_err(|e| ModuleError::Allocation {
+                message: "unable to alloc function bytes",
+                err: e,
+            })?;
 
         unsafe {
             ptr::copy_nonoverlapping(bytes.as_ptr(), ptr, size);
@@ -836,12 +845,18 @@ impl Module for JITModule {
             self.memory
                 .writable
                 .allocate(size, align.unwrap_or(WRITABLE_DATA_ALIGNMENT))
-                .expect("TODO: handle OOM etc.")
+                .map_err(|e| ModuleError::Allocation {
+                    message: "unable to alloc writable data",
+                    err: e,
+                })?
         } else {
             self.memory
                 .readonly
                 .allocate(size, align.unwrap_or(READONLY_DATA_ALIGNMENT))
-                .expect("TODO: handle OOM etc.")
+                .map_err(|e| ModuleError::Allocation {
+                    message: "unable to alloc readonly data",
+                    err: e,
+                })?
         };
 
         match *init {

cranelift/jit/src/memory.rs

@@ -1,3 +1,5 @@
+use cranelift_module::{ModuleError, ModuleResult};
+
 #[cfg(feature = "selinux-fix")]
 use memmap2::MmapMut;
 
@@ -56,10 +58,14 @@ impl PtrLen {
         // Safety: We assert that the size is non-zero above.
         let ptr = unsafe { alloc::alloc(layout) };
 
+        if !ptr.is_null() {
             Ok(Self {
                 ptr,
                 len: alloc_size,
             })
+        } else {
+            Err(io::Error::from(io::ErrorKind::OutOfMemory))
+        }
     }
 
     #[cfg(target_os = "windows")]
@@ -168,7 +174,7 @@ impl Memory {
     }
 
     /// Set all memory allocated in this `Memory` up to now as readable and executable.
-    pub(crate) fn set_readable_and_executable(&mut self) {
+    pub(crate) fn set_readable_and_executable(&mut self) -> ModuleResult<()> {
         self.finish_current();
 
         // Clear all the newly allocated code from cache if the processor requires it
@@ -182,7 +188,7 @@ impl Memory {
             };
         }
 
-        let set_region_readable_and_executable = |ptr, len| {
+        let set_region_readable_and_executable = |ptr, len| -> ModuleResult<()> {
             if self.branch_protection == BranchProtection::BTI {
                 #[cfg(all(target_arch = "aarch64", target_os = "linux"))]
                 if std::arch::is_aarch64_feature_detected!("bti") {
@@ -190,42 +196,54 @@ impl Memory {
 
                     unsafe {
                         if libc::mprotect(ptr as *mut libc::c_void, len, prot) < 0 {
-                            panic!("unable to make memory readable+executable");
+                            return Err(ModuleError::Backend(
+                                anyhow::Error::new(io::Error::last_os_error())
+                                    .context("unable to make memory readable+executable"),
+                            ));
                         }
                     }
 
-                    return;
+                    return Ok(());
                 }
             }
 
             unsafe {
-                region::protect(ptr, len, region::Protection::READ_EXECUTE)
+                region::protect(ptr, len, region::Protection::READ_EXECUTE).map_err(|e| {
-                    .expect("unable to make memory readable+executable");
+                    ModuleError::Backend(
+                        anyhow::Error::new(e).context("unable to make memory readable+executable"),
+                    )
+                })?;
             }
+            Ok(())
         };
 
         for &PtrLen { ptr, len, .. } in self.non_protected_allocations_iter() {
-            set_region_readable_and_executable(ptr, len);
+            set_region_readable_and_executable(ptr, len)?;
         }
 
         // Flush any in-flight instructions from the pipeline
         icache_coherence::pipeline_flush_mt().expect("Failed pipeline flush");
 
         self.already_protected = self.allocations.len();
+        Ok(())
     }
 
     /// Set all memory allocated in this `Memory` up to now as readonly.
-    pub(crate) fn set_readonly(&mut self) {
+    pub(crate) fn set_readonly(&mut self) -> ModuleResult<()> {
         self.finish_current();
 
         for &PtrLen { ptr, len, .. } in self.non_protected_allocations_iter() {
             unsafe {
-                region::protect(ptr, len, region::Protection::READ)
+                region::protect(ptr, len, region::Protection::READ).map_err(|e| {
-                    .expect("unable to make memory readonly");
+                    ModuleError::Backend(
+                        anyhow::Error::new(e).context("unable to make memory readonly"),
+                    )
+                })?;
             }
         }
 
         self.already_protected = self.allocations.len();
+        Ok(())
     }
 
     /// Iterates non protected memory allocations that are of not zero bytes in size.

cranelift/jit/tests/basic.rs

@@ -209,5 +209,5 @@ fn libcall_function() {
 
     module.define_function(func_id, &mut ctx).unwrap();
 
-    module.finalize_definitions();
+    module.finalize_definitions().unwrap();
 }

cranelift/module/src/module.rs

@@ -229,6 +229,14 @@ pub enum ModuleError {
     /// Wraps a `cranelift-codegen` error
     Compilation(CodegenError),
 
+    /// Memory allocation failure from a backend
+    Allocation {
+        /// Tell where the allocation came from
+        message: &'static str,
+        /// Io error the allocation failed with
+        err: std::io::Error,
+    },
+
     /// Wraps a generic error from a backend
     Backend(anyhow::Error),
 }
@@ -250,6 +258,7 @@ impl std::error::Error for ModuleError {
             | Self::DuplicateDefinition { .. }
             | Self::InvalidImportDefinition { .. } => None,
             Self::Compilation(source) => Some(source),
+            Self::Allocation { err: source, .. } => Some(source),
             Self::Backend(source) => Some(&**source),
         }
     }
@@ -284,6 +293,9 @@ impl std::fmt::Display for ModuleError {
             Self::Compilation(err) => {
                 write!(f, "Compilation error: {}", err)
             }
+            Self::Allocation { message, err } => {
+                write!(f, "Allocation error: {}: {}", message, err)
+            }
             Self::Backend(err) => write!(f, "Backend error: {}", err),
         }
     }