[PATCH 1/1] samples/rust: Add jump_label sample demonstrating static branches on s390

[Jan Polensky]

This sample demonstrates the use of static branches (jump labels) in Rust
kernel modules on s390 architecture. It successfully generates __jump_table
entries and uses s390-specific jgnop instructions for runtime code patching.

Key components:
- C helper defining DEFINE_STATIC_KEY_FALSE(rust_sample_key)
- Rust module using kernel::jump_label::arch_static_branch! macro
- Composite module build (C + Rust) with proper symbol binding
- Generates real __jump_table section with relocations
- Uses s390 jgnop instruction (6-byte NOP) for branch patching

Verification shows:
- __jump_table section present in compiled module
- Three relocations: code location, target, and key pointer
- s390 jgnop instruction at offset 0x102

The module tests the static branch in a loop, demonstrating that Rust
can leverage the kernel's jump label infrastructure for optimized
conditional branches that can be patched at runtime.

Signed-off-by: Jan Polensky <polensky@xxxxxxxxxxxxx>
Signed-off-by: Jan Polensky <japo@xxxxxxxxxxxxx>
---
 rust/bindings/bindings_helper.h       |  5 ++
 samples/rust/Kconfig                  | 15 +++++
 samples/rust/Makefile                 |  4 +-
 samples/rust/rust_jump_label_core.rs  | 79 +++++++++++++++++++++++++++
 samples/rust/rust_jump_label_helper.c | 20 +++++++
 5 files changed, 122 insertions(+), 1 deletion(-)
 create mode 100644 samples/rust/rust_jump_label_core.rs
 create mode 100644 samples/rust/rust_jump_label_helper.c

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index 446dbeaf0866..539c8d18a72c 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -164,3 +164,8 @@ const unsigned long RUST_CONST_HELPER_GPU_BUDDY_TRIM_DISABLE = GPU_BUDDY_TRIM_DI
 #include "../../drivers/android/binder/rust_binder.h"
 #include "../../drivers/android/binder/rust_binder_events.h"
 #endif
+
+/* Rust jump label sample - external static key */
+#if IS_ENABLED(CONFIG_SAMPLE_RUST_JUMP_LABEL)
+extern struct static_key_false rust_sample_key;
+#endif
diff --git a/samples/rust/Kconfig b/samples/rust/Kconfig
index 10364db13872..8c579d861988 100644
--- a/samples/rust/Kconfig
+++ b/samples/rust/Kconfig
@@ -203,4 +203,19 @@ config SAMPLE_RUST_TRIGGER
 
 	  If unsure, say N.
 
+config SAMPLE_RUST_JUMP_LABEL
+	tristate "Jump label (static branch)"
+	depends on JUMP_LABEL
+	help
+	  This option builds the Rust jump label sample module.
+
+	  It demonstrates the use of static branches (jump labels) in Rust
+	  kernel modules, particularly testing ARCH_STATIC_BRANCH_ASM
+	  implementation on s390 architecture.
+
+	  To compile this as a module, choose M here:
+	  the module will be called rust_jump_label.
+
+	  If unsure, say N.
+
 endif # SAMPLES_RUST
diff --git a/samples/rust/Makefile b/samples/rust/Makefile
index 037fd43f9878..f91d8a29802a 100644
--- a/samples/rust/Makefile
+++ b/samples/rust/Makefile
@@ -18,8 +18,10 @@ obj-$(CONFIG_SAMPLE_RUST_CONFIGFS)		+= rust_configfs.o
 obj-$(CONFIG_SAMPLE_RUST_SOC)			+= rust_soc.o
 obj-$(CONFIG_SAMPLE_RUST_BUG)			+= rust_bug.o
 obj-$(CONFIG_SAMPLE_RUST_TRIGGER)		+= rust_trigger.o
+obj-$(CONFIG_SAMPLE_RUST_JUMP_LABEL)		+= rust_jump_label.o
 
 rust_print-y := rust_print_main.o rust_print_events.o
 rust_trigger-y := rust_trigger_main.o rust_print_events.o
+rust_jump_label-y := rust_jump_label_helper.o rust_jump_label_core.o
 
-subdir-$(CONFIG_SAMPLE_RUST_HOSTPROGS)		+= hostprogs
+subdir-$(CONFIG_SAMPLE_RUST_HOSTPROGS)		+= hostprogs
\ No newline at end of file
diff --git a/samples/rust/rust_jump_label_core.rs b/samples/rust/rust_jump_label_core.rs
new file mode 100644
index 000000000000..af3efeb3c2a3
--- /dev/null
+++ b/samples/rust/rust_jump_label_core.rs
@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Rust jump label sample.
+//!
+//! Demonstrates static branches (jump labels) using a C-defined static_key.
+//! Generates __jump_table entries and architecture-specific branch instructions.
+
+use kernel::prelude::*;
+use kernel::sync::atomic::{Atomic, Relaxed};
+
+module! {
+    type: RustJumpLabel,
+    name: "rust_jump_label",
+    authors: ["Jan Polensky"],
+    description: "Rust jump label sample",
+    license: "GPL",
+}
+
+struct RustJumpLabel;
+
+impl kernel::Module for RustJumpLabel {
+    fn init(_module: &'static ThisModule) -> Result<Self> {
+        pr_info!("Rust jump label sample (init)\n");
+
+        let counter = Atomic::new(0);
+        Self::test_jump_label(&counter)?;
+
+        Ok(RustJumpLabel)
+    }
+}
+
+impl RustJumpLabel {
+    fn test_jump_label(counter: &Atomic<u64>) -> Result {
+        for i in 0..10 {
+            counter.fetch_add(1, Relaxed);
+
+            // SAFETY: rust_sample_key is defined in C helper and exported.
+            let branch_taken = unsafe {
+                #[cfg(not(CONFIG_JUMP_LABEL))]
+                {
+                    let key_ptr: *const kernel::bindings::static_key_false =
+                        core::ptr::addr_of!(kernel::bindings::rust_sample_key);
+                    let static_key_ptr: *const kernel::bindings::static_key =
+                        core::ptr::addr_of!((*key_ptr).key);
+                    kernel::bindings::static_key_count(static_key_ptr.cast_mut()) > 0
+                }
+
+                #[cfg(CONFIG_JUMP_LABEL)]
+                {
+                    kernel::jump_label::arch_static_branch!(
+                        kernel::bindings::rust_sample_key,
+                        kernel::bindings::static_key_false,
+                        key,
+                        false
+                    )
+                }
+            };
+
+            if branch_taken {
+                pr_info!("Branch taken at iteration {}\n", i);
+            }
+        }
+
+        let final_count = counter.load(Relaxed);
+        if final_count != 10 {
+            pr_err!("Counter mismatch! Expected 10, got {}\n", final_count);
+            return Err(EINVAL);
+        }
+
+        pr_info!("Test completed. Counter: {}\n", final_count);
+        Ok(())
+    }
+}
+
+impl Drop for RustJumpLabel {
+    fn drop(&mut self) {
+        pr_info!("Rust jump label sample (exit)\n");
+    }
+}
diff --git a/samples/rust/rust_jump_label_helper.c b/samples/rust/rust_jump_label_helper.c
new file mode 100644
index 000000000000..a8e6ebb48846
--- /dev/null
+++ b/samples/rust/rust_jump_label_helper.c
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: GPL-2.0
+
+/*
+ * C helper for Rust jump label sample.
+ * Provides a static key that can be used from Rust code.
+ */
+
+#include <linux/module.h>
+#include <linux/jump_label.h>
+
+/*
+ * Define a static key for the Rust module.
+ * This creates the necessary C infrastructure for jump labels.
+ */
+DEFINE_STATIC_KEY_FALSE(rust_sample_key);
+EXPORT_SYMBOL_GPL(rust_sample_key);
+
+MODULE_LICENSE("GPL");
+MODULE_DESCRIPTION("C helper for Rust jump label sample");
+MODULE_AUTHOR("Jan Polensky");
-- 
2.51.0


--67Ci9XMUDsiZd1k+--