// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
// This file is copied from
// https://github.com/ClickHouse/ClickHouse/blob/master/src/base/phdr_cache.cpp
// and modified by Doris

#include "common/phdr_cache.h"

#if defined(__clang__)
#pragma clang diagnostic ignored "-Wreserved-identifier"
#endif

/// This code was based on the code by Fedor Korotkiy https://www.linkedin.com/in/fedor-korotkiy-659a1838/

#if defined(__linux__) && !defined(THREAD_SANITIZER) && !defined(USE_MUSL)
#define USE_PHDR_CACHE 1
#endif

/// Thread Sanitizer uses dl_iterate_phdr function on initialization and fails if we provide our own.
#ifdef USE_PHDR_CACHE

#if defined(__clang__)
#pragma clang diagnostic ignored "-Wreserved-id-macro"
#pragma clang diagnostic ignored "-Wunused-macros"
#endif

#define __msan_unpoison(X, Y) // NOLINT
#if defined(__clang__) && defined(__has_feature)
#if __has_feature(memory_sanitizer)
#undef __msan_unpoison
#include <sanitizer/msan_interface.h>
#endif
#endif

#include <dlfcn.h>
#include <link.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <cstddef>
#include <cstring>
#include <limits>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#if defined(USE_UNWIND) && USE_UNWIND
#ifndef UNW_LOCAL_ONLY
#define UNW_LOCAL_ONLY
#endif
#include <libunwind.h>
#endif

namespace {

// This is adapted from
// https://github.com/scylladb/seastar/blob/master/core/exception_hacks.hh
// https://github.com/scylladb/seastar/blob/master/core/exception_hacks.cc

constexpr size_t MAX_PHDR_CACHE_LOADED_OBJECTS = 4096;
constexpr size_t MAX_PHDR_CACHE_PROGRAM_HEADERS = 128;
constexpr size_t MAX_PHDR_CACHE_OBJECT_NAME = 4096;

using DLIterateFunction = int (*)(int (*callback)(dl_phdr_info* info, size_t size, void* data),
                                  void* data);

DLIterateFunction getOriginalDLIteratePHDR() {
    void* func = dlsym(RTLD_NEXT, "dl_iterate_phdr");
    if (!func) {
        throw std::runtime_error("Cannot find dl_iterate_phdr function with dlsym");
    }
    return reinterpret_cast<DLIterateFunction>(func);
}

struct RawPHDRCacheEntry {
    dl_phdr_info info {};
    std::array<ElfW(Phdr), MAX_PHDR_CACHE_PROGRAM_HEADERS> phdrs {};
    size_t phdr_count = 0;
    std::array<char, MAX_PHDR_CACHE_OBJECT_NAME> name {};
    uintptr_t address_begin = std::numeric_limits<uintptr_t>::max();
    uintptr_t address_end = 0;
};

struct RawPHDRCacheSnapshot {
    std::array<RawPHDRCacheEntry, MAX_PHDR_CACHE_LOADED_OBJECTS> entries {};
    size_t size = 0;
    bool overflow = false;
    bool phdr_truncated = false;
    bool name_truncated = false;
};

struct PHDRCacheEntry {
    dl_phdr_info info {};
    std::vector<ElfW(Phdr)> phdrs;
    std::string name;
    uintptr_t address_begin = std::numeric_limits<uintptr_t>::max();
    uintptr_t address_end = 0;

    bool contains(uintptr_t ip) const {
        return ip == 0 || (address_begin <= ip && ip < address_end);
    }
};

using PHDRCache = std::vector<PHDRCacheEntry>;
std::atomic<PHDRCache*> phdr_cache {};
// This flag is flipped inside the stack-trace signal handler. Force a static TLS access model so
// reading it from our dl_iterate_phdr interposer does not call into the dynamic loader's TLS path.
__thread bool use_phdr_cache __attribute__((tls_model("initial-exec"))) = false;

uintptr_t saturated_segment_end(uintptr_t begin, uintptr_t size) {
    const uintptr_t max_address = std::numeric_limits<uintptr_t>::max();
    return begin > max_address - size ? max_address : begin + size;
}

void copy_object_name(const char* source, RawPHDRCacheEntry* entry,
                      RawPHDRCacheSnapshot* snapshot) {
    if (source == nullptr) {
        entry->name[0] = '\0';
        return;
    }

    size_t length = 0;
    while (length + 1 < entry->name.size() && source[length] != '\0') {
        entry->name[length] = source[length];
        ++length;
    }
    entry->name[length] = '\0';
    if (source[length] != '\0') {
        snapshot->name_truncated = true;
    }
}

int collectPHDRCacheEntry(dl_phdr_info* info, size_t /*size*/, void* data) {
    auto* snapshot = reinterpret_cast<RawPHDRCacheSnapshot*>(data);
    if (snapshot->size >= snapshot->entries.size()) {
        snapshot->overflow = true;
        return 0;
    }

    auto& entry = snapshot->entries[snapshot->size++];
    entry.info = *info;
    copy_object_name(info->dlpi_name, &entry, snapshot);

    const size_t phdr_count = std::min<size_t>(info->dlpi_phnum, entry.phdrs.size());
    if (phdr_count < info->dlpi_phnum) {
        snapshot->phdr_truncated = true;
    }
    entry.phdr_count = phdr_count;
    entry.info.dlpi_phnum = static_cast<ElfW(Half)>(phdr_count);
    entry.info.dlpi_name = nullptr;
    entry.info.dlpi_phdr = nullptr;

    if (info->dlpi_phdr == nullptr) {
        return 0;
    }
    std::memcpy(entry.phdrs.data(), info->dlpi_phdr, phdr_count * sizeof(ElfW(Phdr)));

    for (size_t i = 0; i < phdr_count; ++i) {
        const auto& phdr = entry.phdrs[i];
        if (phdr.p_type != PT_LOAD || phdr.p_memsz == 0) {
            continue;
        }
        const auto begin = static_cast<uintptr_t>(info->dlpi_addr + phdr.p_vaddr);
        const auto end = saturated_segment_end(begin, static_cast<uintptr_t>(phdr.p_memsz));
        entry.address_begin = std::min(entry.address_begin, begin);
        entry.address_end = std::max(entry.address_end, end);
    }
    return 0;
}

PHDRCache* buildPHDRCache(const RawPHDRCacheSnapshot& snapshot) {
    auto* cache = new PHDRCache;
    cache->reserve(snapshot.size);
    for (size_t i = 0; i < snapshot.size; ++i) {
        const auto& raw_entry = snapshot.entries[i];
        PHDRCacheEntry entry;
        entry.info = raw_entry.info;
        entry.phdrs.assign(raw_entry.phdrs.begin(), raw_entry.phdrs.begin() + raw_entry.phdr_count);
        entry.name = raw_entry.name.data();
        entry.address_begin = raw_entry.address_begin;
        entry.address_end = raw_entry.address_end;
        cache->emplace_back(std::move(entry));
    }
    for (auto& entry : *cache) {
        entry.info.dlpi_phdr = entry.phdrs.data();
        entry.info.dlpi_name = entry.name.c_str();
    }
    return cache;
}

int iteratePHDRCache(int (*callback)(dl_phdr_info* info, size_t size, void* data), void* data,
                     uintptr_t ip) {
    auto* current_phdr_cache = phdr_cache.load(std::memory_order_acquire);
    if (current_phdr_cache == nullptr) {
        return 0;
    }

    int result = 0;
    for (auto& entry : *current_phdr_cache) {
        if (!entry.contains(ip)) {
            continue;
        }
        result = callback(&entry.info, sizeof(dl_phdr_info), data);
        if (result != 0) {
            break;
        }
    }
    return result;
}

} // namespace

extern "C"
#ifndef __clang__
        [[gnu::visibility("default")]] [[gnu::externally_visible]]
#endif
        int
        dl_iterate_phdr(int (*callback)(dl_phdr_info* info, size_t size, void* data), void* data) {
    if (!use_phdr_cache) {
        return getOriginalDLIteratePHDR()(callback, data);
    }

    return iteratePHDRCache(callback, data, 0);
}

extern "C"
#ifndef __clang__
        [[gnu::visibility("default")]] [[gnu::externally_visible]]
#endif
        int
        doris_unwind_iterate_phdr(int (*callback)(dl_phdr_info* info, size_t size, void* data),
                                  void* data, uintptr_t ip) {
    return iteratePHDRCache(callback, data, ip);
}

#include "util/debug/leak_annotations.h"

void updatePHDRCache() {
    // Fill out ELF header cache for access without locking.
    // Old snapshots are intentionally kept alive because another thread may already be unwinding
    // through the previous cache when a Doris-controlled dlopen/dlclose refreshes this one.

    auto raw_snapshot = std::make_unique<RawPHDRCacheSnapshot>();
    getOriginalDLIteratePHDR()(
            [](dl_phdr_info* info, size_t size, void* data) {
                // `info` is created by dl_iterate_phdr, which is a non-instrumented
                // libc function, so we have to unpoison it manually.
                __msan_unpoison(info, sizeof(*info));

                return collectPHDRCacheEntry(info, size, data);
            },
            raw_snapshot.get());

    PHDRCache* new_phdr_cache = buildPHDRCache(*raw_snapshot);
    phdr_cache.store(new_phdr_cache, std::memory_order_release);

    /// Memory is intentionally leaked.
    ANNOTATE_LEAKING_OBJECT_PTR(new_phdr_cache);
}

bool hasPHDRCache() {
    return phdr_cache.load(std::memory_order_acquire) != nullptr;
}

void configureLibunwindPHDRCache() {
#if defined(USE_UNWIND) && USE_UNWIND
    static std::once_flag once;
    std::call_once(once, [] {
        unw_context_t context;
        unw_cursor_t cursor;
        (void)unw_getcontext(&context);
        (void)unw_init_local(&cursor, &context);
        // Doris-patched libunwind gets FDEs from the PHDR snapshot. Disable the global DWARF
        // register-state cache so a signal handler cannot self-deadlock on libunwind's cache mutex
        // after interrupting a thread that was already unwinding.
        (void)unw_set_caching_policy(unw_local_addr_space, UNW_CACHE_NONE);
    });
#endif
}

ScopedPHDRCacheRead::ScopedPHDRCacheRead() : _previous(use_phdr_cache) {
    use_phdr_cache = true;
}

ScopedPHDRCacheRead::~ScopedPHDRCacheRead() {
    use_phdr_cache = _previous;
}

#else

void updatePHDRCache() {}

void configureLibunwindPHDRCache() {}

#if defined(USE_MUSL)
/// With statically linked with musl, dl_iterate_phdr is immutable.
bool hasPHDRCache() {
    return true;
}
#else
bool hasPHDRCache() {
    return false;
}
#endif

ScopedPHDRCacheRead::ScopedPHDRCacheRead() = default;

ScopedPHDRCacheRead::~ScopedPHDRCacheRead() = default;

#endif