Merge pull request #988 from goldshtn/usdt-addressed-arg

Support base + index * scale addressing for USDT arguments

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());
   }
 }