Limit number of generated BPF programs

In some cases (e.g. when using the "probe" builtin), the number of
generated BPF functions may be very large which may cause BPFtrace to
hang. This change prevents that by setting a limit on the number of
generated BPF functions. The limit can be changed using the
BPFTRACE_MAX_BPF_PROGS environment variable.

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to
 #### Added
 - Add builtin function: `cgroup_path`
   - [#2055](https://github.com/iovisor/bpftrace/pull/2055)
+- Limit number of generated BPF programs
+  - [#2141](https://github.com/iovisor/bpftrace/pull/2141)
 
 #### Changed
 #### Deprecated

docs/reference_guide.md

@@ -164,6 +164,7 @@ ENVIRONMENT:
     BPFTRACE_NO_CPP_DEMANGLE    [default: 0] disable C++ symbol demangling
     BPFTRACE_MAP_KEYS_MAX       [default: 4096] max keys in a map
     BPFTRACE_MAX_PROBES         [default: 512] max number of probes bpftrace can attach to
+    BPFTRACE_MAX_BPF_PROGS      [default: 512] max number of generated BPF programs
     BPFTRACE_CACHE_USER_SYMBOLS [default: auto] enable user symbol cache
     BPFTRACE_VMLINUX            [default: none] vmlinux path used for kernel symbol resolution
     BPFTRACE_BTF                [default: none] BTF file
@@ -530,6 +531,14 @@ If you're getting a lot of dropped events bpftrace may not be processing events
 fast enough. It may be useful to bump the value higher so more events can be queued up. The tradeoff
 is that bpftrace will use more memory.
 
+### 9.9 `BPFTRACE_MAX_BPF_PROGS`
+
+Default: 512
+
+This is the maximum number of BPF programs (functions) that bpftrace can generate.
+The main purpose of this limit is to prevent bpftrace from hanging since generating a lot of probes
+takes a lot of resources (and it should not happen often).
+
 ## 10. Clang Environment Variables
 
 bpftrace parses header files using libclang, the C interface to Clang. Thus environment variables

man/adoc/bpftrace.adoc

@@ -244,6 +244,14 @@ If you're getting a lot of dropped events bpftrace may not be processing events
 It may be useful to bump the value higher so more events can be queued up.
 The tradeoff is that bpftrace will use more memory.
 
+=== BPFTRACE_MAX_BPF_PROGS
+
+Default: 512
+
+This is the maximum number of BPF programs (functions) that bpftrace can generate.
+The main purpose of this limit is to prevent bpftrace from hanging since generating a lot of probes
+takes a lot of resources (and it should not happen often).
+
 ////
 !
 !

src/ast/passes/codegen_llvm.cpp

@@ -2281,6 +2281,18 @@ void CodegenLLVM::visit(Probe &probe)
         matches = bpftrace_.probe_matcher_->get_matches_for_ap(*attach_point);
       }
 
+      probe_count_ += matches.size();
+      if (probe_count_ > bpftrace_.max_programs_)
+      {
+        throw std::runtime_error(
+            "Your program is trying to generate more than " +
+            std::to_string(probe_count_) +
+            " BPF programs, which exceeds the current limit of " +
+            std::to_string(bpftrace_.max_programs_) +
+            ".\nYou can increase the limit through the BPFTRACE_MAX_BPF_PROGS "
+            "environment variable.");
+      }
+
       tracepoint_struct_ = "";
       for (const auto &m : matches)
       {

src/ast/passes/codegen_llvm.h

@@ -207,6 +207,7 @@ class CodegenLLVM : public Visitor
   AttachPoint *current_attach_point_ = nullptr;
   std::string probefull_;
   std::string tracepoint_struct_;
+  uint64_t probe_count_ = 0;
   std::map<std::string, int> next_probe_index_;
   // Used if there are duplicate USDT entries
   int current_usdt_location_index_{ 0 };

src/bpftrace.h

@@ -168,6 +168,7 @@ class BPFtrace
   uint64_t mapmax_ = 4096;
   size_t cat_bytes_max_ = 10240;
   uint64_t max_probes_ = 512;
+  uint64_t max_programs_ = 512;
   uint64_t log_size_ = 1000000;
   uint64_t perf_rb_pages_ = 64;
   uint64_t max_type_res_iterations = 0;

src/main.cpp

@@ -100,6 +100,7 @@ void usage()
   std::cerr << "    BPFTRACE_MAP_KEYS_MAX       [default: 4096] max keys in a map" << std::endl;
   std::cerr << "    BPFTRACE_CAT_BYTES_MAX      [default: 10k] maximum bytes read by cat builtin" << std::endl;
   std::cerr << "    BPFTRACE_MAX_PROBES         [default: 512] max number of probes" << std::endl;
+  std::cerr << "    BPFTRACE_MAX_BPF_PROGS      [default: 512] max number of generated BPF programs" << std::endl;
   std::cerr << "    BPFTRACE_LOG_SIZE           [default: 1000000] log size in bytes" << std::endl;
   std::cerr << "    BPFTRACE_PERF_RB_PAGES      [default: 64] pages per CPU to allocate for ring buffer" << std::endl;
   std::cerr << "    BPFTRACE_NO_USER_SYMBOLS    [default: 0] disable user symbol resolution" << std::endl;
@@ -266,6 +267,9 @@ static std::optional<struct timespec> get_boottime()
   if (!get_uint64_env_var("BPFTRACE_MAX_PROBES", bpftrace.max_probes_))
     return false;
 
+  if (!get_uint64_env_var("BPFTRACE_MAX_BPF_PROGS", bpftrace.max_programs_))
+    return false;
+
   if (!get_uint64_env_var("BPFTRACE_LOG_SIZE", bpftrace.log_size_))
     return false;
 

tests/runtime/probe

@@ -215,3 +215,11 @@ NAME END_processing_after_exit
 PROG interval:s:1 { exit(); } END { printf("end"); }
 EXPECT end
 TIMEOUT 2
+
+# The "probe" builtin forces bpftrace to generate one BPF function for each
+# match and so we should fail on exceeding BPFTRACE_MAX_BPF_PROGS.
+NAME bpf_programs_limit
+PROG k:* { @[probe] = count(); }
+EXPECT BPFTRACE_MAX_BPF_PROGS
+WILL_FAIL
+TIMEOUT 2