From 8f89f333e7d41d1240bd19d31480c04bf629ab77 Mon Sep 17 00:00:00 2001
From: Andreas Hindborg <a.hindborg@samsung.com>
Date: Mon, 4 Sep 2023 13:35:30 +0200
Subject: rust: allow allocation with gfp_t flags

Signed-off-by: Andreas Hindborg <a.hindborg@samsung.com>
---
 rust/kernel/allocator.rs | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/rust/kernel/allocator.rs b/rust/kernel/allocator.rs
index a8f3d5be1af1..a6fb23601f09 100644
--- a/rust/kernel/allocator.rs
+++ b/rust/kernel/allocator.rs
@@ -2,12 +2,44 @@
 
 //! Allocator support.
 
+use core::alloc::AllocError;
 use core::alloc::{GlobalAlloc, Layout};
 use core::ptr;
+use core::ptr::NonNull;
 
 use crate::bindings;
 
-struct KernelAllocator;
+pub(crate) struct KernelAllocator;
+
+impl KernelAllocator {
+    #[cfg(not(test))]
+    #[cfg(not(testlib))]
+    pub(crate) fn allocate_with_flags(
+        &self,
+        layout: Layout,
+        flags: bindings::gfp_t,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        // `krealloc()` is used instead of `kmalloc()` because the latter is
+        // an inline function and cannot be bound to as a result.
+        let mem = unsafe { bindings::krealloc(ptr::null(), layout.size(), flags) as *mut u8 };
+        if mem.is_null() {
+            return Err(AllocError);
+        }
+        let mem = unsafe { core::slice::from_raw_parts_mut(mem, bindings::ksize(mem as _)) };
+        // Safety: checked for non null above
+        Ok(unsafe { NonNull::new_unchecked(mem) })
+    }
+
+    #[cfg(test)]
+    #[cfg(testlib)]
+    pub(crate) fn allocate_with_flags(
+        &self,
+        layout: Layout,
+        _flags: bindings::gfp_t,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        self.allocate(layout)
+    }
+}
 
 /// Calls `krealloc` with a proper size to alloc a new object aligned to `new_layout`'s alignment.
 ///
@@ -81,7 +113,7 @@ unsafe impl GlobalAlloc for KernelAllocator {
 }
 
 #[global_allocator]
-static ALLOCATOR: KernelAllocator = KernelAllocator;
+pub(crate) static ALLOCATOR: KernelAllocator = KernelAllocator;
 
 // See <https://github.com/rust-lang/rust/pull/86844>.
 #[no_mangle]
-- 
cgit v1.2.3-58-ga151