[PATCH 2/2] rust: kunit: add config to fail kunit if TAINT_WARN is set

Triggering a `BUG: sleeping function called from invalid context` does
not mark the unit test as failed. Add a CONFIG_RUST_TAINT_WARN_CHECK to
enable the assertion of TAINT_WARN being set during a unit test.

Signed-off-by: Malte Wechter <[email protected]>
---
 lib/kunit/Kconfig    | 11 +++++++++++
 rust/macros/kunit.rs | 11 +++++++++++
 2 files changed, 22 insertions(+)

diff --git a/lib/kunit/Kconfig b/lib/kunit/Kconfig
index 30bac00c42ce1..4025b5b305549 100644
--- a/lib/kunit/Kconfig
+++ b/lib/kunit/Kconfig
@@ -154,5 +154,16 @@ config RUST_LOCKDEP_KUNIT_DEBUG_LOCKS
 	  this makes it so the KUnit test does not succeed if the test assertions are
 	  true, but a lockdep warning is triggered.
 
+	  If unsure, say N.
+
+config RUST_KUNIT_TAINT_WARN_CHECK
+	bool "Enable extra taint assertion in KUnit tests"
+	depends on RUST
+	default n
+	help
+	  Adds an extra assertion to each Rust kunit test case that asserts
+	  that the kernel is not tainted. If the kernel becomes tainted with a TAINT_WARN,
+	  Enabling this config marks the test as failed.
+
 	  If unsure, say N.
 endif # KUNIT
diff --git a/rust/macros/kunit.rs b/rust/macros/kunit.rs
index d1cd0349f86f0..90264ef54a3a8 100644
--- a/rust/macros/kunit.rs
+++ b/rust/macros/kunit.rs
@@ -146,6 +146,10 @@ macro_rules! assert_eq {
                 {
                     #[cfg(CONFIG_RUST_LOCKDEP_KUNIT_DEBUG_LOCKS)]
                     let __debug_locks_snapshot = ::kernel::bindings::debug_locks;
+                    #[cfg(CONFIG_RUST_KUNIT_TAINT_WARN_CHECK)]
+                    let __is_tainted_snapshot =
+                        ::kernel::bindings::test_taint(::kernel::bindings::TAINT_WARN);
+
 
                     (*_test).status = ::kernel::bindings::kunit_status_KUNIT_SUCCESS;
                     use ::kernel::kunit::is_test_result_ok;
@@ -157,6 +161,13 @@ macro_rules! assert_eq {
                             ::kernel::bindings::debug_locks == __debug_locks_snapshot;
                         assert!(__debug_locks_ok);
                     }
+                    #[cfg(CONFIG_RUST_KUNIT_TAINT_WARN_CHECK)]
+                    {
+                        let __is_tainted =
+                            ::kernel::bindings::test_taint(::kernel::bindings::TAINT_WARN)
+                            == __is_tainted_snapshot;
+                        assert!(__is_tainted);
+                    }
                 }
             }
         });

-- 
2.51.2