add support for init namespace to pid and tid (bpftrace#4254)

* docs: update descriptions about PID namespaces

The manual didn't reflect the changes in bpftrace#3428.

Signed-off-by: Junxuan Liao <[email protected]>

* add support for init namespace to `pid` and `tid`

Introduce two function variants of the built-ins.

fixes bpftrace#4253.

Signed-off-by: Junxuan Liao <[email protected]>

---------

Signed-off-by: Junxuan Liao <[email protected]>

Author: MikeWalrus

CHANGELOG.md

@@ -55,6 +55,8 @@ and this project adheres to
   - [#4126](https://github.com/bpftrace/bpftrace/pull/4126)
 - For loops now support `break` and `continue`
   - [#4250](https://github.com/bpftrace/bpftrace/pull/4250)
+- Add `pid` and `tid` functions for choosing between the initial or the current namespace
+  - [#4254](https://github.com/bpftrace/bpftrace/pull/4254)
 #### Changed
 - `-p` CLI flag now applies to all probes (except BEGIN/END)
   - [#3800](https://github.com/bpftrace/bpftrace/pull/3800)

docs/migration_guide.md

@@ -24,6 +24,14 @@ config = {
 }
 ```
 
+## Versions 0.22.x (or earlier) to 0.24.x (or later)
+After https://github.com/bpftrace/bpftrace/pull/3428 `pid` and `tid` are no
+longer from the initial namespace.
+
+https://github.com/bpftrace/bpftrace/pull/3428 introduced `pid(init)` and
+`tid(init)` which have the old behavior.
+Unfortunately, there are no equivalent workarounds in 0.23.x.
+
 ## Versions 0.21.x (or earlier) to 0.22.x (or later)
 
 ### Added block scoping for scratch variables

man/adoc/bpftrace.adoc

@@ -1988,8 +1988,8 @@ For string arguments in an action block use the `str()` call to retrieve the val
 | `pid`
 | uint32
 | 4.2
-| get_current_pid_tgid
-| Process ID of the current thread (aka thread group ID), as seen from the init namespace
+| get_current_pid_tgid / get_ns_current_pid_tgid
+| Process ID of the current thread (aka thread group ID), as seen from the PID namespace of bpftrace
 
 | `probe`
 | string
@@ -2018,8 +2018,8 @@ For string arguments in an action block use the `str()` call to retrieve the val
 | `tid`
 | uint32
 | 4.2
-| get_current_pid_tgid
-| Thread ID of the current thread, as seen from the init namespace
+| get_current_pid_tgid / get_ns_current_pid_tgid
+| Thread ID of the current thread, as seen from the PID namespace of bpftrace
 
 | `uid`
 | uint64
@@ -2141,6 +2141,14 @@ Tracing block I/O sizes > 0 bytes
 | Print file content
 | Async
 
+| <<functions-pid, `pid`>>
+| Process ID of the current thread
+| Sync
+
+| <<functions-tid, `tid`>>
+| Thread ID of the current thread
+| Sync
+
 | <<functions-cgroupid, `cgroupid`>>
 | Resolve cgroup ID
 | Compile Time
@@ -2350,6 +2358,28 @@ tracepoint:syscalls:sys_enter_execve {
 55f683ee2000-55f683ee3000 rw-p 00000000 00:00 0
 ----
 
+[#functions-pid]
+=== pid
+.variants
+* `uint32 pid([curr_ns|init])`
+
+Returns the process ID of the current thread.
+Defaults to `curr_ns`.
+
+- `pid(curr_ns)` - The process ID as seen from the PID namespace of bpftrace.
+- `pid(init)` - The process ID as seen from the initial PID namespace.
+
+[#functions-tid]
+=== pid
+.variants
+* `uint32 tid([curr_ns|init])`
+
+Returns the thread ID of the current thread.
+Defaults to `curr_ns`.
+
+- `tid(curr_ns)` - The thread ID as seen from the PID namespace of bpftrace.
+- `tid(init)` - The thread ID as seen from the initial PID namespace.
+
 [#functions-cgroupid]
 === cgroupid
 

src/ast/irbuilderbpf.cpp

@@ -8,6 +8,7 @@
 #include "bpftrace.h"
 #include "globalvars.h"
 #include "log.h"
+#include "types.h"
 #include "util/exceptions.h"
 #include <filesystem>
 #include <llvm/IR/DataLayout.h>
@@ -74,10 +75,10 @@ libbpf::bpf_func_id IRBuilderBPF::selectProbeReadHelper(AddrSpace as, bool str)
 // It represents the inode of the initial (global) PID namespace
 constexpr uint32_t PROC_PID_INIT_INO = 0xeffffffc;
 
-Value *IRBuilderBPF::CreateGetPid(const Location &loc)
+Value *IRBuilderBPF::CreateGetPid(const Location &loc, bool force_init)
 {
   const auto &pidns = bpftrace_.get_pidns_self_stat();
-  if (pidns && pidns->st_ino != PROC_PID_INIT_INO) {
+  if (!force_init && pidns && pidns->st_ino != PROC_PID_INIT_INO) {
     // Get namespaced target PID when we're running in a namespace
     AllocaInst *res = CreateAllocaBPF(BpfPidnsInfoType(), "bpf_pidns_info");
     CreateGetNsPidTgid(
@@ -89,16 +90,16 @@ Value *IRBuilderBPF::CreateGetPid(const Location &loc)
     return pid;
   }
 
-  // Get global target PID when we're in the initial namespace
+  // Get global target PID otherwise
   Value *pidtgid = CreateGetPidTgid(loc);
   Value *pid = CreateTrunc(CreateLShr(pidtgid, 32), getInt32Ty(), "pid");
   return pid;
 }
 
-Value *IRBuilderBPF::CreateGetTid(const Location &loc)
+Value *IRBuilderBPF::CreateGetTid(const Location &loc, bool force_init)
 {
   const auto &pidns = bpftrace_.get_pidns_self_stat();
-  if (pidns && pidns->st_ino != PROC_PID_INIT_INO) {
+  if (!force_init && pidns && pidns->st_ino != PROC_PID_INIT_INO) {
     // Get namespaced target TID when we're running in a namespace
     AllocaInst *res = CreateAllocaBPF(BpfPidnsInfoType(), "bpf_pidns_info");
     CreateGetNsPidTgid(
@@ -110,7 +111,7 @@ Value *IRBuilderBPF::CreateGetTid(const Location &loc)
     return tid;
   }
 
-  // Get global target TID when we're in the initial namespace
+  // Get global target TID otherwise
   Value *pidtgid = CreateGetPidTgid(loc);
   Value *tid = CreateTrunc(pidtgid, getInt32Ty(), "tid");
   return tid;
@@ -128,7 +129,7 @@ AllocaInst *IRBuilderBPF::CreateUSym(Value *val,
   StructType *usym_t = GetStructType("usym_t", elements, false);
   AllocaInst *buf = CreateAllocaBPF(usym_t, "usym");
 
-  Value *pid = CreateGetPid(loc);
+  Value *pid = CreateGetPid(loc, false);
   Value *probe_id_val = Constant::getIntegerValue(getInt32Ty(),
                                                   APInt(32, probe_id));
 

src/ast/irbuilderbpf.h

@@ -199,8 +199,8 @@ class IRBuilderBPF : public IRBuilder<> {
   StructType *GetStructType(const std::string &name,
                             const std::vector<llvm::Type *> &elements,
                             bool packed = false);
-  Value *CreateGetPid(const Location &loc);
-  Value *CreateGetTid(const Location &loc);
+  Value *CreateGetPid(const Location &loc, bool force_init);
+  Value *CreateGetTid(const Location &loc, bool force_init);
   AllocaInst *CreateUSym(Value *val, int probe_id, const Location &loc);
   Value *CreateRegisterRead(Value *ctx, const std::string &builtin);
   Value *CreateRegisterRead(Value *ctx, int offset, const std::string &name);

src/ast/passes/codegen_llvm.cpp

@@ -96,6 +96,11 @@ static auto getTargetMachine()
   return target;
 }
 
+static bool shouldForceInitPidNs(const ExpressionList &args)
+{
+  return args.size() == 1 && args.at(0).as<Identifier>()->ident == "init";
+}
+
 using namespace llvm;
 
 namespace {
@@ -622,7 +627,7 @@ ScopedExpr CodegenLLVM::kstack_ustack(const std::string &ident,
   // ustack keys are special: see IRBuilderBPF::GetStackStructType()
   if (is_ustack) {
     // store pid
-    b_.CreateStore(b_.CreateGetPid(loc),
+    b_.CreateStore(b_.CreateGetPid(loc, false),
                    b_.CreateGEP(stack_key_struct,
                                 stack_key,
                                 { b_.getInt64(0), b_.getInt32(2) }));
@@ -669,9 +674,9 @@ ScopedExpr CodegenLLVM::visit(Builtin &builtin)
                          builtin.builtin_type.stack_type,
                          builtin.loc);
   } else if (builtin.ident == "pid") {
-    return ScopedExpr(b_.CreateGetPid(builtin.loc));
+    return ScopedExpr(b_.CreateGetPid(builtin.loc, false));
   } else if (builtin.ident == "tid") {
-    return ScopedExpr(b_.CreateGetTid(builtin.loc));
+    return ScopedExpr(b_.CreateGetTid(builtin.loc, false));
   } else if (builtin.ident == "cgroup") {
     return ScopedExpr(b_.CreateGetCurrentCgroupId(builtin.loc));
   } else if (builtin.ident == "uid" || builtin.ident == "gid" ||
@@ -1807,6 +1812,14 @@ ScopedExpr CodegenLLVM::visit(Call &call)
     } else {
       return ScopedExpr(b_.CreateGetNs(call.return_type.ts_mode, call.loc));
     }
+  } else if (call.func == "pid") {
+    bool force_init = shouldForceInitPidNs(call.vargs);
+
+    return ScopedExpr(b_.CreateGetPid(call.loc, force_init));
+  } else if (call.func == "tid") {
+    bool force_init = shouldForceInitPidNs(call.vargs);
+
+    return ScopedExpr(b_.CreateGetTid(call.loc, force_init));
   } else {
     LOG(BUG) << "missing codegen for function \"" << call.func << "\"";
     __builtin_unreachable();

src/ast/passes/semantic_analyser.cpp

@@ -590,6 +590,11 @@ static const std::map<std::string, call_spec> CALL_SPEC = {
       .arg_types={
         map_type_spec{},
       } } },
+  { "pid",
+    { .min_args=0,
+      .max_args=1,
+      .discard_ret_warn = true },
+  },
 };
 // clang-format on
 
@@ -754,6 +759,12 @@ void SemanticAnalyser::visit(Identifier &identifier)
     } else {
       identifier.addError() << "Invalid timestamp mode: " << identifier.ident;
     }
+  } else if (func_ == "pid" || func_ == "tid") {
+    if (identifier.ident != "curr_ns" && identifier.ident != "init") {
+      identifier.addError()
+          << "Invalid PID namespace mode: " << identifier.ident
+          << " (expects: curr_ns or init)";
+    }
   } else {
     // Final attempt: try to parse as a stack mode.
     ConfigParser<StackMode> parser;
@@ -1729,6 +1740,16 @@ If you're seeing errors, try clamping the string sizes. For example:
       call.addError() << "Failed to initialize sw_tai in "
                          "userspace. This is very unexpected.";
     }
+  } else if (call.func == "pid" || call.func == "tid") {
+    call.return_type = CreateUInt32();
+    if (call.vargs.size() == 1) {
+      auto &arg = call.vargs.at(0);
+      if (!(arg.as<Identifier>())) {
+        call.addError() << call.func
+                        << "() only supports curr_ns and init as the argument ("
+                        << arg.type().GetTy() << " provided)";
+      }
+    }
   } else {
     call.addError() << "Unknown function: '" << call.func << "'";
   }

tests/codegen/call_pid_tid.cpp

@@ -0,0 +1,51 @@
+#include "common.h"
+
+namespace bpftrace {
+namespace test {
+namespace codegen {
+
+TEST(codegen, call_pid_tid)
+{
+  test("kprobe:f { @x = pid(); @y = tid() }", NAME);
+}
+
+TEST(codegen, call_pid_tid_in_child_ns)
+{
+  auto bpftrace = get_mock_bpftrace();
+  bpftrace->mock_in_init_pid_ns = false;
+  bpftrace->helper_check_level_ = 0;
+
+  test(*bpftrace, "kprobe:f { @x = pid(); @y = tid() }", NAME);
+}
+
+TEST(codegen, call_pid_tid_curr_ns)
+{
+  test("kprobe:f { @x = pid(curr_ns); @y = tid(curr_ns) }", NAME);
+}
+
+TEST(codegen, call_pid_tid_curr_ns_in_child_ns)
+{
+  auto bpftrace = get_mock_bpftrace();
+  bpftrace->mock_in_init_pid_ns = false;
+  bpftrace->helper_check_level_ = 0;
+
+  test(*bpftrace, "kprobe:f { @x = pid(curr_ns); @y = tid(curr_ns) }", NAME);
+}
+
+TEST(codegen, call_pid_tid_init)
+{
+  test("kprobe:f { @x = pid(init); @y = tid(init) }", NAME);
+}
+
+TEST(codegen, call_pid_tid_init_ns_in_child_ns)
+{
+  auto bpftrace = get_mock_bpftrace();
+  bpftrace->mock_in_init_pid_ns = false;
+  bpftrace->helper_check_level_ = 0;
+
+  test(*bpftrace, "kprobe:f { @x = pid(init); @y = tid(init) }", NAME);
+}
+
+} // namespace codegen
+} // namespace test
+} // namespace bpftrace