From 4e5b1a008126b3da391a8301e40b8bf1c4e11dd6 Mon Sep 17 00:00:00 2001
From: Delyan Kratunov <delyank@meta.com>
Date: Thu, 20 Oct 2022 12:14:10 -0700
Subject: [PATCH 022/148] Inline map relocations into BpfProgram

Relocator is not really necessary, so clean it up and move the logic to
BpfProgram, where all relocations will live.
---
 src/CMakeLists.txt |  1 -
 src/bpfprogram.cpp | 45 +++++++++++++++++++++++++++++----------
 src/bpfprogram.h   |  9 +++++---
 src/bpftrace.cpp   |  5 ++---
 src/relocator.cpp  | 52 ----------------------------------------------
 src/relocator.h    | 32 ----------------------------
 6 files changed, 42 insertions(+), 102 deletions(-)
 delete mode 100644 src/relocator.cpp
 delete mode 100644 src/relocator.h

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 9fb3bdaa..5024d24c 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -32,7 +32,6 @@ add_library(runtime
   probe_matcher.cpp
   procmon.cpp
   printf.cpp
-  relocator.cpp
   resolve_cgroupid.cpp
   required_resources.cpp
   struct.cpp
diff --git a/src/bpfprogram.cpp b/src/bpfprogram.cpp
index a3b7e214..d13341c5 100644
--- a/src/bpfprogram.cpp
+++ b/src/bpfprogram.cpp
@@ -1,7 +1,5 @@
 #include "bpfprogram.h"
 
-#include "relocator.h"
-
 #include <optional>
 #include <stdexcept>
 #include <tuple>
@@ -11,19 +9,19 @@ namespace bpftrace {
 std::optional<BpfProgram> BpfProgram::CreateFromBytecode(
     const BpfBytecode &bytecode,
     const std::string &name,
-    BPFtrace &bpftrace)
+    MapManager &maps)
 {
   if (bytecode.find(name) != bytecode.end())
   {
-    return BpfProgram(bytecode, name, bpftrace);
+    return BpfProgram(bytecode, name, maps);
   }
   return std::nullopt;
 }
 
 BpfProgram::BpfProgram(const BpfBytecode &bytecode,
                        const std::string &name,
-                       BPFtrace &bpftrace)
-    : bytecode_(bytecode), bpftrace_(bpftrace), name_(name), code_()
+                       MapManager &maps)
+    : bytecode_(bytecode), maps_(maps), name_(name), code_()
 {
 }
 
@@ -39,11 +37,36 @@ void BpfProgram::assemble()
 
   code_ = bytecode_.at(name_);
 
-  // Perform relocations on the copy of the code for this particular program.
-  auto relocator = Relocator(std::make_tuple(code_.data(), code_.size()),
-                             bpftrace_);
-  if (relocator.relocate())
-    throw std::runtime_error("Could not relocate program, see log");
+  relocateMaps();
+}
+
+void BpfProgram::relocateMaps()
+{
+  struct bpf_insn *insns = reinterpret_cast<struct bpf_insn *>(code_.data());
+  size_t insn_cnt = code_.size() / sizeof(struct bpf_insn);
+  for (uintptr_t i = 0; i < insn_cnt; ++i)
+  {
+    struct bpf_insn *insn = &insns[i];
+
+    // Relocate mapid -> mapfd
+    //
+    // This relocation keeps codegen independent of runtime state (such as FD
+    // numbers). This helps make codegen tests more reliable and enables
+    // features such as AOT compilation.
+    if (insn->code == BPF_DW && (insn->src_reg == BPF_PSEUDO_MAP_FD ||
+                                 insn->src_reg == BPF_PSEUDO_MAP_VALUE))
+    {
+      auto mapid = insn->imm;
+      auto map = maps_[mapid];
+      if (map)
+        insn->imm = static_cast<int32_t>((*map)->mapfd_);
+      else
+        throw std::runtime_error(std::string("Unknown map id ") +
+                                 std::to_string(mapid));
+
+      ++i; // ldimm64 is 2 insns wide
+    }
+  }
 }
 
 } // namespace bpftrace
diff --git a/src/bpfprogram.h b/src/bpfprogram.h
index d595400e..f47cf2ad 100644
--- a/src/bpfprogram.h
+++ b/src/bpfprogram.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include "mapmanager.h"
+
 #include <cstdint>
 #include <optional>
 #include <string>
@@ -19,7 +21,7 @@ public:
   static std::optional<BpfProgram> CreateFromBytecode(
       const BpfBytecode &bytecode,
       const std::string &name,
-      BPFtrace &bpftrace);
+      MapManager &maps);
 
   void assemble();
 
@@ -33,11 +35,12 @@ public:
 private:
   explicit BpfProgram(const BpfBytecode &bytecode,
                       const std::string &name,
-                      BPFtrace &bpftrace);
+                      MapManager &bpftrace);
 
+  void relocateMaps();
 
   const BpfBytecode &bytecode_;
-  BPFtrace &bpftrace_;
+  MapManager &maps_;
   std::string name_;
   std::vector<uint8_t> code_;
 };
diff --git a/src/bpftrace.cpp b/src/bpftrace.cpp
index f4bc5379..ee7e5152 100644
--- a/src/bpftrace.cpp
+++ b/src/bpftrace.cpp
@@ -33,7 +33,6 @@
 #include "bpftrace.h"
 #include "log.h"
 #include "printf.h"
-#include "relocator.h"
 #include "resolve_cgroupid.h"
 #include "triggers.h"
 #include "utils.h"
@@ -942,12 +941,12 @@ std::vector<std::unique_ptr<AttachedProbe>> BPFtrace::attach_probe(
                                               probe.index,
                                               usdt_location_idx);
 
-  auto program = BpfProgram::CreateFromBytecode(bytecode, name, *this);
+  auto program = BpfProgram::CreateFromBytecode(bytecode, name, maps);
   if (!program)
   {
     auto orig_program = BpfProgram::CreateFromBytecode(bytecode,
                                                        orig_name,
-                                                       *this);
+                                                       maps);
     if (orig_program)
       program.emplace(std::move(*orig_program));
   }
diff --git a/src/relocator.cpp b/src/relocator.cpp
deleted file mode 100644
index 8c938b2b..00000000
--- a/src/relocator.cpp
+++ /dev/null
@@ -1,52 +0,0 @@
-#include "relocator.h"
-
-#include "bpftrace.h"
-#include "log.h"
-
-namespace libbpf {
-#include "libbpf/bpf.h"
-} // namespace libbpf
-
-namespace bpftrace {
-
-Relocator::Relocator(std::tuple<uint8_t *, uintptr_t> func, BPFtrace &bpftrace)
-    : insns_(reinterpret_cast<struct bpf_insn *>(std::get<0>(func))),
-      nr_(std::get<1>(func) / sizeof(struct bpf_insn)),
-      bpftrace_(bpftrace)
-{
-}
-
-int Relocator::relocate()
-{
-  for (uintptr_t i = 0; i < nr_; ++i)
-  {
-    struct bpf_insn *insn = &insns_[i];
-
-    // Relocate mapid -> mapfd
-    //
-    // This relocation keeps codegen independent of runtime state (such as FD
-    // numbers). This helps make codegen tests more reliable and enables
-    // features such as AOT compilation.
-    if (insn->code == BPF_DW && (insn->src_reg == BPF_PSEUDO_MAP_FD ||
-                                 insn->src_reg == BPF_PSEUDO_MAP_VALUE))
-    {
-      auto mapid = insn->imm;
-      auto map = bpftrace_.maps[mapid];
-      if (map)
-      {
-        insn->imm = static_cast<int32_t>((*map)->mapfd_);
-      }
-      else
-      {
-        LOG(ERROR) << "Failed to relocate mapid=" << mapid << ": ID unknown";
-        return 1;
-      }
-
-      ++i; // ldimm64 is 2 insns wide
-    }
-  }
-
-  return 0;
-}
-
-} // namespace bpftrace
diff --git a/src/relocator.h b/src/relocator.h
deleted file mode 100644
index c4f94033..00000000
--- a/src/relocator.h
+++ /dev/null
@@ -1,32 +0,0 @@
-#pragma once
-
-#include <cstdint>
-#include <tuple>
-
-struct bpf_insn;
-namespace bpftrace {
-class BPFtrace;
-
-// Relocates BPF bytecode prior to being loaded into the kernel
-//
-// Relocations are needed to customize the bytecode to the host the
-// bytecode is going to be run on.
-class Relocator
-{
-public:
-  Relocator(std::tuple<uint8_t *, uintptr_t> func, BPFtrace &bpftrace);
-  ~Relocator() = default;
-
-  // Perform relocations
-  //
-  // Note this function will modify the instructions passed into the
-  // constructor
-  int relocate();
-
-private:
-  struct bpf_insn *insns_;
-  uintptr_t nr_;
-  BPFtrace &bpftrace_;
-};
-
-} // namespace bpftrace
-- 
2.33.0