Support base + index * scale addressing for USDT arguments

It turns out that some software will have USDT probe arguments
referencing memory using the full `nnn@(%basereg + %idxreg * scale`
syntax. This is represented as `nnn@(%basereg,%idxreg,scale)` in
the `NT_STAPSDT` note, encountered in building a recent version of
PostgreSQL on FC25.

This format is now recognized by the USDT parser, and the correct
BPF code is emitted to retrieve arguments that reference memory
using this full addressing syntax.`

src/cc/bcc_usdt.h

@@ -39,11 +39,13 @@ struct bcc_usdt_location {
     uint64_t address;
 };
 
-#define BCC_USDT_ARGUMENT_NONE          0x0
-#define BCC_USDT_ARGUMENT_CONSTANT      0x1
-#define BCC_USDT_ARGUMENT_DEREF_OFFSET  0x2
-#define BCC_USDT_ARGUMENT_DEREF_IDENT   0x4
-#define BCC_USDT_ARGUMENT_REGISTER_NAME 0x8
+#define BCC_USDT_ARGUMENT_NONE                0x0
+#define BCC_USDT_ARGUMENT_CONSTANT            0x1
+#define BCC_USDT_ARGUMENT_DEREF_OFFSET        0x2
+#define BCC_USDT_ARGUMENT_DEREF_IDENT         0x4
+#define BCC_USDT_ARGUMENT_BASE_REGISTER_NAME  0x8
+#define BCC_USDT_ARGUMENT_INDEX_REGISTER_NAME 0x10
+#define BCC_USDT_ARGUMENT_SCALE               0x20
 
 struct bcc_usdt_argument {
     int size;
@@ -51,7 +53,9 @@ struct bcc_usdt_argument {
     int constant;
     int deref_offset;
     const char *deref_ident;
-    const char *register_name;
+    const char *base_register_name;
+    const char *index_register_name;
+    int scale;
 };
 
 typedef void (*bcc_usdt_cb)(struct bcc_usdt *);

src/cc/usdt.cc

@@ -360,48 +360,56 @@ void bcc_usdt_foreach(void *usdt, bcc_usdt_cb callback) {
 
 int bcc_usdt_get_location(void *usdt, const char *probe_name,
                           int index, struct bcc_usdt_location *location) {
-    USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
-    USDT::Probe *probe = ctx->get(probe_name);
-    if (!probe)
-        return -1;
-    if (index < 0 || (size_t)index >= probe->num_locations())
-        return -1;
-    location->address = probe->address(index);
-    return 0;
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  USDT::Probe *probe = ctx->get(probe_name);
+  if (!probe)
+    return -1;
+  if (index < 0 || (size_t)index >= probe->num_locations())
+    return -1;
+  location->address = probe->address(index);
+  return 0;
 }
 
 int bcc_usdt_get_argument(void *usdt, const char *probe_name,
                           int location_index, int argument_index,
                           struct bcc_usdt_argument *argument) {
-    USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
-    USDT::Probe *probe = ctx->get(probe_name);
-    if (!probe)
-        return -1;
-    if (argument_index < 0 || (size_t)argument_index >= probe->num_arguments())
-        return -1;
-    if (location_index < 0 || (size_t)location_index >= probe->num_locations())
-        return -1;
-    auto const &location = probe->location(location_index);
-    auto const &arg = location.arguments_[argument_index];
-    argument->size = arg.arg_size();
-    argument->valid = BCC_USDT_ARGUMENT_NONE;
-    if (arg.constant()) {
-        argument->valid |= BCC_USDT_ARGUMENT_CONSTANT;
-        argument->constant = *(arg.constant());
-    }
-    if (arg.deref_offset()) {
-        argument->valid |= BCC_USDT_ARGUMENT_DEREF_OFFSET;
-        argument->deref_offset = *(arg.deref_offset());
-    }
-    if (arg.deref_ident()) {
-        argument->valid |= BCC_USDT_ARGUMENT_DEREF_IDENT;
-        argument->deref_ident = arg.deref_ident()->c_str();
-    }
-    if (arg.register_name()) {
-        argument->valid |= BCC_USDT_ARGUMENT_REGISTER_NAME;
-        argument->register_name = arg.register_name()->c_str();
-    }
-    return 0;
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  USDT::Probe *probe = ctx->get(probe_name);
+  if (!probe)
+    return -1;
+  if (argument_index < 0 || (size_t)argument_index >= probe->num_arguments())
+    return -1;
+  if (location_index < 0 || (size_t)location_index >= probe->num_locations())
+    return -1;
+  auto const &location = probe->location(location_index);
+  auto const &arg = location.arguments_[argument_index];
+  argument->size = arg.arg_size();
+  argument->valid = BCC_USDT_ARGUMENT_NONE;
+  if (arg.constant()) {
+    argument->valid |= BCC_USDT_ARGUMENT_CONSTANT;
+    argument->constant = *(arg.constant());
+  }
+  if (arg.deref_offset()) {
+    argument->valid |= BCC_USDT_ARGUMENT_DEREF_OFFSET;
+    argument->deref_offset = *(arg.deref_offset());
+  }
+  if (arg.deref_ident()) {
+    argument->valid |= BCC_USDT_ARGUMENT_DEREF_IDENT;
+    argument->deref_ident = arg.deref_ident()->c_str();
+  }
+  if (arg.base_register_name()) {
+    argument->valid |= BCC_USDT_ARGUMENT_BASE_REGISTER_NAME;
+    argument->base_register_name = arg.base_register_name()->c_str();
+  }
+  if (arg.index_register_name()) {
+    argument->valid |= BCC_USDT_ARGUMENT_INDEX_REGISTER_NAME;
+    argument->index_register_name = arg.index_register_name()->c_str();
+  }
+  if (arg.scale()) {
+    argument->valid |= BCC_USDT_ARGUMENT_SCALE;
+    argument->scale = *(arg.scale());
+  }
+  return 0;
 }
 
 void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback) {

src/cc/usdt.h

@@ -40,7 +40,9 @@ private:
   optional<int> constant_;
   optional<int> deref_offset_;
   optional<std::string> deref_ident_;
-  optional<std::string> register_name_;
+  optional<std::string> base_register_name_;
+  optional<std::string> index_register_name_;
+  optional<int> scale_;
 
   bool get_global_address(uint64_t *address, const std::string &binpath,
                           const optional<int> &pid) const;
@@ -57,7 +59,13 @@ public:
   std::string ctype() const;
 
   const optional<std::string> &deref_ident() const { return deref_ident_; }
-  const optional<std::string> &register_name() const { return register_name_; }
+  const optional<std::string> &base_register_name() const {
+    return base_register_name_;
+  }
+  const optional<std::string> &index_register_name() const {
+    return index_register_name_;
+  }
+  const optional<int> scale() const { return scale_; }
   const optional<int> constant() const { return constant_; }
   const optional<int> deref_offset() const { return deref_offset_; }
 
@@ -68,12 +76,18 @@ class ArgumentParser {
   const char *arg_;
   ssize_t cur_pos_;
 
+  void skip_whitespace_from(size_t pos);
+  void skip_until_whitespace_from(size_t pos);
+
 protected:
   virtual bool normalize_register(std::string *reg, int *reg_size) = 0;
 
+  ssize_t parse_register(ssize_t pos, std::string &name, int &size);
   ssize_t parse_number(ssize_t pos, optional<int> *number);
   ssize_t parse_identifier(ssize_t pos, optional<std::string> *ident);
-  ssize_t parse_register(ssize_t pos, Argument *dest);
+  ssize_t parse_base_register(ssize_t pos, Argument *dest);
+  ssize_t parse_index_register(ssize_t pos, Argument *dest);
+  ssize_t parse_scale(ssize_t pos, Argument *dest);
   ssize_t parse_expr(ssize_t pos, Argument *dest);
   ssize_t parse_1(ssize_t pos, Argument *dest);
 

src/cc/usdt_args.cc

@@ -61,20 +61,28 @@ bool Argument::assign_to_local(std::ostream &stream,
 
   if (!deref_offset_) {
     tfm::format(stream, "%s = (%s)ctx->%s;", local_name, ctype(),
-                *register_name_);
+                *base_register_name_);
     return true;
   }
 
   if (deref_offset_ && !deref_ident_) {
+    tfm::format(stream, "{ u64 __addr = ctx->%s + (%d)",
+                *base_register_name_, *deref_offset_);
+    if (index_register_name_) {
+      int scale = scale_.value_or(1);
+      tfm::format(stream, " + (ctx->%s * %d);", *index_register_name_, scale);
+    } else {
+      tfm::format(stream, ";");
+    }
     tfm::format(stream,
-                "{ u64 __addr = ctx->%s + (%d); %s __res = 0x0; "
+                "%s __res = 0x0; "
                 "bpf_probe_read(&__res, sizeof(__res), (void *)__addr); "
                 "%s = __res; }",
-                *register_name_, *deref_offset_, ctype(), local_name);
+                ctype(), local_name);
     return true;
   }
 
-  if (deref_offset_ && deref_ident_ && *register_name_ == "ip") {
+  if (deref_offset_ && deref_ident_ && *base_register_name_ == "ip") {
     uint64_t global_address;
     if (!get_global_address(&global_address, binpath, pid))
       return false;
@@ -109,7 +117,8 @@ ssize_t ArgumentParser::parse_identifier(ssize_t pos,
   return pos;
 }
 
-ssize_t ArgumentParser::parse_register(ssize_t pos, Argument *dest) {
+ssize_t ArgumentParser::parse_register(ssize_t pos, std::string &name,
+                                       int &size) {
   ssize_t start = ++pos;
   if (arg_[start - 1] != '%')
     return -start;
@@ -117,16 +126,41 @@ ssize_t ArgumentParser::parse_register(ssize_t pos, Argument *dest) {
   while (isalnum(arg_[pos])) pos++;
 
   std::string regname(arg_ + start, pos - start);
-  int regsize = 0;
-
-  if (!normalize_register(&regname, &regsize))
+  if (!normalize_register(&regname, &size))
     return -start;
 
-  dest->register_name_ = regname;
+  name = regname;
+  return pos;
+}
+
+ssize_t ArgumentParser::parse_base_register(ssize_t pos, Argument *dest) {
+  int size;
+  std::string name;
+  ssize_t res = parse_register(pos, name, size);
+  if (res < 0)
+      return res;
+
+  dest->base_register_name_ = name;
   if (!dest->arg_size_)
-    dest->arg_size_ = regsize;
+    dest->arg_size_ = size;
 
-  return pos;
+  return res;
+}
+
+ssize_t ArgumentParser::parse_index_register(ssize_t pos, Argument *dest) {
+  int size;
+  std::string name;
+  ssize_t res = parse_register(pos, name, size);
+  if (res < 0)
+      return res;
+
+  dest->index_register_name_ = name;
+
+  return res;
+}
+
+ssize_t ArgumentParser::parse_scale(ssize_t pos, Argument *dest) {
+  return parse_number(pos, &dest->scale_);
 }
 
 ssize_t ArgumentParser::parse_expr(ssize_t pos, Argument *dest) {
@@ -134,7 +168,7 @@ ssize_t ArgumentParser::parse_expr(ssize_t pos, Argument *dest) {
     return parse_number(pos + 1, &dest->constant_);
 
   if (arg_[pos] == '%')
-    return parse_register(pos, dest);
+    return parse_base_register(pos, dest);
 
   if (isdigit(arg_[pos]) || arg_[pos] == '-') {
     pos = parse_number(pos, &dest->deref_offset_);
@@ -154,10 +188,22 @@ ssize_t ArgumentParser::parse_expr(ssize_t pos, Argument *dest) {
   if (arg_[pos] != '(')
     return -pos;
 
-  pos = parse_register(pos + 1, dest);
+  pos = parse_base_register(pos + 1, dest);
   if (pos < 0)
     return pos;
 
+  if (arg_[pos] == ',') {
+    pos = parse_index_register(pos + 1, dest);
+    if (pos < 0)
+      return pos;
+
+    if (arg_[pos] == ',') {
+      pos = parse_scale(pos + 1, dest);
+      if (pos < 0)
+        return pos;
+    }
+  }
+
   return (arg_[pos] == ')') ? pos + 1 : -pos;
 }
 
@@ -180,6 +226,17 @@ void ArgumentParser::print_error(ssize_t pos) {
   fputc('\n', stderr);
 }
 
+void ArgumentParser::skip_whitespace_from(size_t pos) {
+    while (isspace(arg_[pos])) pos++;
+    cur_pos_ = pos;
+}
+
+void ArgumentParser::skip_until_whitespace_from(size_t pos) {
+    while (arg_[pos] != '\0' && !isspace(arg_[pos]))
+        pos++;
+    cur_pos_ = pos;
+}
+
 bool ArgumentParser::parse(Argument *dest) {
   if (done())
     return false;
@@ -187,16 +244,15 @@ bool ArgumentParser::parse(Argument *dest) {
   ssize_t res = parse_1(cur_pos_, dest);
   if (res < 0) {
     print_error(-res);
-    cur_pos_ = -res;
+    skip_whitespace_from(-res + 1);
     return false;
   }
   if (!isspace(arg_[res]) && arg_[res] != '\0') {
     print_error(res);
-    cur_pos_ = res;
+    skip_until_whitespace_from(res);
     return false;
   }
-  while (isspace(arg_[res])) res++;
-  cur_pos_ = res;
+  skip_whitespace_from(res);
   return true;
 }
 

src/python/bcc/libbcc.py

@@ -92,7 +92,7 @@ lib.bpf_attach_kprobe.argtypes = [ct.c_int, ct.c_int, ct.c_char_p, ct.c_char_p,
 lib.bpf_detach_kprobe.restype = ct.c_int
 lib.bpf_detach_kprobe.argtypes = [ct.c_char_p]
 lib.bpf_attach_uprobe.restype = ct.c_void_p
-lib.bpf_attach_uprobe.argtypes = [ct.c_int, ct.c_int, ct.c_char_p, ct.c_char_p, 
+lib.bpf_attach_uprobe.argtypes = [ct.c_int, ct.c_int, ct.c_char_p, ct.c_char_p,
         ct.c_ulonglong, ct.c_int, ct.c_int, ct.c_int, _CB_TYPE, ct.py_object]
 lib.bpf_detach_uprobe.restype = ct.c_int
 lib.bpf_detach_uprobe.argtypes = [ct.c_char_p]
@@ -197,7 +197,9 @@ class BCC_USDT_ARGUMENT_FLAGS(object):
     CONSTANT = 0x1
     DEREF_OFFSET = 0x2
     DEREF_IDENT = 0x4
-    REGISTER_NAME = 0x8
+    BASE_REGISTER_NAME = 0x8
+    INDEX_REGISTER_NAME = 0x10
+    SCALE = 0x20
 
 class bcc_usdt_argument(ct.Structure):
     _fields_ = [
@@ -206,7 +208,9 @@ class bcc_usdt_argument(ct.Structure):
             ('constant', ct.c_int),
             ('deref_offset', ct.c_int),
             ('deref_ident', ct.c_char_p),
-            ('register_name', ct.c_char_p)
+            ('base_register_name', ct.c_char_p),
+            ('index_register_name', ct.c_char_p),
+            ('scale', ct.c_int)
         ]
 
 _USDT_CB = ct.CFUNCTYPE(None, ct.POINTER(bcc_usdt))

src/python/bcc/usdt.py

@@ -32,8 +32,12 @@ class USDTProbeArgument(object):
             self.deref_offset = argument.deref_offset
         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT != 0:
             self.deref_ident = argument.deref_ident
-        if self.valid & BCC_USDT_ARGUMENT_FLAGS.REGISTER_NAME != 0:
-            self.register_name = argument.register_name
+        if self.valid & BCC_USDT_ARGUMENT_FLAGS.BASE_REGISTER_NAME != 0:
+            self.base_register_name = argument.base_register_name
+        if self.valid & BCC_USDT_ARGUMENT_FLAGS.INDEX_REGISTER_NAME != 0:
+            self.index_register_name = argument.index_register_name
+        if self.valid & BCC_USDT_ARGUMENT_FLAGS.SCALE != 0:
+            self.scale = argument.scale
 
     def _size_prefix(self):
         return "%d %s bytes" % \
@@ -45,16 +49,22 @@ class USDTProbeArgument(object):
         if self.valid & BCC_USDT_ARGUMENT_FLAGS.CONSTANT != 0:
             return "%d" % self.constant
         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET == 0:
-            return "%s" % self.register_name
+            return "%s" % self.base_register_name
         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0 and \
            self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT == 0:
+            if self.valid & BCC_USDT_ARGUMENT_FLAGS.INDEX_REGISTER_NAME != 0:
+                index_offset = " + %s" % self.index_register_name
+                if self.valid & BCC_USDT_ARGUMENT_FLAGS.SCALE != 0:
+                    index_offset += " * %d" % self.scale
+            else:
+                index_offset = ""
             sign = '+' if self.deref_offset >= 0 else '-'
-            return "*(%s %s %d)" % (self.register_name,
-                                    sign, abs(self.deref_offset))
+            return "*(%s %s %d%s)" % (self.base_register_name,
+                                    sign, abs(self.deref_offset), index_offset)
         if self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_OFFSET != 0 and \
            self.valid & BCC_USDT_ARGUMENT_FLAGS.DEREF_IDENT != 0 and \
-           self.valid & BCC_USDT_ARGUMENT_FLAGS.REGISTER_NAME != 0 and \
-           self.register_name == "ip":
+           self.valid & BCC_USDT_ARGUMENT_FLAGS.BASE_REGISTER_NAME != 0 and \
+           self.base_register_name == "ip":
             sign = '+' if self.deref_offset >= 0 else '-'
             return "*(&%s %s %d)" % (self.deref_ident,
                                      sign, abs(self.deref_offset))

tests/cc/test_usdt_args.cc

@@ -35,16 +35,21 @@ static void verify_register(USDT::ArgumentParser_x64 &parser, int arg_size,
 static void verify_register(USDT::ArgumentParser_x64 &parser, int arg_size,
                             const std::string &regname,
                             optional<int> deref_offset = nullopt,
-                            optional<std::string> deref_ident = nullopt) {
+                            optional<std::string> deref_ident = nullopt,
+                            optional<std::string> index_regname = nullopt,
+                            optional<int> scale = nullopt) {
   USDT::Argument arg;
   REQUIRE(parser.parse(&arg));
   REQUIRE(arg.arg_size() == arg_size);
 
-  REQUIRE(arg.register_name());
-  REQUIRE(arg.register_name() == regname);
+  REQUIRE(arg.base_register_name());
+  REQUIRE(arg.base_register_name() == regname);
 
   REQUIRE(arg.deref_offset() == deref_offset);
   REQUIRE(arg.deref_ident() == deref_ident);
+
+  REQUIRE(arg.index_register_name() == index_regname);
+  REQUIRE(arg.scale() == scale);
 }
 
 TEST_CASE("test usdt argument parsing", "[usdt]") {
@@ -66,7 +71,9 @@ TEST_CASE("test usdt argument parsing", "[usdt]") {
         "-4@global_max_action(%rip) "
         "8@24+mp_(%rip) "
         "-4@CheckpointStats+40(%rip) "
-        "4@glob-2(%rip) ");
+        "4@glob-2(%rip) "
+        "8@(%rax,%rdx,8) "
+        "4@(%rbx,%rcx)");
 
     verify_register(parser, -4, 0);
     verify_register(parser, 8, 1234);
@@ -85,6 +92,9 @@ TEST_CASE("test usdt argument parsing", "[usdt]") {
     verify_register(parser, -4, "ip", 40, std::string("CheckpointStats"));
     verify_register(parser, 4, "ip", -2, std::string("glob"));
 
+    verify_register(parser, 8, "ax", 0, nullopt, std::string("dx"), 8);
+    verify_register(parser, 4, "bx", 0, nullopt, std::string("cx"));
+
     REQUIRE(parser.done());
   }
 }