[PATCH v2] rust: add a sample alloc usage

[Hui Zhu <hui.zhu@linux.dev>]

From: Hui Zhu <zhuhui@kylinos.cn>

Add a sample to the samples memory alloc usage.
It will samples the usage of VBox and KVec.

It also includes a test for allocating alignments larger than PAGE_SIZE
using VBox.
Since this feature is currently unsupported, it will report errors even
when successful.

Changelog:
v2:
According to the comments of Danilo, updated the commit to samples the
usage of VBox and KVec.

Co-developed-by: Geliang Tang <geliang@kernel.org>
Signed-off-by: Geliang Tang <geliang@kernel.org>
Signed-off-by: Hui Zhu <zhuhui@kylinos.cn>
---
 samples/rust/Kconfig       | 10 +++++
 samples/rust/Makefile      |  1 +
 samples/rust/rust_alloc.rs | 88 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 samples/rust/rust_alloc.rs

diff --git a/samples/rust/Kconfig b/samples/rust/Kconfig
index 7f7371a004ee..03a86f819768 100644
--- a/samples/rust/Kconfig
+++ b/samples/rust/Kconfig
@@ -105,6 +105,16 @@ config SAMPLE_RUST_DRIVER_AUXILIARY
 
 	  If unsure, say N.
 
+config SAMPLE_RUST_ALLOC
+	tristate "Alloc Test Driver"
+	help
+	  This option builds the Rust alloc Test driver sample.
+
+	  To compile this as a module, choose M here:
+	  the module will be called rust_alloc.
+
+	  If unsure, say N.
+
 config SAMPLE_RUST_HOSTPROGS
 	bool "Host programs"
 	help
diff --git a/samples/rust/Makefile b/samples/rust/Makefile
index bd2faad63b4f..7c3e68d9ada5 100644
--- a/samples/rust/Makefile
+++ b/samples/rust/Makefile
@@ -10,6 +10,7 @@ obj-$(CONFIG_SAMPLE_RUST_DRIVER_PLATFORM)	+= rust_driver_platform.o
 obj-$(CONFIG_SAMPLE_RUST_DRIVER_FAUX)		+= rust_driver_faux.o
 obj-$(CONFIG_SAMPLE_RUST_DRIVER_AUXILIARY)	+= rust_driver_auxiliary.o
 obj-$(CONFIG_SAMPLE_RUST_CONFIGFS)		+= rust_configfs.o
+obj-$(CONFIG_SAMPLE_RUST_ALLOC)		+= rust_alloc.o
 
 rust_print-y := rust_print_main.o rust_print_events.o
 
diff --git a/samples/rust/rust_alloc.rs b/samples/rust/rust_alloc.rs
new file mode 100644
index 000000000000..61efde37b5f2
--- /dev/null
+++ b/samples/rust/rust_alloc.rs
@@ -0,0 +1,88 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (c) 2025, Kylin Software
+
+//! Rust alloc sample.
+
+use kernel::bindings;
+use kernel::prelude::*;
+
+module! {
+    type: RustAlloc,
+    name: "rust_alloc",
+    authors: ["Rust for Linux Contributors"],
+    description: "Rust alloc sample",
+    license: "GPL",
+}
+
+const VBOX_SIZE: usize = 1024;
+const VBOX_LARGE_ALIGN: usize = bindings::PAGE_SIZE * 4;
+const KVEC_VAL: [usize; 3] = [10, 20, 30];
+
+#[repr(align(128))]
+struct VboxBlob([u8; VBOX_SIZE]);
+
+// This structure is used to test the allocation of alignments larger
+// than PAGE_SIZE.
+// Since this is not yet supported, avoid accessing the contents of
+// the structure for now.
+#[allow(dead_code)]
+#[repr(align(8192))]
+struct VboxLargeAlignBlob([u8; VBOX_LARGE_ALIGN]);
+
+struct RustAlloc {
+    vbox_blob: VBox<VboxBlob>,
+    kvec_blob: KVec<usize>,
+}
+
+fn check_align(addr: usize, align: usize) -> bool {
+    debug_assert!(align.is_power_of_two());
+    if addr & (align - 1) != 0 {
+        pr_err!("Address {:#x} is not aligned with {:#x}.\n", addr, align);
+        false
+    } else {
+        true
+    }
+}
+
+impl kernel::Module for RustAlloc {
+    fn init(_module: &'static ThisModule) -> Result<Self> {
+        pr_info!("Rust allocator sample (init)\n");
+
+        let vbox_blob = VBox::<VboxBlob>::new_uninit(GFP_KERNEL)?;
+        if !check_align(vbox_blob.as_ptr() as usize, 128) {
+            return Err(EINVAL);
+        }
+        let vbox_blob = vbox_blob.write(VboxBlob([0xfeu8; VBOX_SIZE]));
+
+        if let Ok(_) = VBox::<VboxLargeAlignBlob>::new_uninit(GFP_KERNEL) {
+            pr_err!("Allocations for VBox with alignments larger than PAGE_SIZE should fail, but here it succeeded.\n");
+            return Err(EINVAL);
+        }
+
+        let mut kvec_blob = KVec::new();
+        kvec_blob.extend_from_slice(&KVEC_VAL, GFP_KERNEL)?;
+
+        Ok(Self {
+            vbox_blob,
+            kvec_blob,
+        })
+    }
+}
+
+impl Drop for RustAlloc {
+    fn drop(&mut self) {
+        pr_info!("Rust allocator sample (exit)\n");
+
+        // check the values
+        for b in self.vbox_blob.0.as_slice().iter() {
+            if *b != 0xfeu8 {
+                pr_err!("vbox_blob contains wrong values\n");
+            }
+        }
+
+        if self.kvec_blob.as_slice() != KVEC_VAL {
+            pr_err!("kvec_blob contains wrong values\n");
+        }
+    }
+}
-- 
2.43.0