* Author: Andrey Ryabinin <a.ryabinin@samsung.com>
  */
 
+#define pr_fmt(fmt) "kasan_test: " fmt
+
+#include <kunit/test.h>
 #include <linux/bitops.h>
 #include <linux/delay.h>
+#include <linux/io.h>
 #include <linux/kasan.h>
 #include <linux/kernel.h>
 #include <linux/mm.h>
 #include <linux/module.h>
 #include <linux/printk.h>
 #include <linux/random.h>
+#include <linux/set_memory.h>
 #include <linux/slab.h>
 #include <linux/string.h>
+#include <linux/tracepoint.h>
 #include <linux/uaccess.h>
-#include <linux/io.h>
 #include <linux/vmalloc.h>
-#include <linux/set_memory.h>
+#include <trace/events/printk.h>
 
 #include <asm/page.h>
 
-#include <kunit/test.h>
-
 #include "kasan.h"
 
 #define OOB_TAG_OFF (IS_ENABLED(CONFIG_KASAN_GENERIC) ? 0 : KASAN_GRANULE_SIZE)
 
+static bool multishot;
+
+/* Fields set based on lines observed in the console. */
+static struct {
+       bool report_found;
+       bool async_fault;
+} test_status;
+
 /*
  * Some tests use these global variables to store return values from function
  * calls that could otherwise be eliminated by the compiler as dead code.
 void *kasan_ptr_result;
 int kasan_int_result;
 
-static struct kunit_resource resource;
-static struct kunit_kasan_status test_status;
-static bool multishot;
+/* Probe for console output: obtains test_status lines of interest. */
+static void probe_console(void *ignore, const char *buf, size_t len)
+{
+       if (strnstr(buf, "BUG: KASAN: ", len))
+               WRITE_ONCE(test_status.report_found, true);
+       else if (strnstr(buf, "Asynchronous fault: ", len))
+               WRITE_ONCE(test_status.async_fault, true);
+}
 
-/*
- * Temporarily enable multi-shot mode. Otherwise, KASAN would only report the
- * first detected bug and panic the kernel if panic_on_warn is enabled. For
- * hardware tag-based KASAN also allow tag checking to be reenabled for each
- * test, see the comment for KUNIT_EXPECT_KASAN_FAIL().
- */
-static int kasan_test_init(struct kunit *test)
+static void register_tracepoints(struct tracepoint *tp, void *ignore)
+{
+       check_trace_callback_type_console(probe_console);
+       if (!strcmp(tp->name, "console"))
+               WARN_ON(tracepoint_probe_register(tp, probe_console, NULL));
+}
+
+static void unregister_tracepoints(struct tracepoint *tp, void *ignore)
+{
+       if (!strcmp(tp->name, "console"))
+               tracepoint_probe_unregister(tp, probe_console, NULL);
+}
+
+static int kasan_suite_init(struct kunit_suite *suite)
 {
        if (!kasan_enabled()) {
-               kunit_err(test, "can't run KASAN tests with KASAN disabled");
+               pr_err("Can't run KASAN tests with KASAN disabled");
                return -1;
        }
 
+       /*
+        * Temporarily enable multi-shot mode. Otherwise, KASAN would only
+        * report the first detected bug and panic the kernel if panic_on_warn
+        * is enabled.
+        */
        multishot = kasan_save_enable_multi_shot();
-       test_status.report_found = false;
-       test_status.sync_fault = false;
-       kunit_add_named_resource(test, NULL, NULL, &resource,
-                                       "kasan_status", &test_status);
+
+       /*
+        * Because we want to be able to build the test as a module, we need to
+        * iterate through all known tracepoints, since the static registration
+        * won't work here.
+        */
+       for_each_kernel_tracepoint(register_tracepoints, NULL);
        return 0;
 }
 
-static void kasan_test_exit(struct kunit *test)
+static void kasan_suite_exit(struct kunit_suite *suite)
 {
        kasan_restore_multi_shot(multishot);
-       KUNIT_EXPECT_FALSE(test, test_status.report_found);
+       for_each_kernel_tracepoint(unregister_tracepoints, NULL);
+       tracepoint_synchronize_unregister();
+}
+
+static void kasan_test_exit(struct kunit *test)
+{
+       KUNIT_EXPECT_FALSE(test, READ_ONCE(test_status.report_found));
 }
 
 /**
        if (IS_ENABLED(CONFIG_KASAN_HW_TAGS) &&                         \
            kasan_sync_fault_possible()) {                              \
                if (READ_ONCE(test_status.report_found) &&              \
-                   READ_ONCE(test_status.sync_fault))                  \
+                   !READ_ONCE(test_status.async_fault))                \
                        kasan_enable_tagging();                         \
                migrate_enable();                                       \
        }                                                               \
        WRITE_ONCE(test_status.report_found, false);                    \
+       WRITE_ONCE(test_status.async_fault, false);                     \
 } while (0)
 
 #define KASAN_TEST_NEEDS_CONFIG_ON(test, config) do {                  \
 
 static struct kunit_suite kasan_kunit_test_suite = {
        .name = "kasan",
-       .init = kasan_test_init,
        .test_cases = kasan_kunit_test_cases,
        .exit = kasan_test_exit,
+       .suite_init = kasan_suite_init,
+       .suite_exit = kasan_suite_exit,
 };
 
 kunit_test_suite(kasan_kunit_test_suite);
 
 
 #include <asm/sections.h>
 
-#include <kunit/test.h>
-
 #include "kasan.h"
 #include "../slab.h"
 
 
 #endif
 
-#if IS_ENABLED(CONFIG_KASAN_KUNIT_TEST)
-static void update_kunit_status(bool sync)
-{
-       struct kunit *test;
-       struct kunit_resource *resource;
-       struct kunit_kasan_status *status;
-
-       test = current->kunit_test;
-       if (!test)
-               return;
-
-       resource = kunit_find_named_resource(test, "kasan_status");
-       if (!resource) {
-               kunit_set_failure(test);
-               return;
-       }
-
-       status = (struct kunit_kasan_status *)resource->data;
-       WRITE_ONCE(status->report_found, true);
-       WRITE_ONCE(status->sync_fault, sync);
-
-       kunit_put_resource(resource);
-}
-#else
-static void update_kunit_status(bool sync) { }
-#endif
-
 static DEFINE_SPINLOCK(report_lock);
 
 static void start_report(unsigned long *flags, bool sync)
 {
        /* Respect the /proc/sys/kernel/traceoff_on_warning interface. */
        disable_trace_on_warning();
-       /* Update status of the currently running KASAN test. */
-       update_kunit_status(sync);
        /* Do not allow LOCKDEP mangling KASAN reports. */
        lockdep_off();
        /* Make sure we don't end up in loop. */