Enable printf, cat, system to have more args

This enables printf, cat and system to have more than 7 arguments. It can be more user-friendly to allow any reasonable number of arguments for these builtin calls which use formatted strings. Resolves #1336.
bpftrace · Jul 8, 2020 · f2cdf87 · f2cdf87
1 parent 24205c9
commit f2cdf87
Show file tree

Hide file tree

Showing 7 changed files with 89 additions and 81 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -50,6 +50,8 @@ and this project adheres to
   - [#1377](https://github.com/iovisor/bpftrace/pull/1377)
 - Add support for non-map print()
   - [#1381](https://github.com/iovisor/bpftrace/pull/1381)
+- Enable `printf`, `cat` and `system` to have more than 7 arguments
+  - [#1404](https://github.com/iovisor/bpftrace/pull/1404)
 - Enable the `ternary` operator to evaluate builtin calls
   - [#1405](https://github.com/iovisor/bpftrace/pull/1405)
 

diff --git a/src/ast/semantic_analyser.cpp b/src/ast/semantic_analyser.cpp
@@ -766,7 +766,7 @@ void SemanticAnalyser::visit(Call &call)
   else if (call.func == "printf" || call.func == "system" || call.func == "cat")
   {
     check_assignment(call, false, false, false);
-    if (check_varargs(call, 1, 7))
+    if (check_varargs(call, 1, 128))
     {
       check_arg(call, Type::string, 0, true);
       if (is_final_pass())

diff --git a/src/bpftrace.cpp b/src/bpftrace.cpp
@@ -48,53 +48,69 @@ namespace bpftrace {
 DebugLevel bt_debug = DebugLevel::kNone;
 bool bt_verbose = false;
 volatile sig_atomic_t BPFtrace::exitsig_recv = false;
+const int FMT_BUF_SZ = 512;
 
-int format(char * s, size_t n, const char * fmt, std::vector<std::unique_ptr<IPrintable>> &args) {
+std::string format(std::string fmt,
+                   std::vector<std::unique_ptr<IPrintable>> &args)
+{
+  std::string retstr;
+  auto buffer = std::vector<char>(FMT_BUF_SZ);
+  auto check_snprintf_ret = [](int r) {
+    if (r < 0)
+    {
+      std::cerr << "format() error occurred: " << std::strerror(errno)
+                << std::endl;
+      abort();
+    }
+  };
   // Args have been made safe for printing by now, so replace nonstandard format
   // specifiers with %s
-  std::string str = std::string(fmt);
   size_t start_pos = 0;
-  while ((start_pos = str.find("%r", start_pos)) != std::string::npos)
+  while ((start_pos = fmt.find("%r", start_pos)) != std::string::npos)
   {
-    str.replace(start_pos, 2, "%s");
+    fmt.replace(start_pos, 2, "%s");
     start_pos += 2;
   }
-  fmt = str.c_str();
 
-  int ret = -1;
-  switch(args.size()) {
-    case 0:
-      ret = snprintf(s, n, "%s", fmt);
-      break;
-    case 1:
-      ret = snprintf(s, n, fmt, args.at(0)->value());
-      break;
-    case 2:
-      ret = snprintf(s, n, fmt, args.at(0)->value(), args.at(1)->value());
-      break;
-    case 3:
-      ret = snprintf(s, n, fmt, args.at(0)->value(), args.at(1)->value(), args.at(2)->value());
-      break;
-    case 4:
-      ret = snprintf(s, n, fmt, args.at(0)->value(), args.at(1)->value(), args.at(2)->value(), args.at(3)->value());
-      break;
-    case 5:
-      ret = snprintf(s, n, fmt, args.at(0)->value(), args.at(1)->value(), args.at(2)->value(),
-        args.at(3)->value(), args.at(4)->value());
-      break;
-    case 6:
-      ret = snprintf(s, n, fmt, args.at(0)->value(), args.at(1)->value(), args.at(2)->value(),
-        args.at(3)->value(), args.at(4)->value(), args.at(5)->value());
-      break;
-    default:
-      std::cerr << "format() can only take up to 7 arguments (" << args.size() << ") provided" << std::endl;
-      abort();
-  }
-  if (ret < 0 && errno != 0) {
-    std::cerr << "format() error occurred: " << std::strerror(errno) << std::endl;
-    abort();
-  }
-  return ret;
+  auto tokens_begin = std::sregex_iterator(fmt.begin(),
+                                           fmt.end(),
+                                           format_specifier_re);
+  auto tokens_end = std::sregex_iterator();
+
+  // replace format string tokens with args one by one
+  int literal_text_pos = 0; // starting pos of literal text (text that is not
+                            // format specifier)
+  int i = 0;                // args index
+  while (tokens_begin != tokens_end)
+  {
+    // take out the literal text
+    retstr += fmt.substr(literal_text_pos,
+                         tokens_begin->position() - literal_text_pos);
+    // replace current specifier with an arg
+    int r = snprintf(buffer.data(),
+                     buffer.capacity(),
+                     tokens_begin->str().c_str(),
+                     args.at(i)->value());
+    check_snprintf_ret(r);
+    if (static_cast<size_t>(r) >= buffer.capacity())
+    {
+      // the buffer is not big enough to hold the string, resize it
+      buffer.resize(r + 1);
+      r = snprintf(buffer.data(),
+                   buffer.capacity(),
+                   tokens_begin->str().c_str(),
+                   args.at(i)->value());
+      check_snprintf_ret(r);
+    }
+    retstr += std::string(buffer.data());
+    // move to the next literal text
+    literal_text_pos = tokens_begin->position() + tokens_begin->length();
+    ++tokens_begin;
+    ++i;
+  }
+  // append whatever is left
+  retstr += fmt.substr(literal_text_pos);
+  return retstr;
 }
 
 BPFtrace::~BPFtrace()
@@ -573,66 +589,35 @@ void perf_event_printer(void *cb_cookie, void *data, int size __attribute__((unu
     }
 
     auto id = printf_id - asyncactionint(AsyncAction::syscall);
-    auto fmt = std::get<0>(bpftrace->system_args_[id]).c_str();
+    auto fmt = std::get<0>(bpftrace->system_args_[id]);
     auto args = std::get<1>(bpftrace->system_args_[id]);
     auto arg_values = bpftrace->get_arg_values(args, arg_data);
 
-    const int BUFSIZE = 512;
-    char buffer[BUFSIZE];
-    int sz = format(buffer, BUFSIZE, fmt, arg_values);
-    // Return value is required size EXCLUDING null byte
-    if (sz >= BUFSIZE)
-    {
-      std::cerr << "syscall() command to long (" << sz << " bytes): ";
-      std::cerr << buffer << std::endl;
-      return;
-    }
-    bpftrace->out_->message(MessageType::syscall, exec_system(buffer), false);
+    bpftrace->out_->message(MessageType::syscall,
+                            exec_system(format(fmt, arg_values).c_str()),
+                            false);
     return;
   }
   else if ( printf_id >= asyncactionint(AsyncAction::cat))
   {
     auto id = printf_id - asyncactionint(AsyncAction::cat);
-    auto fmt = std::get<0>(bpftrace->cat_args_[id]).c_str();
+    auto fmt = std::get<0>(bpftrace->cat_args_[id]);
     auto args = std::get<1>(bpftrace->cat_args_[id]);
     auto arg_values = bpftrace->get_arg_values(args, arg_data);
 
-    const int BUFSIZE = 512;
-    char buffer[BUFSIZE];
-    int sz = format(buffer, BUFSIZE, fmt, arg_values);
-    // Return value is required size EXCLUDING null byte
-    if (sz >= BUFSIZE)
-    {
-      std::cerr << "cat() command to long (" << sz << " bytes): ";
-      std::cerr << buffer << std::endl;
-      return;
-    }
     std::stringstream buf;
-    cat_file(buffer, bpftrace->cat_bytes_max_, buf);
+    cat_file(format(fmt, arg_values).c_str(), bpftrace->cat_bytes_max_, buf);
     bpftrace->out_->message(MessageType::cat, buf.str(), false);
 
     return;
   }
 
   // printf
-  auto fmt = std::get<0>(bpftrace->printf_args_[printf_id]).c_str();
+  auto fmt = std::get<0>(bpftrace->printf_args_[printf_id]);
   auto args = std::get<1>(bpftrace->printf_args_[printf_id]);
   auto arg_values = bpftrace->get_arg_values(args, arg_data);
 
-  // First try with a stack buffer, if that fails use a heap buffer
-  const int BUFSIZE=512;
-  char buffer[BUFSIZE];
-  int required_size = format(buffer, BUFSIZE, fmt, arg_values);
-  // Return value is required size EXCLUDING null byte
-  if (required_size < BUFSIZE) {
-    bpftrace->out_->message(MessageType::printf, std::string(buffer), false);
-  } else {
-    auto buf = std::make_unique<char[]>(required_size+1);
-    // if for some reason the size is still wrong the string
-    // will just be silently truncated
-    format(buf.get(), required_size, fmt, arg_values);
-    bpftrace->out_->message(MessageType::printf, std::string(buf.get()), false);
-  }
+  bpftrace->out_->message(MessageType::printf, format(fmt, arg_values), false);
 }
 
 std::vector<std::unique_ptr<IPrintable>> BPFtrace::get_arg_values(const std::vector<Field> &args, uint8_t* arg_data)

diff --git a/src/printf.cpp b/src/printf.cpp
@@ -9,9 +9,10 @@ namespace bpftrace {
 std::string verify_format_string(const std::string &fmt, std::vector<Field> args)
 {
   std::stringstream message;
-  const std::regex re("%-?[0-9]*(\\.[0-9]+)?[a-zA-Z]+");
 
-  auto tokens_begin = std::sregex_iterator(fmt.begin(), fmt.end(), re);
+  auto tokens_begin = std::sregex_iterator(fmt.begin(),
+                                           fmt.end(),
+                                           format_specifier_re);
   auto tokens_end = std::sregex_iterator();
 
   auto num_tokens = std::distance(tokens_begin, tokens_end);

diff --git a/src/printf.h b/src/printf.h
@@ -7,6 +7,8 @@
 
 namespace bpftrace {
 
+const std::regex format_specifier_re("%-?[0-9]*(\\.[0-9]+)?[a-zA-Z]+");
+
 struct Field;
 
 std::string verify_format_string(const std::string& fmt,

diff --git a/tests/runtime/call b/tests/runtime/call
@@ -13,6 +13,11 @@ RUN bpftrace -v -e 'struct Foo { int a; char b[10]; } uprobe:testprogs/uprobe_te
 EXPECT 123 hello 456 world
 TIMEOUT 5
 
+NAME printf_more_arguments
+RUN bpftrace -v -e 'BEGIN { printf("%d: %s; %d: %s;; %d: %s;;; %d: %s;;;;\\n", 1, "a", 2, "ab", 3, "abc", 4, "abcd"); exit(); }'
+EXPECT 1: a; 2: ab;; 3: abc;;; 4: abcd;;;;
+TIMEOUT 5
+
 NAME time
 RUN bpftrace -v -e 'i:ms:1 { time("%H:%M:%S\n"); exit();}'
 EXPECT [0-9]*:[0-9]*:[0-9]*
@@ -87,6 +92,11 @@ RUN bpftrace --unsafe -v -e 'i:ms:100 { system("echo 'ok_system'"); exit();}'
 EXPECT ok_system
 TIMEOUT 5
 
+NAME system_more_args
+RUN bpftrace --unsafe -v -e 'i:ms:100 { system("echo %d %d %d %d %d %d %d", 1, 2, 3, 4, 5, 6, 7); exit();}'
+EXPECT 1 2 3 4 5 6 7
+TIMEOUT 5
+
 NAME count
 RUN bpftrace -v -e 'i:ms:100 { @ = count(); exit();}'
 EXPECT @:\s[0-9]+
@@ -151,6 +161,11 @@ RUN bpftrace -v -e 'i:ms:1 { cat("/proc/uptime"); exit();}'
 EXPECT [0-9]*.[0-9]* [0-9]*.[0-9]*
 TIMEOUT 5
 
+NAME cat_more_args
+RUN bpftrace -v -e 'i:ms:1 { cat("/%s%s%s%s/%s%s%s", "p", "r", "o", "c", "u", "p", "time"); exit();}'
+EXPECT [0-9]*.[0-9]* [0-9]*.[0-9]*
+TIMEOUT 5
+
 NAME uaddr
 RUN bpftrace -v -e 'uprobe:testprogs/uprobe_test:function1 { printf("0x%lx -- 0x%lx\n", *uaddr("GLOBAL_A"), *uaddr("GLOBAL_C")); exit(); }' -c ./testprogs/uprobe_test
 EXPECT 0x55555555 -- 0x33333333

diff --git a/tests/semantic_analyser.cpp b/tests/semantic_analyser.cpp
@@ -902,6 +902,9 @@ TEST(semantic_analyser, printf)
   test("kprobe:f { printf(\"%A\", comm) }", 10);
   test("kprobe:f { @x = printf(\"hi\") }", 1);
   test("kprobe:f { $x = printf(\"hi\") }", 1);
+  test("kprobe:f { printf(\"%d %d %d %d %d %d %d %d %d\", 1, 2, 3, 4, 5, 6, 7, "
+       "8, 9); }",
+       0);
 }
 
 TEST(semantic_analyser, system)