Support accessing up to 255 USDT probe args (bpftrace#4118)

Support arg0-arg255 for usdt

Signed-off-by: gray <greyschwinger@gmail.com>

Author: jschwinger233

CHANGELOG.md

@@ -39,6 +39,8 @@ and this project adheres to
   - [#4037](https://github.com/bpftrace/bpftrace/pull/4037)
 - Add warning when unset or empty positional parameters are used
   - [#4095](https://github.com/bpftrace/bpftrace/pull/4095)
+- Support accessing up to 255 USDT probe args
+  - [#4118](https://github.com/bpftrace/bpftrace/pull/4118)
 #### Changed
 - `-p` CLI flag now applies to all probes (except BEGIN/END)
   - [#3800](https://github.com/bpftrace/bpftrace/pull/3800)

src/ast/ast.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <charconv>
 #include <cstddef>
 #include <cstdint>
 #include <string>
@@ -340,11 +341,24 @@ class Builtin : public Node {
     return builtin_type;
   }
 
-  // Check if the builtin is 'arg0' - 'arg9'
+  // Check if the builtin is 'arg0' - 'arg255'
   bool is_argx() const
   {
-    return !ident.compare(0, 3, "arg") && ident.size() == 4 &&
-           ident.at(3) >= '0' && ident.at(3) <= '9';
+    if (ident.size() < 4 || ident.size() > 6 || !ident.starts_with("arg"))
+      return false;
+
+    std::string_view num_part = ident.substr(3);
+
+    // no leading zeros
+    if (num_part.size() > 1 && num_part.front() == '0')
+      return false;
+
+    int arg_num = 0;
+    auto [ptr, ec] = std::from_chars(num_part.data(),
+                                     num_part.data() + num_part.size(),
+                                     arg_num);
+    return ec == std::errc() && ptr == num_part.data() + num_part.size() &&
+           arg_num >= 0 && arg_num < 256;
   }
 
   std::string ident;

src/ast/irbuilderbpf.cpp

@@ -11,6 +11,7 @@
 #include "bpftrace.h"
 #include "globalvars.h"
 #include "log.h"
+#include "util/exceptions.h"
 
 namespace libbpf {
 #include "libbpf/bpf.h"
@@ -1491,9 +1492,8 @@ Value *IRBuilderBPF::CreateUSDTReadArgument(Value *ctx,
   }
 
   if (usdt == nullptr) {
-    builtin.addError() << "failed to initialize usdt context for probe "
-                       << attach_point->target;
-    exit(-1);
+    throw util::FatalUserException(
+        "failed to initialize usdt context for probe " + attach_point->target);
   }
 
   std::string ns = attach_point->usdt.provider;
@@ -1505,10 +1505,10 @@ Value *IRBuilderBPF::CreateUSDTReadArgument(Value *ctx,
                             usdt_location_index,
                             arg_num,
                             &argument) != 0) {
-    builtin.addError() << "couldn't get argument " << arg_num << " for "
-                       << attach_point->target << ":" << attach_point->ns << ":"
-                       << attach_point->func;
-    exit(-2);
+    throw util::FatalUserException("couldn't get argument " +
+                                   std::to_string(arg_num) + " for " +
+                                   attach_point->target + ":" +
+                                   attach_point->ns + ":" + attach_point->func);
   }
 
   Value *result = CreateUSDTReadArgument(ctx, &argument, builtin, as, loc);

src/ast/passes/semantic_analyser.cpp

@@ -950,18 +950,19 @@ void SemanticAnalyser::visit(Builtin &builtin)
       return;
     ProbeType pt = probetype(probe->attach_points[0]->provider);
     AddrSpace addrspace = find_addrspace(pt);
+    int arg_num = atoi(builtin.ident.substr(3).c_str());
     for (auto *attach_point : probe->attach_points) {
       ProbeType type = probetype(attach_point->provider);
       if (type != ProbeType::kprobe && type != ProbeType::uprobe &&
           type != ProbeType::usdt && type != ProbeType::rawtracepoint)
         builtin.addError() << "The " << builtin.ident
                            << " builtin can only be used with "
                            << "'kprobes', 'uprobes' and 'usdt' probes";
+      // argx in USDT probes doesn't need to check against arch::max_arg()
+      if (type != ProbeType::usdt && arg_num > arch::max_arg())
+        builtin.addError() << arch::name() << " doesn't support "
+                           << builtin.ident;
     }
-    int arg_num = atoi(builtin.ident.substr(3).c_str());
-    if (arg_num > arch::max_arg())
-      builtin.addError() << arch::name() << " doesn't support "
-                         << builtin.ident;
     builtin.builtin_type = CreateUInt64();
     builtin.builtin_type.SetAS(addrspace);
   } else if (!builtin.ident.compare(0, 4, "sarg") &&

src/lexer.l

@@ -44,7 +44,7 @@ hspace   [ \t]
 vspace   [\n\r]
 space    {hspace}|{vspace}
 path     :(\\.|[_\-\./a-zA-Z0-9#$+\*])+
-builtin  arg[0-9]|args|cgroup|comm|cpid|numaid|cpu|ncpus|ctx|curtask|elapsed|func|gid|pid|probe|rand|retval|sarg[0-9]|tid|uid|username|jiffies|nsecs|kstack|ustack
+builtin  arg[0-9]+|args|cgroup|comm|cpid|numaid|cpu|ncpus|ctx|curtask|elapsed|func|gid|pid|probe|rand|retval|sarg[0-9]|tid|uid|username|jiffies|nsecs|kstack|ustack
 
 int_type        bool|(u)?int(8|16|32|64)
 builtin_type    void|(u)?(min|max|sum|count|avg|stats)_t|probe_t|username|lhist_t|hist_t|usym_t|ksym_t|timestamp|macaddr_t|cgroup_path_t|strerror_t|kstack_t|ustack_t

tests/runtime/usdt

@@ -132,7 +132,7 @@ EXPECT Attaching 3 probes...
 
 NAME usdt probes - attach to probe on multiple files by wildcard
 PROG usdt:./testprogs/usdt*::* { printf("here\n" ); exit(); }
-EXPECT Attaching 48 probes...
+EXPECT Attaching 49 probes...
 
 NAME usdt probes - attach to probe on multiple providers by wildcard and pid
 RUN {{BPFTRACE}} -e 'usdt:::*probe2 { printf("here\n" ); exit(); }' -p {{BEFORE_PID}}
@@ -287,3 +287,8 @@ RUN {{BPFTRACE}} runtime/scripts/usdt_multi_modules.bt
 EXPECT Attaching 2 probes...
 TIMEOUT 1
 REQUIRES ./testprogs/systemtap_sys_sdt_check
+
+NAME usdt multi arguments
+RUN {{BPFTRACE}} -e 'usdt:./testprogs/usdt_multi_args:usdt_multi_args:probe1 { printf("%llx %llx %llx %llx %llx %llx %llx %llx %llx %llx %lld %llx\n", arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11); exit(); }' -c ./testprogs/usdt_multi_args
+EXPECT deadbeef 1 ffffffffffffffff cafebabe 8badf00d feedface 123456789abcdef 0 7fffffffffffffff 5555555555555555 42 aaaaaaaaaaaaaaaa
+REQUIRES ./testprogs/systemtap_sys_sdt_check

tests/testprogs/usdt_multi_args.c

@@ -0,0 +1,29 @@
+#ifdef HAVE_SYSTEMTAP_SYS_SDT_H
+#include <sys/sdt.h>
+#else
+#define DTRACE_PROBE1(a, b, c) (void)0
+#endif
+#include <stdint.h>
+
+int main()
+{
+  uint32_t a = 0xdeadbeef;
+  uint32_t b = 1;
+  uint64_t c = UINT64_MAX;
+  uint32_t d = 0xcafebabe;
+  uint32_t e = 0x8badf00d;
+  uint32_t f = 0xfeedface;
+  uint64_t g = 0x0123456789abcdefULL;
+  uint64_t h = 0x0ULL;
+  uint64_t i = 0x7fffffffffffffffULL;
+  uint64_t j = 0x5555555555555555ULL;
+  uint32_t k = 42;
+  uint64_t l = 0xaaaaaaaaaaaaaaaaULL;
+
+  while (1) {
+    DTRACE_PROBE12(usdt_multi_args, probe1, a, b, c, d, e, f, g, h, i, j, k, l);
+  }
+
+  return 0;
+}
+