LSPlt/lsplt/src/main/jni/lsplt.cc

#include "include/lsplt.hpp"

#include <regex.h>
#include <sys/mman.h>

#include <array>
#include <cinttypes>
#include <list>
#include <map>
#include <mutex>
#include <vector>

#include "elf_util.hpp"

namespace {
struct BaseInfo {
    regex_t regex;
    ino_t inode;
    std::string symbol;
    ~BaseInfo() { regfree(&regex); }
};

struct RegisterInfo : BaseInfo {
    void *callback;
    void **backup;
    ~RegisterInfo() { regfree(&regex); }
};

struct IgnoreInfo : BaseInfo {};

struct MapInfo {
    std::string path;
    ino_t inode;
    uintptr_t start;
    uintptr_t end;
    std::map<uintptr_t, uintptr_t> hooks;
    uintptr_t backup;
    std::unique_ptr<Elf> elf;
    [[nodiscard]] bool Match(const BaseInfo &info) const {
        return (info.inode != 0 && info.inode == inode) ||
               (info.inode == 0 && regexec(&info.regex, path.c_str(), 0, nullptr, 0) == 0);
    }
};

class MapInfos : public std::map<uintptr_t, MapInfo, std::greater<>> {
public:
    static MapInfos ScanMapInfo() {
        constexpr static auto kPermLength = 5;
        constexpr static auto kMapEntry = 5;
        MapInfos info;
        auto maps =
            std::unique_ptr<FILE, decltype(&fclose)>{fopen("/proc/self/maps", "r"), &fclose};
        if (maps) {
            char *line = nullptr;
            size_t len = 0;
            ssize_t read;
            while ((read = getline(&line, &len, maps.get())) != -1) {
                uintptr_t start = 0;
                uintptr_t end = 0;
                uintptr_t off = 0;
                ino_t inode = 0;
                std::array<char, kPermLength> perm{'\0'};
                int path_off;
                if (sscanf(line, "%" PRIxPTR "-%" PRIxPTR " %4s %" PRIxPTR " %*x:%*x %lu %n%*s",
                           &start, &end, perm.data(), &off, &inode, &path_off) != kMapEntry) {
                    continue;
                }
                // we basically only care about r--p entry
                // and for offset == 0 it's an ELF header
                // and for offset != 0 it's what we hook
                // if (perm[0] != 'r') continue;
                if (perm[3] != 'p') continue;
                if (perm[2] == 'x') continue;
                // if (off != 0) continue;
                while (path_off < read && isspace(line[path_off])) path_off++;
                if (path_off >= read) continue;
                std::string path{line + path_off};
                if (path.empty()) continue;
                if (path[0] == '[') continue;

                info.emplace(start, MapInfo{std::move(path), inode, start, end, {}, 0, nullptr});
            }
            free(line);
        }
        return info;
    }

    // fiter out ignored
    void Filter(const std::list<RegisterInfo> &register_info,
                const std::list<IgnoreInfo> &ignore_info) {
        for (auto iter = begin(); iter != end();) {
            const auto &info = iter->second;
            bool matched = false;
            for (const auto &reg : register_info) {
                if (!info.Match(reg)) {
                    continue;
                }
                bool ignored = false;
                for (const auto &ign : ignore_info) {
                    if (!info.Match(ign)) {
                        continue;
                    }
                    if (ign.symbol == reg.symbol) {
                        ignored = true;
                        break;
                    }
                }
                if (!ignored) {
                    matched = true;
                    break;
                }
            }
            if (matched) {
                ++iter;
            } else {
                iter = erase(iter);
            }
        }
    }

    void Merge(MapInfos &old) {
        // merge with old map info
        for (auto &info : old) {
            if (info.second.backup) {
                erase(info.second.backup);
            }
            if (auto iter = find(info.first); iter != end()) {
                iter->second = std::move(info.second);
            } else {
                emplace(info.first, std::move(info.second));
            }
        }
    }

    bool DoHook(uintptr_t addr, uintptr_t callback, uintptr_t *backup) {
        auto iter = lower_bound(addr);
        if (iter == end()) return false;
        // iter.first < addr
        auto &info = iter->second;
        if (info.end <= addr) return false;
        if (!iter->second.backup) {
            auto len = info.end - info.start;
            // let os find a suitable address
            auto *backup_addr = mmap(nullptr, len, PROT_NONE, MAP_PRIVATE | MAP_ANON, -1, 0);
            if (backup_addr == MAP_FAILED) return false;
            if (auto *new_addr = mremap(reinterpret_cast<void *>(info.start), len, len,
                                        MREMAP_FIXED | MREMAP_MAYMOVE, backup_addr);
                new_addr == MAP_FAILED || new_addr != backup_addr) {
                return false;
            }
            if (auto *new_addr =
                    mmap(reinterpret_cast<void *>(info.start), len, PROT_READ | PROT_WRITE,
                         MAP_PRIVATE | MAP_FIXED | MAP_ANON, -1, 0);
                new_addr == MAP_FAILED) {
                return false;
            }
            memcpy(reinterpret_cast<void *>(info.start), backup_addr, len);
            info.backup = reinterpret_cast<uintptr_t>(backup_addr);
        }
        auto *the_addr = reinterpret_cast<uintptr_t *>(addr);
        *backup = *the_addr;
        *the_addr = callback;
        if (auto hook_iter = info.hooks.find(addr); hook_iter != info.hooks.end()) {
            if (hook_iter->second == callback) info.hooks.erase(hook_iter);
        } else {
            info.hooks.emplace(addr, *backup);
        }
        if (info.hooks.empty()) {
            auto len = info.end - info.start;
            if (auto *new_addr =
                    mremap(reinterpret_cast<void *>(info.backup), len, len,
                           MREMAP_FIXED | MREMAP_MAYMOVE, reinterpret_cast<void *>(info.start));
                new_addr == MAP_FAILED || reinterpret_cast<uintptr_t>(new_addr) != info.start) {
                return false;
            }
            info.backup = 0;
        }
        return true;
    }

    bool DoHook(std::list<RegisterInfo> &register_info,
                const std::list<IgnoreInfo> &ignore_info) {
        bool res = true;
        for (auto &[_, info] : *this) {
            for (auto iter = register_info.begin(); iter != register_info.end();) {
                const auto &reg = *iter;
                if (!info.Match(reg)) {
                    continue;
                }
                bool ignored = false;
                for (const auto &ign : ignore_info) {
                    if (!info.Match(ign)) {
                        continue;
                    }
                    if (ign.symbol == reg.symbol) {
                        ignored = true;
                        break;
                    }
                }
                if (ignored) {
                    ++iter;
                    continue;
                }
                if (info.start == 0 && !info.elf) {
                    info.elf = std::make_unique<Elf>(info.start);
                }
                if (info.elf && info.elf->Valid()) {
                    for (auto addr: info.elf->FindPltAddr(reg.symbol)) {
                        res = DoHook(addr, reinterpret_cast<uintptr_t>(reg.callback),
                                     reinterpret_cast<uintptr_t *>(reg.backup)) && res;
                    }
                }
                if (reg.inode != 0) {
                    iter = register_info.erase(iter);
                } else {
                    ++iter;
                }
            }
        }
        return res;
    }

    bool InvalidateBackup() {
        bool res = true;
        for (auto &[_, info] : *this) {
            if (info.backup) {
                for (auto &[addr, backup] : info.hooks) {
                    // store new address to backup since we don't need backup
                    backup = *reinterpret_cast<uintptr_t *>(addr);
                }
                auto len = info.end - info.start;
                if (auto *new_addr =
                        mremap(reinterpret_cast<void *>(info.backup), len, len,
                               MREMAP_FIXED | MREMAP_MAYMOVE, reinterpret_cast<void *>(info.start));
                    new_addr == MAP_FAILED || reinterpret_cast<uintptr_t>(new_addr) != info.start) {
                    res = false;
                    info.hooks.clear();
                    continue;
                }
                for (auto &[addr, backup] : info.hooks) {
                    *reinterpret_cast<uintptr_t *>(addr) = backup;
                }
                info.hooks.clear();
                info.backup = 0;
            }
        }
        return res;
    }
};

std::mutex register_mutex;
std::list<RegisterInfo> register_info;
std::mutex ignore_mutex;
std::list<IgnoreInfo> ignore_info;
MapInfos map_info;
}  // namespace

namespace lsplt {
[[maybe_unused]] bool RegisterHook(std::string_view regex_str, std::string_view symbol,
                                   void *callback, void **backup) {
    if (regex_str.empty() || symbol.empty() || !callback) return false;

    regex_t regex;

    if (regcomp(&regex, regex_str.data(), REG_NOSUB) != 0) return false;

    std::unique_lock lock(register_mutex);
    register_info.emplace_back(RegisterInfo{{regex, 0, std::string{symbol}}, callback, backup});

    return true;
}

[[maybe_unused]] bool RegisterHook(ino_t ino, std::string_view symbol, void *callback,
                                   void **backup) {
    if (symbol.empty() || !callback) return false;

    std::unique_lock lock(register_mutex);
    register_info.emplace_back(RegisterInfo{{{}, ino, std::string{symbol}}, callback, backup});

    return true;
}

[[maybe_unused]] void IgnoreHook(std::string_view regex, std::string_view symbol) {
    if (regex.empty() || symbol.empty()) return;

    regex_t reg;

    if (regcomp(&reg, regex.data(), REG_NOSUB) != 0) return;

    std::unique_lock lock(ignore_mutex);
    ignore_info.emplace_back(IgnoreInfo{{reg, 0, std::string{symbol}}});
}

[[maybe_unused]] void IgnoreHook(ino_t ino, std::string_view symbol) {
    if (symbol.empty()) return;

    std::unique_lock lock(ignore_mutex);
    ignore_info.emplace_back(IgnoreInfo{{regex_t{}, ino, std::string{symbol}}});
}

[[maybe_unused]] bool CommitHook() {
    std::unique_lock lock(register_mutex);
    if (register_info.empty()) return true;
    std::unique_lock lock2(ignore_mutex);

    auto new_map_info = MapInfos::ScanMapInfo();
    if (new_map_info.empty()) return false;

    new_map_info.Filter(register_info, ignore_info);

    new_map_info.Merge(map_info);
    // update to new map info
    map_info = std::move(new_map_info);

    return map_info.DoHook(register_info, ignore_info);
}

[[gnu::destructor]] [[maybe_unused]] bool InvalidateBackup() {
    std::unique_lock lock(register_mutex);
    std::unique_lock lock2(ignore_mutex);
    return map_info.InvalidateBackup();
}
}  // namespace lsplt