Support usdt args with index register and scale

CHANGELOG.md

@@ -52,6 +52,8 @@ and this project adheres to
   - [#1656](https://github.com/iovisor/bpftrace/pull/1656)
 - Add builtin function: `macaddr`
   - [#1647](https://github.com/iovisor/bpftrace/pull/1647)
+- Add support for usdt arguments utilising the index register and scale
+  - [#1684](https://github.com/iovisor/bpftrace/pull/1684)
 
 #### Changed
 - Warn if using `print` on `stats` maps with top and div arguments
@@ -140,6 +142,8 @@ and this project adheres to
   - [#1600](https://github.com/iovisor/bpftrace/pull/1600)
 - Fix attaching to multiple usdt probe locations with the same label
   - [#1681](https://github.com/iovisor/bpftrace/pull/1681)
+- Fix signed extension of usdt arguments to the internal 64-bit integer type
+  - [#1684](https://github.com/iovisor/bpftrace/pull/1684)
 
 #### Tools
 - Hook up execsnoop.bt script onto `execveat` call

src/ast/irbuilderbpf.cpp

@@ -555,20 +555,43 @@ Value *IRBuilderBPF::CreateUSDTReadArgument(Value *ctx,
                                             const location &loc)
 {
   assert(ctx && ctx->getType() == getInt8PtrTy());
-  // TODO (mmarchini): Handle base + index * scale addressing.
-  // https://github.com/iovisor/bcc/pull/988
-  if (argument->valid & BCC_USDT_ARGUMENT_INDEX_REGISTER_NAME)
-    LOG(ERROR) << "index register is not handled yet ["
-               << argument->index_register_name << "]";
-  if (argument->valid & BCC_USDT_ARGUMENT_SCALE)
-    LOG(ERROR) << "scale is not handled yet [" << argument->scale << "]";
+  // Argument size must be 1, 2, 4, or 8. See
+  // https://sourceware.org/systemtap/wiki/UserSpaceProbeImplementation
+  int abs_size = std::abs(argument->size);
+  assert(abs_size == 1 || abs_size == 2 || abs_size == 4 || abs_size == 8);
   if (argument->valid & BCC_USDT_ARGUMENT_DEREF_IDENT)
-    LOG(ERROR) << "defer ident is not handled yet [" << argument->deref_ident
+    LOG(ERROR) << "deref ident is not handled yet [" << argument->deref_ident
                << "]";
+  // USDT arguments can be any valid gas (GNU asm) operand.
+  // BCC normalises these into the bcc_usdt_argument and supports most
+  // valid gas operands.
+  //
+  // This code handles the following argument types:
+  // * A constant (ARGUMENT_CONSTANT)
+  //
+  // * The value of a register (ARGUMENT_BASE_REGISTER_NAME without
+  // ARGUMENT_DEREF_OFFSET set).
+  //
+  // * The value at address: base_register + offset + (index_register * scale)
+  // Where index_register and scale are optional.
+  // Note: Offset is optional in the gas operand, however will be set as zero
+  // if the register needs to be dereferenced.
 
   if (argument->valid & BCC_USDT_ARGUMENT_CONSTANT)
-    return getInt64(argument->constant);
+  {
+    // Correctly sign extend and convert to a 64-bit int
+    return CreateIntCast(getIntN(abs_size * 8, argument->constant),
+                         getInt64Ty(),
+                         argument->size < 0);
+  }
 
+  if (argument->valid & BCC_USDT_ARGUMENT_INDEX_REGISTER_NAME &&
+      !(argument->valid & BCC_USDT_ARGUMENT_BASE_REGISTER_NAME))
+  {
+    // Invalid combination??
+    LOG(ERROR) << "index register set without base register;"
+               << " this case is not yet handled";
+  }
   Value *result = nullptr;
   if (argument->valid & BCC_USDT_ARGUMENT_BASE_REGISTER_NAME) {
     int offset = 0;
@@ -579,29 +602,46 @@ Value *IRBuilderBPF::CreateUSDTReadArgument(Value *ctx,
                  << " not known";
     }
 
-    // Argument size must be 1, 2, 4, or 8. See
-    // https://sourceware.org/systemtap/wiki/UserSpaceProbeImplementation
-    int abs_size = std::abs(argument->size);
-    assert(abs_size == 1 || abs_size == 2 || abs_size == 4 || abs_size == 8);
 
     // bpftrace's args are internally represented as 64 bit integers. However,
     // the underlying argument (of the target program) may be less than 64
     // bits. So we must be careful to zero out unused bits.
     Value* reg = CreateGEP(ctx, getInt64(offset * sizeof(uintptr_t)), "load_register");
     AllocaInst *dst = CreateAllocaBPF(builtin.type, builtin.ident);
+    Value *index_offset = nullptr;
+    if (argument->valid & BCC_USDT_ARGUMENT_INDEX_REGISTER_NAME)
+    {
+      int ioffset = arch::offset(argument->index_register_name);
+      if (ioffset < 0)
+      {
+        LOG(FATAL) << "offset for register " << argument->index_register_name
+                   << " not known";
+      }
+      index_offset = CreateGEP(ctx,
+                               getInt64(ioffset * sizeof(uintptr_t)),
+                               "load_register");
+      index_offset = CreateLoad(getInt64Ty(), index_offset);
+      if (argument->valid & BCC_USDT_ARGUMENT_SCALE)
+      {
+        index_offset = CreateMul(index_offset, getInt64(argument->scale));
+      }
+    }
     if (argument->valid & BCC_USDT_ARGUMENT_DEREF_OFFSET) {
       Value *ptr = CreateAdd(CreateLoad(getInt64Ty(), reg),
                              getInt64(argument->deref_offset));
-      // Zero out `dst` here in case we read less than 64 bits
-      CreateStore(getInt64(0), dst);
+      if (index_offset)
+      {
+        ptr = CreateAdd(ptr, index_offset);
+      }
       CreateProbeRead(ctx, dst, abs_size, ptr, as, loc);
-      result = CreateLoad(dst);
+      result = CreateLoad(getIntNTy(abs_size * 8), dst);
     }
     else
     {
-      result = CreateLoad(GetType(CreateInt(abs_size * 8)), reg);
-      result = CreateIntCast(result, getInt64Ty(), argument->size < 0);
+      result = CreateLoad(getIntNTy(abs_size * 8), reg);
     }
+    // Sign extend and convert to a bpftools standard 64-bit integer type
+    result = CreateIntCast(result, getInt64Ty(), argument->size < 0);
     CreateLifetimeEnd(dst);
   }
   return result;

tests/runtime/usdt

@@ -130,7 +130,7 @@ TIMEOUT 5
 
 NAME "usdt probes - attach to probe on multiple files by wildcard"
 RUN bpftrace -e 'usdt:./testprogs/usdt*::* { printf("here\n" ); exit(); }'
-EXPECT Attaching 9 probes...
+EXPECT Attaching 41 probes...
 TIMEOUT 5
 
 NAME "usdt probes - attach to probe on multiple providers by wildcard and pid"
@@ -259,6 +259,32 @@ EXPECT ^1$
 TIMEOUT 5
 BEFORE ./testprogs/usdt_sized_args
 
+NAME "usdt constant arguments"
+RUN bpftrace -e 'usdt:./testprogs/usdt_args:usdt_args:const_* { printf("%lld ", arg0); }' -c ./testprogs/usdt_args
+EXPECT 4092785136 -202182160 61936 -3600 240 -16
+#EXPECT -579005069656919568 -579005069656919568 4092785136 -202182160 61936 -3600 240 -16
+# Bug in bcc, constants are stored in a 32-bit int, so 64-bit values are truncated
+TIMEOUT 5
+REQUIRES ./testprogs/usdt_args should_not_skip
+
+NAME "usdt reg arguments"
+RUN bpftrace -e 'usdt:./testprogs/usdt_args:usdt_args:reg_* { printf("%lld ", arg0); }' -c ./testprogs/usdt_args
+EXPECT -579005069656919568 -579005069656919568 4092785136 -202182160 61936 -3600 240 -16
+TIMEOUT 5
+REQUIRES ./testprogs/usdt_args should_not_skip
+
+NAME "usdt addr arguments"
+RUN bpftrace -e 'usdt:./testprogs/usdt_args:usdt_args:addr_* { printf("%lld ", arg0); }' -c ./testprogs/usdt_args
+EXPECT -579005069656919568 -579005069656919568 4092785136 -202182160 61936 -3600 240 -16
+TIMEOUT 5
+REQUIRES ./testprogs/usdt_args should_not_skip
+
+NAME "usdt addr+index arguments"
+RUN bpftrace -e 'usdt:./testprogs/usdt_args:usdt_args:index_* { printf("%lld ", arg0); }' -c ./testprogs/usdt_args
+EXPECT -579005069656919568 -579005069656919568 4092785136 -202182160 61936 -3600 240 -16
+TIMEOUT 5
+REQUIRES ./testprogs/usdt_args should_not_skip
+
 # USDT probes can be inlined which creates duplicate identical probes. We must
 # attach to all of them
 NAME "usdt duplicated markers"

tests/testprogs/usdt_args.c

@@ -0,0 +1,159 @@
+#ifdef HAVE_SYSTEMTAP_SYS_SDT_H
+#include <sys/sdt.h>
+#else
+#define DTRACE_PROBE1(a, b, d) (void)0
+#endif
+#include <stddef.h>
+#include <stdint.h>
+/* For best results compile using gcc -O1 (or higher)
+ *
+ * Clang likes to always put the argument on stack rather than outputting an
+ * operand to access it directly.
+ *
+ * Note: The macros might not always output the type of probe claimed as this
+ * varies depending on compiler and optimisation settings.
+ *
+ * So I've used macros to make the repeated code easier to fix in the
+ * future.
+ */
+
+#if (defined(__GNUC__) && !defined(__clang__))
+#pragma GCC optimize("O1")
+#else
+#pragma message("non-gcc compiler: the correct probes might not be installed")
+#endif
+
+#define test_value 0xf7f6f5f4f3f2f1f0
+volatile uint64_t one = 1;
+
+typedef union all_types
+{
+  int8_t i_8;
+  uint8_t i_u8;
+  uint16_t i_u16;
+  int16_t i_16;
+  uint32_t i_u32;
+  int32_t i_32;
+  uint64_t i_u64;
+  int64_t i_64;
+} all_types_t;
+
+#ifdef __x86_64__
+/* Force the value to be stored in a register, works for both gcc and clang
+ *
+ * PROBE_REG(al, u, 8) expands to:
+ * register uint8_t _reg_u8 asm("al") = (uint8_t)test_value;
+ * DTRACE_PROBE1(usdt_args, reg_u8, _reg_u8)
+ */
+#define PROBE_REG(register_, sign, size)                                       \
+  register sign##int##size##_t _reg_##sign##size asm(                          \
+      #register_) = (sign##int##size##_t)test_value;                           \
+  DTRACE_PROBE1(usdt_args, reg_##sign##size, _reg_##sign##size)
+#else
+/* Force a calculation to get the value into a register
+ * works for gcc (-O1+) but not clang
+ *
+ * PROBE_REG(al, u, 8) expands to:
+ * uint8_t _reg_u8 = (uint8_t)(test_value * one);
+ * DTRACE_PROBE1(usdt_args, reg_u8, _reg_u8);
+ * */
+#define PROBE_REG(register_, sign, size)                                       \
+  sign##int##size##_t _reg_##sign##size = (sign##int##size##_t)(test_value *   \
+                                                                one);          \
+  DTRACE_PROBE1(usdt_args, reg_##sign##size, _reg_##sign##size)
+#endif
+
+/* Read a value of the stack, expect an offset here
+ * Works with optimisation -O1+
+ *
+ * PROBE_ADDRESS(u, 8) expands to:
+ * array[1].i_u8 = (uint8_t) test_value;
+ * DTRACE_PROBE1(usdt_args, addr_u8, array[1].i_u8)
+ */
+#define PROBE_ADDRESS(sign, size)                                              \
+  array[1].i_##sign##size = (sign##int##size##_t)test_value;                   \
+  DTRACE_PROBE1(usdt_args, addr_##sign##size, array[1].i_##sign##size)
+
+/* Read a value of the stack, expect an offset here
+ * Works only with gcc optimisation -O1+
+ *
+ * PROBE_INDEX(u, 8) expands to:
+ * DTRACE_PROBE1(usdt_args, index_u8, array[i].i_u8)
+ */
+#define PROBE_INDEX(sign, size)                                                \
+  DTRACE_PROBE1(usdt_args, index_##sign##size, array[i].i_##sign##size)
+
+int main(int argc, char **argv)
+{
+  (void)argv;
+  volatile all_types_t array[10];
+
+  if (argc > 1)
+  // If we don't have Systemtap headers, we should skip USDT tests. Returning 1
+  // can be used as validation in the REQUIRE
+#ifndef HAVE_SYSTEMTAP_SYS_SDT_H
+    return 1;
+#else
+    return 0;
+#endif
+
+  for (volatile size_t i = 0; i < sizeof(array) / sizeof(array[0]); i++)
+  {
+    array[i].i_u64 = test_value;
+  }
+
+#ifdef HAVE_SYSTEMTAP_SYS_SDT_H
+  /* Constants */
+  DTRACE_PROBE1(usdt_args, const_u64, (uint64_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_64, (int64_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_u32, (uint32_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_32, (int32_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_u16, (uint16_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_16, (int16_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_u8, (uint8_t)test_value);
+  DTRACE_PROBE1(usdt_args, const_8, (int8_t)test_value);
+
+  /* Direct register reads - start from 64 and work down
+   * to verify the correct number of bytes are read */
+  PROBE_REG(rax, u, 64);
+  PROBE_REG(rax, , 64);
+  PROBE_REG(eax, u, 32);
+  PROBE_REG(eax, , 32);
+  PROBE_REG(ax, u, 16);
+  PROBE_REG(ax, , 16);
+  PROBE_REG(al, u, 8);
+  PROBE_REG(al, , 8);
+
+  /* Base address most likely with and offset(aka. displacement) */
+  PROBE_ADDRESS(u, 64);
+  PROBE_ADDRESS(, 64);
+  PROBE_ADDRESS(u, 32);
+  PROBE_ADDRESS(, 32);
+  PROBE_ADDRESS(u, 16);
+  PROBE_ADDRESS(, 16);
+  PROBE_ADDRESS(u, 8);
+  PROBE_ADDRESS(, 8);
+
+  /* Base address + offset + (index * scale) */
+  for (volatile int i = 7; i <= 7; i++)
+  {
+    PROBE_INDEX(u, 64);
+    PROBE_INDEX(, 64);
+    PROBE_INDEX(u, 32);
+    PROBE_INDEX(, 32);
+    PROBE_INDEX(u, 16);
+    PROBE_INDEX(, 16);
+    PROBE_INDEX(u, 8);
+    PROBE_INDEX(, 8);
+  }
+
+/* TLS not yet supported, need label and segment support */
+#if 0
+  volatile static __thread uint64_t tls_64 = (uint64_t)test_value;
+  DTRACE_PROBE1(usdt_args, tls_u64, tls_64);
+#endif
+#endif /* HAVE_SYSTEMTAP_SYS_SDT_H */
+
+  return 0;
+}
+