Prevent banned kernel funcs from all probes

Previously we only banned kretprobes from
using banned kernel functions but a bpftrace
script at Meta saw a crash from utilizing
one of these functions in kfunc/kretfunc.

This changes prevents any probe from attaching
to one of these functions.

Issue: bpftrace#3144

src/attached_probe.cpp

@@ -37,7 +37,7 @@ namespace bpftrace {
  * Kernel functions that are unsafe to trace are excluded in the Kernel with
  * `notrace`. However, the ones below are not excluded.
  */
-const std::set<std::string> banned_kretprobes = {
+const std::set<std::string> banned_kernel_funcs = {
   "_raw_spin_lock",
   "_raw_spin_lock_irqsave",
   "_raw_spin_unlock_irqrestore",
@@ -104,10 +104,10 @@ std::string progtypeName(libbpf::bpf_prog_type t)
   }
 }
 
-void check_banned_kretprobes(std::string const &kprobe_name)
+void check_banned_kernel_funcs(std::string const &kernel_func)
 {
-  if (banned_kretprobes.find(kprobe_name) != banned_kretprobes.end()) {
-    LOG(FATAL) << "kretprobe:" << kprobe_name
+  if (banned_kernel_funcs.find(kernel_func) != banned_kernel_funcs.end()) {
+    LOG(FATAL) << "kernel function: " << kernel_func
                << " can't be used as it might lock up your system.";
   }
 }
@@ -178,7 +178,13 @@ AttachedProbe::AttachedProbe(Probe &probe,
     : probe_(probe), prog_(std::move(prog)), btf_(btf)
 {
   load_prog(feature);
+  if (probe_.type == ProbeType::kretprobe || probe_.type == ProbeType::kprobe ||
+      probe_.type == ProbeType::kfunc || probe_.type == ProbeType::kretfunc) {
+    check_banned_kernel_funcs(probe_.attach_point);
+  }
+
   LOG(V1) << "Attaching " << probe_.orig_name;
+
   switch (probe_.type) {
     case ProbeType::special:
       // If BPF_PROG_TYPE_RAW_TRACEPOINT is available, no need to attach prog
@@ -190,7 +196,6 @@ AttachedProbe::AttachedProbe(Probe &probe,
       attach_kprobe(safe_mode);
       break;
     case ProbeType::kretprobe:
-      check_banned_kretprobes(probe_.attach_point);
       attach_kprobe(safe_mode);
       break;
     case ProbeType::tracepoint:

tests/runtime/banned_probes

@@ -1,23 +1,45 @@
+# Check the list of banned functions
+
 NAME kretprobe:_raw_spin_lock is banned
 PROG kretprobe:_raw_spin_lock { exit(); }
-EXPECT ERROR: kretprobe:_raw_spin_lock can't be used as it might lock up your system.
+EXPECT ERROR: kernel function: _raw_spin_lock can't be used as it might lock up your system.
 TIMEOUT 1
 WILL_FAIL
 
 NAME kretprobe:queued_spin_lock_slowpath is banned
 PROG kretprobe:queued_spin_lock_slowpath { exit(); }
-EXPECT ERROR: kretprobe:queued_spin_lock_slowpath can't be used as it might lock up your system.
+EXPECT ERROR: kernel function: queued_spin_lock_slowpath can't be used as it might lock up your system.
 TIMEOUT 1
 WILL_FAIL
 
 NAME kretprobe:_raw_spin_unlock_irqrestore is banned
 PROG kretprobe:_raw_spin_unlock_irqrestore { exit(); }
-EXPECT ERROR: kretprobe:_raw_spin_unlock_irqrestore can't be used as it might lock up your system.
+EXPECT ERROR: kernel function: _raw_spin_unlock_irqrestore can't be used as it might lock up your system.
 TIMEOUT 1
 WILL_FAIL
 
 NAME kretprobe:_raw_spin_lock_irqsave is banned
 PROG kretprobe:_raw_spin_lock_irqsave { exit(); }
-EXPECT ERROR: kretprobe:_raw_spin_lock_irqsave can't be used as it might lock up your system.
+EXPECT ERROR: kernel function: _raw_spin_lock_irqsave can't be used as it might lock up your system.
+TIMEOUT 1
+WILL_FAIL
+
+# Check the list of banned probes for those functions
+
+NAME kprobe:_raw_spin_lock is banned
+PROG kprobe:_raw_spin_lock { exit(); }
+EXPECT ERROR: kernel function: _raw_spin_lock can't be used as it might lock up your system.
+TIMEOUT 1
+WILL_FAIL
+
+NAME kfunc:queued_spin_lock_slowpath is banned
+PROG kfunc:queued_spin_lock_slowpath { exit(); }
+EXPECT ERROR: kernel function: queued_spin_lock_slowpath can't be used as it might lock up your system.
+TIMEOUT 1
+WILL_FAIL
+
+NAME kretfunc:_raw_spin_unlock_irqrestore is banned
+PROG kretfunc:_raw_spin_unlock_irqrestore { exit(); }
+EXPECT ERROR: kernel function: _raw_spin_unlock_irqrestore can't be used as it might lock up your system.
 TIMEOUT 1
 WILL_FAIL