1 // Copyright (c) Meta Platforms, Inc. and affiliates. 2 // All rights reserved. 3 // 4 // This source code is licensed under the BSD-style license found in the 5 // LICENSE file in the root directory of this source tree. 6 7 #pragma once 8 9 #include <elf.h> 10 #include <fcntl.h> 11 #include <fmt/format.h> 12 #include <sys/mman.h> 13 #include <sys/stat.h> 14 #include <torch/csrc/profiler/unwind/lexer.h> 15 #include <torch/csrc/profiler/unwind/unwind_error.h> 16 #include <unistd.h> 17 #include <cerrno> 18 #include <cstdio> 19 #include <cstring> 20 21 namespace torch::unwind { 22 23 struct Section { 24 char* data = nullptr; 25 size_t size = 0; stringSection26 const char* string(size_t offset) { 27 return lexer(offset).readCString(); 28 } lexerSection29 CheckedLexer lexer(size_t offset) { 30 return CheckedLexer(data + offset, data, data + size); 31 } 32 }; 33 34 /// Memory maps a file into the address space read-only, and manages the 35 /// lifetime of the mapping. Here are a few use cases: 36 /// 1. Used in the loader to read in initial image, and to inspect 37 // ELF files for dependencies before callling dlopen. 38 /// 39 /// 2. Used in unity to load the elf file. 40 struct MemFile { MemFileMemFile41 explicit MemFile(const char* filename_) 42 : fd_(open(filename_, O_RDONLY)), name_(filename_) { 43 UNWIND_CHECK( 44 fd_ != -1, "failed to open {}: {}", filename_, strerror(errno)); 45 // NOLINTNEXTLINE 46 struct stat s; 47 if (-1 == fstat(fd_, &s)) { 48 close(fd_); // destructors don't run during exceptions 49 UNWIND_CHECK(false, "failed to stat {}: {}", filename_, strerror(errno)); 50 } 51 n_bytes_ = s.st_size; 52 UNWIND_CHECK( 53 n_bytes_ > sizeof(Elf64_Ehdr), "empty shared library: {}", filename_); 54 mem_ = (char*)mmap(nullptr, n_bytes_, PROT_READ, MAP_SHARED, fd_, 0); 55 if (MAP_FAILED == mem_) { 56 close(fd_); 57 UNWIND_CHECK(false, "failed to mmap {}: {}", filename_, strerror(errno)); 58 } 59 ehdr_ = (Elf64_Ehdr*)mem_; 60 #define ELF_CHECK(cond) UNWIND_CHECK(cond, "not an ELF file: {}", filename_) 61 ELF_CHECK(ehdr_->e_ident[EI_MAG0] == ELFMAG0); 62 ELF_CHECK(ehdr_->e_ident[EI_MAG1] == ELFMAG1); 63 ELF_CHECK(ehdr_->e_ident[EI_MAG2] == ELFMAG2); 64 ELF_CHECK(ehdr_->e_ident[EI_MAG3] == ELFMAG3); 65 ELF_CHECK(ehdr_->e_ident[EI_CLASS] == ELFCLASS64); 66 ELF_CHECK(ehdr_->e_ident[EI_VERSION] == EV_CURRENT); 67 ELF_CHECK(ehdr_->e_version == EV_CURRENT); 68 ELF_CHECK(ehdr_->e_machine == EM_X86_64); 69 #undef ELF_CHECK 70 UNWIND_CHECK( 71 ehdr_->e_shoff + sizeof(Elf64_Shdr) * ehdr_->e_shnum <= n_bytes_, 72 "invalid section header table {} {} {}", 73 ehdr_->e_shoff + sizeof(Elf64_Shdr) * ehdr_->e_shnum, 74 n_bytes_, 75 ehdr_->e_shnum); 76 shdr_ = (Elf64_Shdr*)(mem_ + ehdr_->e_shoff); 77 UNWIND_CHECK( 78 ehdr_->e_shstrndx < ehdr_->e_shnum, "invalid strtab section offset"); 79 auto& strtab_hdr = shdr_[ehdr_->e_shstrndx]; 80 strtab_ = getSection(strtab_hdr); 81 } 82 83 MemFile(const MemFile&) = delete; 84 MemFile& operator=(const MemFile&) = delete; dataMemFile85 [[nodiscard]] const char* data() const { 86 return (const char*)mem_; 87 } 88 89 /// Returns whether or not the file descriptor 90 /// of the underlying file is valid. validMemFile91 int valid() { 92 return fcntl(fd_, F_GETFD) != -1 || errno != EBADF; 93 } 94 ~MemFileMemFile95 ~MemFile() { 96 if (mem_) { 97 munmap((void*)mem_, n_bytes_); 98 } 99 if (fd_ >= 0) { 100 close(fd_); 101 } 102 } 103 104 /// Returns the size of the underlying file defined by the `MemFile` sizeMemFile105 size_t size() { 106 return n_bytes_; 107 } fdMemFile108 [[nodiscard]] int fd() const { 109 return fd_; 110 } 111 getSectionMemFile112 Section getSection(const Elf64_Shdr& shdr) { 113 UNWIND_CHECK(shdr.sh_offset + shdr.sh_size <= n_bytes_, "invalid section"); 114 return Section{mem_ + shdr.sh_offset, shdr.sh_size}; 115 } 116 getSectionMemFile117 Section getSection(const char* name, bool optional) { 118 for (int i = 0; i < ehdr_->e_shnum; i++) { 119 if (strcmp(strtab_.string(shdr_[i].sh_name), name) == 0) { 120 return getSection(shdr_[i]); 121 } 122 } 123 UNWIND_CHECK(optional, "{} has no section {}", name_, name); 124 return Section{nullptr, 0}; 125 } 126 strtabMemFile127 Section strtab() { 128 return strtab_; 129 } 130 131 private: 132 template <typename T> loadMemFile133 T* load(size_t offset) { 134 UNWIND_CHECK(offset < n_bytes_, "out of range"); 135 return (T*)(mem_ + offset); 136 } 137 int fd_; 138 char* mem_{nullptr}; 139 size_t n_bytes_{0}; 140 std::string name_; 141 Elf64_Ehdr* ehdr_; 142 Elf64_Shdr* shdr_; 143 Section strtab_ = {nullptr, 0}; 144 }; 145 146 } // namespace torch::unwind 147