From cbc0896091506bc2eaf2db06fc249eb0a87bbb84 Mon Sep 17 00:00:00 2001
From: Viktor Malik <viktor.malik@gmail.com>
Date: Wed, 5 Jun 2024 12:08:17 +0200
Subject: [PATCH] Add unit tests for wildcarded kprobe:mod:func

When the kprobe:mod:func syntax is used with wildcards in either 'mod'
or 'func', kprobe_multi expansion should not be used even if it is
available on the system. The reason is that kprobe_multi doesn't support
specifying the kernel module.

Add unit tests checking that full expansion is always done, even if
kprobe_multi is enabled.
---
 tests/bpftrace.cpp | 61 ++++++++++++++++++++++++++++++++++++++++++++++
 tests/mocks.cpp    | 10 ++++++++
 2 files changed, 71 insertions(+)

diff --git a/tests/bpftrace.cpp b/tests/bpftrace.cpp
index 7dab9632667..5faad3d4098 100644
--- a/tests/bpftrace.cpp
+++ b/tests/bpftrace.cpp
@@ -361,6 +361,67 @@ TEST(bpftrace, add_probes_specify_kernel_module)
                "kernel_mod");
 }
 
+TEST(bpftrace, add_probes_kernel_module_wildcard)
+{
+  auto bpftrace = get_strict_mock_bpftrace();
+  // We enable kprobe_multi here but it doesn't support the module:function
+  // syntax so full expansion should be done anyways.
+  bpftrace->feature_ = std::make_unique<MockBPFfeature>(true);
+  EXPECT_CALL(*bpftrace->mock_probe_matcher,
+              get_symbols_from_traceable_funcs(true))
+      .Times(1);
+
+  parse_probe("kprobe:*kernel_mod:* {}", *bpftrace);
+
+  ASSERT_EQ(3U, bpftrace->get_probes().size());
+  ASSERT_EQ(0U, bpftrace->get_special_probes().size());
+
+  std::string probe_orig_name = "kprobe:*kernel_mod:*";
+  check_kprobe(bpftrace->get_probes().at(0),
+               "func_in_mod",
+               probe_orig_name,
+               0,
+               "kernel_mod");
+  check_kprobe(bpftrace->get_probes().at(1),
+               "other_func_in_mod",
+               probe_orig_name,
+               0,
+               "kernel_mod");
+  check_kprobe(bpftrace->get_probes().at(2),
+               "func_in_mod",
+               probe_orig_name,
+               0,
+               "other_kernel_mod");
+}
+
+TEST(bpftrace, add_probes_kernel_module_function_wildcard)
+{
+  auto bpftrace = get_strict_mock_bpftrace();
+  // We enable kprobe_multi here but it doesn't support the module:function
+  // syntax so full expansion should be done anyways.
+  bpftrace->feature_ = std::make_unique<MockBPFfeature>(true);
+  EXPECT_CALL(*bpftrace->mock_probe_matcher,
+              get_symbols_from_traceable_funcs(true))
+      .Times(1);
+
+  parse_probe("kprobe:kernel_mod:*func_in_mod {}", *bpftrace);
+
+  ASSERT_EQ(2U, bpftrace->get_probes().size());
+  ASSERT_EQ(0U, bpftrace->get_special_probes().size());
+
+  std::string probe_orig_name = "kprobe:kernel_mod:*func_in_mod";
+  check_kprobe(bpftrace->get_probes().at(0),
+               "func_in_mod",
+               probe_orig_name,
+               0,
+               "kernel_mod");
+  check_kprobe(bpftrace->get_probes().at(1),
+               "other_func_in_mod",
+               probe_orig_name,
+               0,
+               "kernel_mod");
+}
+
 TEST(bpftrace, add_probes_offset)
 {
   auto offset = 10;
diff --git a/tests/mocks.cpp b/tests/mocks.cpp
index b2855e8ce58..7793b406236 100644
--- a/tests/mocks.cpp
+++ b/tests/mocks.cpp
@@ -23,6 +23,16 @@ void setup_mock_probe_matcher(MockProbeMatcher &matcher)
         return myval;
       });
 
+  ON_CALL(matcher, get_symbols_from_traceable_funcs(true))
+      .WillByDefault([](void) {
+        std::string ksyms = "kernel_mod:func_in_mod\n"
+                            "kernel_mod:other_func_in_mod\n"
+                            "other_kernel_mod:func_in_mod\n";
+        auto myval = std::unique_ptr<std::istream>(
+            new std::istringstream(ksyms));
+        return myval;
+      });
+
   ON_CALL(matcher, get_symbols_from_file(tracefs::available_events()))
       .WillByDefault([](const std::string &) {
         std::string tracepoints = "sched:sched_one\n"