add elf parser

include/elf_parser.h

@@ -0,0 +1,157 @@
+#include <elf.h>
+#include <link.h>
+
+#include <fstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "logger/logger.h"
+#include "macro.h"
+
+namespace hook {
+
+class HOOK_API CachedSymbolTable {
+   public:
+    struct StringRefIterator {
+        StringRefIterator(const char *str);
+
+        StringRefIterator &operator++() &;
+        StringRefIterator operator++(int) &;
+
+        adt::StringRef operator*();
+
+        bool operator==(const StringRefIterator &other) const;
+        bool operator!=(const StringRefIterator &other) const;
+        const void *data() const;
+
+       private:
+        const char *str_;
+    };
+
+    template <typename T>
+    struct OwnerBuf {
+        OwnerBuf() = default;
+        OwnerBuf(const OwnerBuf &other) = delete;
+        OwnerBuf &operator=(const OwnerBuf &other) = delete;
+
+        OwnerBuf(OwnerBuf &&other) { *this = std::move(other); }
+        OwnerBuf &operator=(OwnerBuf &&other) {
+            std::swap(buf_, other.buf_);
+            std::swap(size_, other.size_);
+            return *this;
+        }
+
+        static OwnerBuf alloc(size_t size) {
+            OwnerBuf buf;
+            buf.buf_ = reinterpret_cast<T *>(malloc(size * sizeof(T)));
+            buf.size_ = size;
+            return buf;
+        }
+
+        T &operator[](size_t index) {
+            assert(index < size_);
+            assert(buf_);
+            return buf_[index];
+        }
+
+        T *data() const { return buf_; }
+
+        size_t size() const { return size_; }
+
+        std::tuple<T *, size_t> release() {
+            auto result = std::make_tuple(buf_, size_);
+            buf_ = 0;
+            size_ = 0;
+            return result;
+        }
+
+        template <typename N>
+        friend class OwnerBuf;
+
+        template <typename N>
+        OwnerBuf<N> cast() {
+            OwnerBuf<N> result;
+            auto [buf, size] = release();
+            result.buf_ = reinterpret_cast<N *>(buf);
+            result.size_ = size * sizeof(T) / sizeof(N);
+            return result;
+        }
+
+        ~OwnerBuf() {
+            if (buf_) {
+                free(buf_);
+#ifndef NDEBUG
+                buf_ = nullptr;
+                size_ = 0;
+#endif
+            }
+        }
+
+       private:
+        T *buf_ = nullptr;
+        size_t size_ = 0;
+    };
+
+    StringRefIterator strtab_begin(const char *str) const;
+
+    StringRefIterator strtab_end(const char *str) const;
+
+    std::tuple<StringRefIterator, StringRefIterator> strtab_range(
+        const char *str, size_t size) const {
+        return std::make_tuple(strtab_begin(str), strtab_begin(str + size));
+    }
+
+    CachedSymbolTable(const std::string &name, const void *base_address,
+                      const std::vector<std::string> &section_names = {});
+
+    void move_to_section_header();
+
+    void move_to_section_header(size_t index);
+
+    adt::StringRef getSectionName(size_t index) const;
+    size_t find_section(adt::StringRef name) const;
+
+    void load_symbol_table();
+    void parse_section_header();
+    void parse_named_section();
+
+    template <typename T>
+    OwnerBuf<T> load_section_data(adt::StringRef name) {
+        auto buf = load_section_data(name);
+        return buf.cast<T>();
+    }
+
+    OwnerBuf<char> load_section_data(adt::StringRef name);
+
+    const std::string &lookUpSymbol(const void *func) const;
+
+    const std::unordered_map<size_t, std::string> &getSymbolTable() const {
+        return symbol_table;
+    }
+
+    size_t min_addrtess() const { return min_address_; }
+    size_t max_addrtess() const { return max_address_; }
+
+    const std::vector<ElfW(Shdr)> &sections() const { return sections_; }
+
+   private:
+    std::string libName;
+    std::ifstream ifs;
+    ElfW(Ehdr) elf_header;
+    std::vector<char> section_header_str;
+    std::vector<ElfW(Shdr)> sections_;
+    std::vector<std::string> strtab;
+    std::unordered_map<size_t, std::string> symbol_table;
+    const void *base_address;
+    size_t min_address_ = -1;
+    size_t max_address_ = 0;
+    std::vector<std::string> section_names;
+};
+
+CachedSymbolTable *createSymbolTable(const std::string &lib,
+                                     const void *address);
+
+CachedSymbolTable *getSymbolTable(const std::string &lib);
+
+}  // namespace hook

include/env_util.h

@@ -18,6 +18,8 @@ namespace hook {
 // if we use the function template we need implement all of the str2value_impl
 // functions to support more type, but sometimes we need implement the
 // str2value_impl near the type define
+
+// TODO: replace C style const char* and length with StringRef
 template <typename T>
 struct str2value_impl {
     void operator()(T& value, const char* str, size_t len = std::string::npos,
@@ -57,7 +59,7 @@ struct str2value_impl<std::pair<K, V>> {
         for (; i < len && str[i] != '\0'; ++i) {
             if (str[i] == '=') {
                 pair.first = str2value<K>()(str, i);
-                pair.second = str2value<V>()(str + i + 1);
+                pair.second = str2value<V>()(str + i + 1, len - i - 1);
                 break;
             }
         }

include/hook.h

@@ -29,6 +29,7 @@ void install_hook();
 struct OriginalInfo {
     const char* libName = nullptr;
     const void* basePtr = nullptr;
+    const void* baseHeadPtr = nullptr;
     void* relaPtr = nullptr;
     void* oldFuncPtr = nullptr;
     void** pltTablePtr = nullptr;
@@ -40,6 +41,7 @@ struct OriginalInfo {
         relaPtr = info.relaPtr;
         oldFuncPtr = info.oldFuncPtr;
         pltTablePtr = info.pltTablePtr;
+        baseHeadPtr = info.baseHeadPtr;
         return *this;
     }
 };
@@ -276,6 +278,7 @@ struct HookFeatureBase {
     void* newFunc;
     void** oldFunc;
     std::function<bool(void)> filter_;
+    std::function<void(const OriginalInfo&)> getNewCallback_;
 };
 
 using WrapFuncGenerator = std::function<void*(const char*, const char*, void*)>;
@@ -340,6 +343,12 @@ struct __HookFeature : public HookFeatureBase {
         return newFuncGenerator(libName, symName.c_str(), newFunc);
     }
 
+    __HookFeature& setGetNewCallback(
+        const std::function<void(const OriginalInfo&)>& getNewCallback) {
+        getNewCallback_ = getNewCallback;
+        return *this;
+    }
+
     std::function<void*(size_t)> findUniqueFunc;
     WrapFuncGenerator newFuncGenerator;
 };
@@ -372,6 +381,9 @@ struct MemberDetector<DerivedT,
         // TODO: if std::get<2>(*iter) is a pointer and it's point to
         // std::get<1>(*iter) then there will return nullptr
         *iter->oldFunc = info.oldFuncPtr;
+        if (iter->getNewCallback_) {
+            iter->getNewCallback_(info);
+        }
         return iter->getNewFunc(info.libName);
     }
 };

include/logger/StringRef.h

@@ -24,7 +24,8 @@ class StringRef {
 
     StringRef(const char* str) : size_(strlen(str)), str_(str) {}
 
-    StringRef(const char* str, size_t size) : size_(size), str_(str) {}
+    StringRef(const char* str, size_t size)
+        : size_(size != std::string::npos ? size : strlen(str_)), str_(str) {}
 
     StringRef(const_iterator begin, const_iterator end)
         : size_(std::distance(begin, end)), str_(begin) {}
@@ -46,7 +47,9 @@ class StringRef {
 
     size_t size() const { return size_; }
 
-    std::string str() const { return std::string(str_, str_ + size_); }
+    std::string str() const {
+        return size_ != 0 ? std::string(str_, str_ + size_) : std::string();
+    }
 
     const char* data() const { return str_; }
 

include/logger/logger.h

@@ -40,10 +40,13 @@ enum class LogModule { profile, trace, hook, python, memory, debug, last };
 class LogModuleHelper {
    public:
     static auto& enum_strs() {
-        static std::array<const char*, 6> strs = {"PROFILE", "TRACE", "HOOK",
-                                                  "PYTHON", "MEMORY", "LAST"};
+        static std::array<const char*, 7> strs = {
+            "PROFILE", "TRACE", "HOOK", "PYTHON", "MEMORY", "DEBUG", "LAST"};
+        static_assert(sizeof(strs) / sizeof(const char*) ==
+                      static_cast<size_t>(LogModule::last) + 1);
         return strs;
     }
+
     static auto begin() { return enum_strs().begin(); }
     static auto end() { return enum_strs().end(); }
 

include/util.h

@@ -0,0 +1,7 @@
+#include <string>
+
+namespace hook {
+
+std::string prettyFormatSize(size_t bytes);
+
+}  // namespace hook

lib/CMakeLists.txt

@@ -12,12 +12,14 @@
 find_package(PythonLibs REQUIRED)
 
 add_library(cuda_mock STATIC
+    util.cpp
     env_util.cpp
     GlobalVarMgr.cpp
     logger.cpp
     env_mgr.cpp
     backtrace.cpp
     statistic.cpp
+    elf_parser.cpp
     hook_context.cpp
     hook.cpp
     cuda_op_tracer.cpp