HACK: arm64: Enable rust modules at EL2

Just a bit of "fun": hack the minimal rust module to print a message
at EL2 (via the pre-existing pl011 driver) instead of adding numbers
together.

Signed-off-by: Will Deacon <willdeacon@google.com>
diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 460761d..3e94ded 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -185,6 +185,7 @@
 	select HAVE_DEBUG_KMEMLEAK
 	select HAVE_DMA_CONTIGUOUS
 	select HAVE_DYNAMIC_FTRACE
+	select HAVE_RUST
 	select FTRACE_MCOUNT_USE_PATCHABLE_FUNCTION_ENTRY \
 		if DYNAMIC_FTRACE_WITH_REGS
 	select HAVE_EFFICIENT_UNALIGNED_ACCESS
diff --git a/arch/arm64/kvm/hyp/nvhe/Makefile.nvhe b/arch/arm64/kvm/hyp/nvhe/Makefile.nvhe
index 381e610..6c0184b 100644
--- a/arch/arm64/kvm/hyp/nvhe/Makefile.nvhe
+++ b/arch/arm64/kvm/hyp/nvhe/Makefile.nvhe
@@ -22,15 +22,17 @@
 ## file containing all nVHE hyp code and data.
 ##
 
-hyp-obj := $(patsubst %.o,%.nvhe.o,$(hyp-obj-y))
+hyp-obj := $(patsubst %.o,%_nvhe.o,$(hyp-obj-y))
 targets += $(hyp-obj) kvm_nvhe.tmp.o kvm_nvhe.rel.o hyp.lds hyp-reloc.S hyp-reloc.o
 
 # 1) Compile all source files to `.nvhe.o` object files. The file extension
 #    avoids file name clashes for files shared with VHE.
-$(obj)/%.nvhe.o: $(src)/%.c FORCE
+$(obj)/%_nvhe.o: $(src)/%.c FORCE
 	$(call if_changed_rule,cc_o_c)
-$(obj)/%.nvhe.o: $(src)/%.S FORCE
+$(obj)/%_nvhe.o: $(src)/%.S FORCE
 	$(call if_changed_rule,as_o_S)
+$(obj)/%_nvhe.o: $(src)/%.rs FORCE
+	$(call if_changed_dep,rustc_o_rs)
 
 # 2) Partially link all '.nvhe.o' files and apply the linker script.
 #    Prefixes names of ELF sections with '.hyp', eg. '.hyp.text'.
diff --git a/rust/Makefile b/rust/Makefile
index 7700d38..906a4c9 100644
--- a/rust/Makefile
+++ b/rust/Makefile
@@ -253,6 +253,7 @@
 
 # Derived from `scripts/Makefile.clang`.
 BINDGEN_TARGET_x86	:= x86_64-linux-gnu
+BINDGEN_TARGET_arm64	:= aarch64-linux-gnu
 BINDGEN_TARGET		:= $(BINDGEN_TARGET_$(SRCARCH))
 
 # All warnings are inhibited since GCC builds are very experimental,
diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index c48bc28..3b81175 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -7,6 +7,7 @@
  */
 
 #include <linux/slab.h>
+#include <asm/kvm_pkvm_module.h>
 
 /* `bindgen` gets confused at certain things. */
 const gfp_t BINDINGS_GFP_KERNEL = GFP_KERNEL;
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index abd4626..b9525cc 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -23,6 +23,7 @@
 #[cfg(not(testlib))]
 mod allocator;
 pub mod error;
+pub mod pkvm;
 pub mod prelude;
 pub mod print;
 pub mod str;
diff --git a/rust/kernel/pkvm.rs b/rust/kernel/pkvm.rs
new file mode 100644
index 0000000..3febe37
--- /dev/null
+++ b/rust/kernel/pkvm.rs
@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! pKVM hyp module interfaces
+
+use crate::error::{code::ENOMEM, Result};
+
+/// Module ops structure passed to initialisation function at EL2.
+pub type El2ModuleOps = *const bindings::pkvm_module_ops;
+
+type El2InitFn = unsafe extern "C" fn(El2ModuleOps) -> core::ffi::c_int;
+
+/// Load a hypervisor module at EL2.
+pub fn load_el2_module(
+	module: &'static crate::ThisModule,
+	init_fn: El2InitFn,
+	token: *mut u64,
+) -> Result {
+
+	unsafe {
+		(*module.0).arch.hyp.init = Some(init_fn);
+		let rc = bindings::__pkvm_load_el2_module(module.0, token);
+
+		if rc != 0 {
+			return Err(ENOMEM);
+		}
+	}
+
+	Ok(())
+}
diff --git a/samples/rust/Makefile b/samples/rust/Makefile
index 1daba5f..f4307d9 100644
--- a/samples/rust/Makefile
+++ b/samples/rust/Makefile
@@ -1,5 +1,14 @@
 # SPDX-License-Identifier: GPL-2.0
 
+$(obj)/hyp/kvm_nvhe.o: FORCE
+	$(Q)$(MAKE) $(build)=$(obj)/hyp $(obj)/hyp/kvm_nvhe.o
+
+clean-files := hyp/hyp.lds hyp/hyp-reloc.S
+
 obj-$(CONFIG_SAMPLE_RUST_MINIMAL)		+= rust_minimal.o
 
+RUSTFLAGS_rust_minimal_host.o += -A improper-ctypes # For __uint_128_t
+
+rust_minimal-y := rust_minimal_host.o hyp/kvm_nvhe.o
+
 subdir-$(CONFIG_SAMPLE_RUST_HOSTPROGS)		+= hostprogs
diff --git a/samples/rust/hyp/Makefile b/samples/rust/hyp/Makefile
new file mode 100644
index 0000000..a14d858
--- /dev/null
+++ b/samples/rust/hyp/Makefile
@@ -0,0 +1,2 @@
+hyp-obj-y := rust_minimal_hyp.o
+include $(srctree)/arch/arm64/kvm/hyp/nvhe/Makefile.module
diff --git a/samples/rust/hyp/rust_minimal_hyp.rs b/samples/rust/hyp/rust_minimal_hyp.rs
new file mode 100644
index 0000000..e220b46
--- /dev/null
+++ b/samples/rust/hyp/rust_minimal_hyp.rs
@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Simple EL2 module written in rust. This is the hyp part.
+
+/// EL2 initialisation function called by the hypervisor.
+#[no_mangle]
+pub extern "C" fn init(ops: kernel::pkvm::El2ModuleOps) -> core::ffi::c_int {
+	let msg = b"oh no, EL2 is rusting away!\n\0";
+
+	unsafe {
+		let puts = (*ops).puts;
+
+		match puts {
+			Some(func) => func(msg.as_ptr() as _),
+			None => (),
+		}
+	}
+	0
+}
diff --git a/samples/rust/rust_minimal.rs b/samples/rust/rust_minimal.rs
deleted file mode 100644
index 54ad176..0000000
--- a/samples/rust/rust_minimal.rs
+++ /dev/null
@@ -1,38 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0
-
-//! Rust minimal sample.
-
-use kernel::prelude::*;
-
-module! {
-    type: RustMinimal,
-    name: b"rust_minimal",
-    author: b"Rust for Linux Contributors",
-    description: b"Rust minimal sample",
-    license: b"GPL",
-}
-
-struct RustMinimal {
-    numbers: Vec<i32>,
-}
-
-impl kernel::Module for RustMinimal {
-    fn init(_module: &'static ThisModule) -> Result<Self> {
-        pr_info!("Rust minimal sample (init)\n");
-        pr_info!("Am I built-in? {}\n", !cfg!(MODULE));
-
-        let mut numbers = Vec::new();
-        numbers.try_push(72)?;
-        numbers.try_push(108)?;
-        numbers.try_push(200)?;
-
-        Ok(RustMinimal { numbers })
-    }
-}
-
-impl Drop for RustMinimal {
-    fn drop(&mut self) {
-        pr_info!("My numbers are {:?}\n", self.numbers);
-        pr_info!("Rust minimal sample (exit)\n");
-    }
-}
diff --git a/samples/rust/rust_minimal_host.rs b/samples/rust/rust_minimal_host.rs
new file mode 100644
index 0000000..9f37af14
--- /dev/null
+++ b/samples/rust/rust_minimal_host.rs
@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Rust minimal sample.
+
+use kernel::prelude::*;
+use kernel::pkvm;
+
+use core::ptr;
+
+module! {
+    type: RustMinimal,
+    name: b"rust_minimal",
+    author: b"Rust for Linux Contributors",
+    description: b"Rust minimal sample",
+    license: b"GPL",
+}
+
+struct RustMinimal;
+
+extern "C" {
+	fn __kvm_nvhe_init(ops: pkvm::El2ModuleOps) -> core::ffi::c_int;
+}
+
+impl kernel::Module for RustMinimal {
+    fn init(module: &'static ThisModule) -> Result<Self> {
+	pkvm::load_el2_module(module, __kvm_nvhe_init, ptr::null_mut())?;
+        Ok(RustMinimal)
+    }
+}
diff --git a/scripts/generate_rust_target.rs b/scripts/generate_rust_target.rs
index 3c6cbe2..5eaf6a5 100644
--- a/scripts/generate_rust_target.rs
+++ b/scripts/generate_rust_target.rs
@@ -161,6 +161,17 @@
         ts.push("features", features);
         ts.push("llvm-target", "x86_64-linux-gnu");
         ts.push("target-pointer-width", "64");
+    } else if cfg.has("ARM64") {
+        ts.push("arch", "aarch64");
+        ts.push(
+            "data-layout",
+            "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128",
+        );
+        ts.push("disable-redzone", true);
+        ts.push("features", "+strict-align,-neon,-fp-armv8");
+        ts.push("llvm-target", "aarch64-linux-gnu");
+        ts.push("max-atomic-width", 128);
+        ts.push("target-pointer-width", "64");
     } else {
         panic!("Unsupported architecture");
     }