Add PEM command line flag for enforcing a max file size for ElfReader (#2194)

Summary: Add PEM command line flag for enforcing a max file size for
ElfReader

While memory profiling the PEM, I noticed that large binaries caused the
PEM's ElfReader code to make > 100 MiB allocations. This PR adds a cli
flag (`--elf_reader_max_file_size`) that allows memory conscious users
to prevent the PEM from parsing binaries that would cause these large
memory allocations.

Relevant Issues: N/A

Type of change: /kind feature

Test Plan: Existing tests and skaffold

Changelog Message: Added `--elf_reader_max_file_size` flag to PEM. This
flag allows memory conscious users to prevent the PEM from parsing large
binaries in order to avoid large memory allocations.

---------

Signed-off-by: Dom Del Nano <[email protected]>

Author: ddelnano

src/common/fs/fs_wrapper.cc

@@ -175,6 +175,15 @@ StatusOr<int64_t> SpaceAvailableInBytes(const std::filesystem::path& path) {
   return si.available;
 }
 
+StatusOr<std::uintmax_t> GetFileSize(const std::string& binary_path) {
+  PX_ASSIGN_OR_RETURN(const auto stat, fs::Stat(binary_path));
+  if (stat.st_size < 0) {
+    return error::Internal("stat() returned negative file size $0 for file $1", stat.st_size,
+                           binary_path);
+  }
+  return stat.st_size;
+}
+
 StatusOr<bool> IsEmpty(const std::filesystem::path& f) {
   std::error_code ec;
   bool val = std::filesystem::is_empty(f, ec);

src/common/fs/fs_wrapper.h

@@ -69,6 +69,8 @@ Status Chown(const std::filesystem::path& path, const uid_t uid, const gid_t gid
 StatusOr<struct stat> Stat(const std::filesystem::path& path);
 StatusOr<int64_t> SpaceAvailableInBytes(const std::filesystem::path& path);
 
+StatusOr<std::uintmax_t> GetFileSize(const std::string& binary_path);
+
 StatusOr<bool> IsEmpty(const std::filesystem::path& path);
 
 StatusOr<std::filesystem::path> Absolute(const std::filesystem::path& path);

src/stirling/obj_tools/elf_reader.cc

@@ -32,6 +32,10 @@
 #include "src/common/fs/fs_wrapper.h"
 #include "src/stirling/obj_tools/init.h"
 
+DEFINE_int64(elf_reader_max_file_size, 0,
+             "Maximum file size in bytes for ELF files. Default value of 0 means all files will "
+             "be parsed regardless of size.");
+
 namespace px {
 namespace stirling {
 namespace obj_tools {
@@ -46,6 +50,41 @@ struct LowercaseHex {
 };
 }  // namespace
 
+StatusOr<std::unique_ptr<ElfReader>> ElfReader::CreateImpl(
+    const std::string& binary_path, const std::filesystem::path& debug_file_dir) {
+  VLOG(1) << absl::Substitute("Creating ElfReader, [binary=$0] [debug_file_dir=$1]", binary_path,
+                              debug_file_dir.string());
+
+  auto elf_reader = std::unique_ptr<ElfReader>(new ElfReader);
+  elf_reader->binary_path_ = binary_path;
+
+  if (!elf_reader->elf_reader_.load_header_and_sections(binary_path)) {
+    return error::Internal("Can't find or process ELF file $0", binary_path);
+  }
+
+  // Check for external debug symbols.
+  Status s = elf_reader->LocateDebugSymbols(debug_file_dir);
+  if (s.ok()) {
+    std::string debug_symbols_path = elf_reader->debug_symbols_path_.string();
+
+    bool internal_debug_symbols =
+        fs::Equivalent(elf_reader->debug_symbols_path_, binary_path).ConsumeValueOr(true);
+
+    // If external debug symbols were found, load that ELF info instead.
+    if (!internal_debug_symbols) {
+      std::string debug_symbols_path = elf_reader->debug_symbols_path_.string();
+      LOG(INFO) << absl::Substitute("Found debug symbols file $0 for binary $1", debug_symbols_path,
+                                    binary_path);
+      elf_reader->elf_reader_.load_header_and_sections(debug_symbols_path);
+      return elf_reader;
+    }
+  }
+
+  // Debug symbols were either in the binary, or no debug symbols were found,
+  // so return original elf_reader.
+  return elf_reader;
+}
+
 Status ElfReader::LocateDebugSymbols(const std::filesystem::path& debug_file_dir) {
   std::string build_id;
   std::string debug_link;
@@ -158,40 +197,22 @@ Status ElfReader::LocateDebugSymbols(const std::filesystem::path& debug_file_dir
   return error::Internal("Could not find debug symbols for $0", binary_path_);
 }
 
-// TODO(oazizi): Consider changing binary_path to std::filesystem::path.
 StatusOr<std::unique_ptr<ElfReader>> ElfReader::Create(
     const std::string& binary_path, const std::filesystem::path& debug_file_dir) {
-  VLOG(1) << absl::Substitute("Creating ElfReader, [binary=$0] [debug_file_dir=$1]", binary_path,
-                              debug_file_dir.string());
-  auto elf_reader = std::unique_ptr<ElfReader>(new ElfReader);
-
-  elf_reader->binary_path_ = binary_path;
-
-  if (!elf_reader->elf_reader_.load_header_and_sections(binary_path)) {
-    return error::Internal("Can't find or process ELF file $0", binary_path);
-  }
-
-  // Check for external debug symbols.
-  Status s = elf_reader->LocateDebugSymbols(debug_file_dir);
-  if (s.ok()) {
-    std::string debug_symbols_path = elf_reader->debug_symbols_path_.string();
-
-    bool internal_debug_symbols =
-        fs::Equivalent(elf_reader->debug_symbols_path_, binary_path).ConsumeValueOr(true);
-
-    // If external debug symbols were found, load that ELF info instead.
-    if (!internal_debug_symbols) {
-      std::string debug_symbols_path = elf_reader->debug_symbols_path_.string();
-      LOG(INFO) << absl::Substitute("Found debug symbols file $0 for binary $1", debug_symbols_path,
-                                    binary_path);
-      elf_reader->elf_reader_.load_header_and_sections(debug_symbols_path);
-      return elf_reader;
+  if (FLAGS_elf_reader_max_file_size != 0) {
+    PX_ASSIGN_OR_RETURN(int64_t file_size, fs::GetFileSize(binary_path));
+    if (file_size > FLAGS_elf_reader_max_file_size) {
+      return error::Internal(
+          "File size $0 exceeds ElfReader's max file size $1. Refusing to process file", file_size,
+          FLAGS_elf_reader_max_file_size);
     }
   }
+  return CreateImpl(binary_path, debug_file_dir);
+}
 
-  // Debug symbols were either in the binary, or no debug symbols were found,
-  // so return original elf_reader.
-  return elf_reader;
+StatusOr<std::unique_ptr<ElfReader>> ElfReader::CreateUncapped(
+    const std::string& binary_path, const std::filesystem::path& debug_file_dir) {
+  return CreateImpl(binary_path, debug_file_dir);
 }
 
 StatusOr<ELFIO::section*> ElfReader::SymtabSection() {

src/stirling/obj_tools/elf_reader.h

@@ -37,6 +37,8 @@ namespace px {
 namespace stirling {
 namespace obj_tools {
 
+constexpr std::string_view kDebugFileDir = "/usr/lib/debug";
+
 class ElfReader {
  public:
   /**
@@ -50,8 +52,14 @@ class ElfReader {
    * @return error if could not setup elf reader.
    */
   static StatusOr<std::unique_ptr<ElfReader>> Create(
-      const std::string& binary_path,
-      const std::filesystem::path& debug_file_dir = "/usr/lib/debug");
+      const std::string& binary_path, const std::filesystem::path& debug_file_dir = kDebugFileDir);
+
+  /**
+   * Creates an ElfReader that does not enforce the max file size limit. This is useful for cases
+   * where the binary size is known in advance or the binary must be loaded regardless of size.
+   */
+  static StatusOr<std::unique_ptr<ElfReader>> CreateUncapped(
+      const std::string& binary_path, const std::filesystem::path& debug_file_dir = kDebugFileDir);
 
   std::filesystem::path& debug_symbols_path() { return debug_symbols_path_; }
 
@@ -207,6 +215,9 @@ class ElfReader {
   }
 
  private:
+  static StatusOr<std::unique_ptr<ElfReader>> CreateImpl(
+      const std::string& binary_path, const std::filesystem::path& debug_file_dir);
+
   ElfReader() = default;
 
   StatusOr<ELFIO::section*> SymtabSection();

src/stirling/source_connectors/socket_tracer/socket_trace_bpf_tables.cc

@@ -51,7 +51,9 @@ ConnInfoMapManager::ConnInfoMapManager(bpf_tools::BCCWrapper* bcc)
     : conn_info_map_(WrappedBCCMap<uint64_t, struct conn_info_t>::Create(bcc, "conn_info_map")),
       conn_disabled_map_(WrappedBCCMap<uint64_t, uint64_t>::Create(bcc, "conn_disabled_map")) {
   std::filesystem::path self_path = GetSelfPath().ValueOrDie();
-  auto elf_reader_or_s = obj_tools::ElfReader::Create(self_path.string());
+  // Opt out of the max file size restriction as this is a self-probe and must attach for
+  // stirling to work.
+  auto elf_reader_or_s = obj_tools::ElfReader::CreateUncapped(self_path.string());
   if (!elf_reader_or_s.ok()) {
     LOG(FATAL) << "Failed to create ElfReader for self probe";
   }