add elf parser (#192)

.github/workflows/clang-format-check.yml

@@ -41,7 +41,7 @@ jobs:
       - name: Install clang-format
         uses: aminya/setup-cpp@v1
         with:
-          clangformat: 18.1.7
+          clangformat: 13.0.1
 
       - name: Run clang-format
         env:

include/elf_parser.h

@@ -0,0 +1,156 @@
+#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::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,75 +18,66 @@ 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
+
 template <typename T>
 struct str2value_impl {
-    void operator()(T& value, const char* str, size_t len = std::string::npos,
+    void operator()(T& value, adt::StringRef str,
                     std::enable_if_t<std::is_same<T, std::string>::value ||
                                      std::is_integral<T>::value>* = nullptr) {
         std::stringstream ss;
-        if (len != std::string::npos) {
-            ss << std::string(str, str + len);
-        } else {
-            ss << str;
-        }
+        ss << str;
         ss >> value;
     }
 };
 
 // TODO: return bool value to check parse result
 template <>
 struct str2value_impl<int> {
-    void operator()(int& value, const char* cstr,
-                    size_t len = std::string::npos);
+    void operator()(int& value, adt::StringRef str);
 };
 
 template <typename T>
 struct str2value {
-    T operator()(const char* str, size_t len = std::string::npos) {
+    T operator()(adt::StringRef str) {
         T ret;
-        str2value_impl<T>()(ret, str, len);
+        str2value_impl<T>()(ret, str);
         return ret;
     }
 };
 
 template <typename K, typename V>
 struct str2value_impl<std::pair<K, V>> {
-    void operator()(std::pair<K, V>& pair, const char* str,
-                    size_t len = std::string::npos) {
+    void operator()(std::pair<K, V>& pair, adt::StringRef str) {
         size_t i = 0;
-        for (; i < len && str[i] != '\0'; ++i) {
+        for (; i < str.size() && str[i] != '\0'; ++i) {
             if (str[i] == '=') {
-                pair.first = str2value<K>()(str, i);
-                pair.second = str2value<V>()(str + i + 1);
+                pair.first = str2value<K>()(str.slice(0, i));
+                pair.second = str2value<V>()(str.slice(i + 1));
                 break;
             }
         }
-        if (i == '\0' || i == len) {
-            pair.first = str2value<K>()(str, i);
-        }
     }
 };
 
 template <typename V>
 struct str2value_impl<std::vector<V>> {
-    void operator()(std::vector<V>& vec, const char* str,
-                    size_t len = std::string::npos) {
+    void operator()(std::vector<V>& vec, adt::StringRef str) {
         size_t i = 0, j = 0;
-        for (; j < len && str[j] != '\0'; ++j) {
+        for (; j < str.size() && str[j] != '\0'; ++j) {
             if (str[j] == ',') {
-                vec.push_back(str2value<V>()(str + i, j));
+                vec.push_back(str2value<V>()(str.slice(i, j)));
                 i = j + 1;
             }
         }
-        vec.push_back(str2value<V>()(str + i, j - i));
+        vec.push_back(str2value<V>()(str.slice(i, j)));
     }
 };
 
 template <typename T>
-T get_env_value(const char* str,
+T get_env_value(adt::StringRef str,
                 std::__void_t<decltype(std::declval<std::istringstream>() >>
                                        std::declval<T&>())>* = nullptr) {
-    auto env_value_str = std::getenv(str);
+    auto env_value_str = std::getenv(str.data());
     if (!env_value_str) {
         return {};
     }
@@ -96,8 +87,8 @@ T get_env_value(const char* str,
 template <typename T, typename K = typename T::first_type,
           typename V = typename T::second_type>
 typename std::enable_if<std::is_same<T, std::pair<K, V>>::value, T>::type
-get_env_value(const char* str) {
-    auto env_value_str = std::getenv(str);
+get_env_value(adt::StringRef str) {
+    auto env_value_str = std::getenv(str.data());
     if (!env_value_str) {
         return {};
     }
@@ -106,8 +97,8 @@ get_env_value(const char* str) {
 
 template <typename T, typename V = typename T::value_type>
 typename std::enable_if<std::is_same<T, std::vector<V>>::value, T>::type
-get_env_value(const char* str) {
-    auto env_value_str = std::getenv(str);
+get_env_value(adt::StringRef str) {
+    auto env_value_str = std::getenv(str.data());
     if (!env_value_str) {
         return {};
     }
@@ -116,8 +107,18 @@ get_env_value(const char* str) {
 
 template <typename T>
 inline typename std::enable_if<std::is_same<T, const char*>::value, T>::type
-get_env_value(const char* str) {
-    return std::getenv(str);
+get_env_value(adt::StringRef str) {
+    return std::getenv(str.data());
+}
+
+template <typename T>
+inline typename std::enable_if<std::is_same<T, adt::StringRef>::value, T>::type
+get_env_value(adt::StringRef str) {
+    auto env_value_str = std::getenv(str.data());
+    if (!env_value_str) {
+        return {};
+    }
+    return adt::StringRef(env_value_str);
 }
 
 }  // namespace hook

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;
     }
 };
@@ -120,20 +122,27 @@ std::string args_to_string(Args... args) {
     return ss.str();
 }
 
-#define IF_ENABLE_LOG_TRACE_AND_ARGS(func)                                    \
-    do {                                                                      \
-        int ctrl = enable_log_backtrace(func);                                \
-        if (ctrl) {                                                           \
-            if (ctrl & 0b10) {                                                \
-                MLOG(TRACE, WARN) << func << ": " << args_to_string(args...); \
-            }                                                                 \
-            if (ctrl & 0b01) {                                                \
-                trace::CallFrames callFrames;                                 \
-                callFrames.CollectNative();                                   \
-                callFrames.CollectPython();                                   \
-                MLOG(TRACE, WARN) << func << " with frame:\n" << callFrames;  \
-            }                                                                 \
-        }                                                                     \
+#define IF_ENABLE_LOG_TRACE_AND_ARGS(func)                                   \
+    do {                                                                     \
+        int ctrl = enable_log_backtrace((func));                             \
+        if (ctrl) {                                                          \
+            if (ctrl & 0b10) {                                               \
+                auto parser_func =                                           \
+                    HookRuntimeContext::instance().lookUpArgsParser((func)); \
+                MLOG(TRACE, WARN)                                            \
+                    << func << ": "                                          \
+                    << (parser_func                                          \
+                            ? reinterpret_cast<std::string (*)(Args...)>(    \
+                                  parser_func)(args...)                      \
+                            : args_to_string(args...));                      \
+            }                                                                \
+            if (ctrl & 0b01) {                                               \
+                trace::CallFrames callFrames;                                \
+                callFrames.CollectNative();                                  \
+                callFrames.CollectPython();                                  \
+                MLOG(TRACE, WARN) << func << " with frame:\n" << callFrames; \
+            }                                                                \
+        }                                                                    \
     } while (0)
 
 template <typename StrT, size_t UniqueId, typename R, typename... Args>
@@ -276,6 +285,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 +350,19 @@ struct __HookFeature : public HookFeatureBase {
         return newFuncGenerator(libName, symName.c_str(), newFunc);
     }
 
+    __HookFeature& setGetNewCallback(
+        const std::function<void(const OriginalInfo&)>& getNewCallback) {
+        getNewCallback_ = getNewCallback;
+        return *this;
+    }
+
+    template <typename... Args>
+    __HookFeature& setArgsParser(std::string (*parser)(Args...)) {
+        HookRuntimeContext::instance().argsParserMap().emplace(
+            symName, reinterpret_cast<void*>(parser));
+        return *this;
+    }
+
     std::function<void*(size_t)> findUniqueFunc;
     WrapFuncGenerator newFuncGenerator;
 };
@@ -372,6 +395,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);
     }
 };