Add support of Cgroup Array in C++

examples/cpp/RandomRead.cc

@@ -19,6 +19,10 @@ const std::string BPF_PROGRAM = R"(
 #include <linux/sched.h>
 #include <uapi/linux/ptrace.h>
 
+#ifndef CGROUP_FILTER
+#define CGROUP_FILTER 0
+#endif
+
 struct urandom_read_args {
   // See /sys/kernel/debug/tracing/events/random/urandom_read/format
   uint64_t common__unused;
@@ -35,8 +39,12 @@ struct event_t {
 };
 
 BPF_PERF_OUTPUT(events);
+BPF_CGROUP_ARRAY(cgroup, 1);
 
 int on_urandom_read(struct urandom_read_args* attr) {
+  if (CGROUP_FILTER && (cgroup.check_current_task(0) != 1))
+    return 0;
+
   struct event_t event = {};
   event.pid = bpf_get_current_pid_tgid();
   bpf_get_current_comm(&event.comm, sizeof(event.comm));
@@ -72,12 +80,29 @@ void signal_handler(int s) {
 }
 
 int main(int argc, char** argv) {
+  if (argc != 1 && argc != 2) {
+    std::cerr << "USAGE: RandomRead [cgroup2_path]" << std::endl;
+    return 1;
+  }
+
+  std::vector<std::string> cflags = {};
+  if (argc == 2)
+    cflags.emplace_back("-DCGROUP_FILTER=1");
+
   bpf = new ebpf::BPF();
-  auto init_res = bpf->init(BPF_PROGRAM);
+  auto init_res = bpf->init(BPF_PROGRAM, cflags, {});
   if (init_res.code() != 0) {
     std::cerr << init_res.msg() << std::endl;
     return 1;
   }
+  if (argc == 2) {
+    auto cgroup_array = bpf->get_cgroup_array("cgroup");
+    auto update_res = cgroup_array.update_value(0, argv[1]);
+    if (update_res.code() != 0) {
+      std::cerr << update_res.msg() << std::endl;
+      return 1;
+    }
+  }
 
   auto attach_res =
       bpf->attach_tracepoint("random:urandom_read", "on_urandom_read");

src/cc/api/BPF.cc

@@ -585,6 +585,13 @@ BPFProgTable BPF::get_prog_table(const std::string& name) {
   return BPFProgTable({});
 }
 
+BPFCgroupArray BPF::get_cgroup_array(const std::string& name) {
+  TableStorage::iterator it;
+  if (bpf_module_->table_storage().Find(Path({bpf_module_->id(), name}), it))
+    return BPFCgroupArray(it->second);
+  return BPFCgroupArray({});
+}
+
 BPFStackTable BPF::get_stack_table(const std::string& name, bool use_debug_file,
                                    bool check_debug_file_crc) {
   TableStorage::iterator it;

src/cc/api/BPF.h

@@ -132,6 +132,8 @@ class BPF {
 
   BPFProgTable get_prog_table(const std::string& name);
 
+  BPFCgroupArray get_cgroup_array(const std::string& name);
+
   BPFStackTable get_stack_table(const std::string& name,
                                 bool use_debug_file = true,
                                 bool check_debug_file_crc = true);

src/cc/api/BPFTable.cc

@@ -17,6 +17,7 @@
 #include <linux/elf.h>
 #include <linux/perf_event.h>
 #include <sys/epoll.h>
+#include <fcntl.h>
 #include <unistd.h>
 #include <cerrno>
 #include <cinttypes>
@@ -30,6 +31,7 @@
 #include "bcc_exception.h"
 #include "bcc_syms.h"
 #include "common.h"
+#include "file_desc.h"
 #include "libbpf.h"
 #include "perf_reader.h"
 
@@ -433,4 +435,39 @@ BPFPerfEventArray::~BPFPerfEventArray() {
               << std::endl;
   }
 }
+
+StatusTuple BPFProgTable::update_value(const int& index, const int& prog_fd) {
+  if (!this->update(const_cast<int*>(&index), const_cast<int*>(&prog_fd)))
+    return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFProgTable::remove_value(const int& index) {
+  if (!this->remove(const_cast<int*>(&index)))
+    return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFCgroupArray::update_value(const int& index,
+                                         const int& cgroup2_fd) {
+  if (!this->update(const_cast<int*>(&index), const_cast<int*>(&cgroup2_fd)))
+    return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFCgroupArray::update_value(const int& index,
+                                         const std::string& cgroup2_path) {
+  FileDesc f(::open(cgroup2_path.c_str(), O_RDONLY | O_CLOEXEC));
+  if ((int)f < 0)
+    return StatusTuple(-1, "Unable to open %s", cgroup2_path.c_str());
+  TRY2(update_value(index, (int)f));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFCgroupArray::remove_value(const int& index) {
+  if (!this->remove(const_cast<int*>(&index)))
+    return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
 }  // namespace ebpf

src/cc/api/BPFTable.h

@@ -343,18 +343,21 @@ public:
       throw std::invalid_argument("Table '" + desc.name + "' is not a prog table");
   }
 
-  // updates an element
-  StatusTuple update_value(const int& index, const int& value) {
-    if (!this->update(const_cast<int*>(&index), const_cast<int*>(&value)))
-      return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
-    return StatusTuple(0);
-  }
+  StatusTuple update_value(const int& index, const int& prog_fd);
+  StatusTuple remove_value(const int& index);
+};
 
-  StatusTuple remove_value(const int& index) {
-    if (!this->remove(const_cast<int*>(&index)))
-      return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
-    return StatusTuple(0);
+class BPFCgroupArray : public BPFTableBase<int, int> {
+public:
+  BPFCgroupArray(const TableDesc& desc)
+      : BPFTableBase<int, int>(desc) {
+    if (desc.type != BPF_MAP_TYPE_CGROUP_ARRAY)
+      throw std::invalid_argument("Table '" + desc.name + "' is not a cgroup array");
   }
+
+  StatusTuple update_value(const int& index, const int& cgroup2_fd);
+  StatusTuple update_value(const int& index, const std::string& cgroup2_path);
+  StatusTuple remove_value(const int& index);
 };
 
 }  // namespace ebpf