/*
 * Copyright (C) 2017 ~ 2018 Deepin Technology Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "vtablehook.h"

#include <QFileInfo>
#include <algorithm>

#ifdef Q_OS_LINUX
#include <sys/mman.h>
#include <unistd.h>
#include <dlfcn.h>
#include <mutex>

QT_BEGIN_NAMESPACE
QFunctionPointer qt_linux_find_symbol_sys(const char *symbol);
QT_END_NAMESPACE

#endif

DPP_BEGIN_NAMESPACE

QMap<quintptr**, quintptr*> VtableHook::objToOriginalVfptr;
QMap<const void*, quintptr*> VtableHook::objToGhostVfptr;
QMap<const void*, quintptr> VtableHook::objDestructFun;
static std::once_flag exitFlag;

bool VtableHook::copyVtable(quintptr **obj)
{
    int vtable_size = getVtableSize(obj);

    if (vtable_size == 0)
        return false;

    // 多开辟一个元素, 新的虚表结构如下:
    // 假设obj对象原虚表长度为2, 表结构为:
    // ┏━━┳━━┳━━┓其中v1 v2为虚函数地址, 0为数组结尾
    // ┃v1┃v2┃\0┃
    // ┗━━┻━━┻━━┛
    // 则新的表结构为:
    // ┏━━┳━━┳━━┳━━┓其中前三个元素为原虚表的复制, sv为原虚表入口地址
    // ┃v1┃v2┃\0┃sv┃
    // ┗━━┻━━┻━━┻━━┛
    vtable_size += 2;

    quintptr *new_vtable = new quintptr[vtable_size];

    memcpy(new_vtable, *obj, (vtable_size - 1) * sizeof(quintptr));

    //! save original vfptr
    objToOriginalVfptr[obj] = *obj;
    // 存储对象原虚表入口地址
    new_vtable[vtable_size - 1] = quintptr(*obj);

    *obj = new_vtable;
    //! save ghost vfptr
    objToGhostVfptr[obj] = new_vtable;

    return true;
}

bool VtableHook::clearGhostVtable(const void *obj)
{
    objToOriginalVfptr.remove((quintptr**)obj);
    objDestructFun.remove(obj);

    quintptr *vtable = objToGhostVfptr.take(obj);

    if (vtable) {
        delete[] vtable;

        return true;
    }

    return false;
}

void VtableHook::clearAllGhostVtable()
{
    const QList<const void *> _objects = objToGhostVfptr.keys();

    for (const void *_obj : _objects)
        clearGhostVtable(_obj);
}

/*!
 * \brief 通过遍历尝试找到析构函数在虚表中的位置
 * \param obj
 * \param destoryObjFun
 * \return
 */
int VtableHook::getDestructFunIndex(quintptr **obj, std::function<void(void)> destoryObjFun)
{
    class _DestoryProbe
    {
    public:
        static quintptr probe(quintptr obj) {
            static quintptr _obj = 0;

            if (obj == 0) {
                obj = _obj;
                _obj = 0;
            } else {
                _obj = obj;
            }

            return obj;
        }

        static void nothing() {

        }
    };

    quintptr *vtable = *obj;
    int vtable_size = getVtableSize(obj);

    if (vtable_size == 0)
        return -1;

    quintptr *new_vtable = new quintptr[vtable_size];
    std::fill(new_vtable, new_vtable + vtable_size, quintptr(&_DestoryProbe::nothing));

    // 给对象设置新的虚表
    *obj = new_vtable;

    int index = -1;

    for (int i = 0; i < vtable_size; ++i) {
        new_vtable[i] = quintptr(&_DestoryProbe::probe);

        // 尝试销毁此对象, 且观察_DestoryProbe::probe是否被调用
        // 如果被调用, 则证明覆盖此虚函数能达到监控对象被销毁的目的
        destoryObjFun();

        if (_DestoryProbe::probe(0) == quintptr(obj)) {
            index = i;
            break;
        }
    }

    // 恢复旧的虚表
    *obj = vtable;
    // 销毁临时虚表
    delete[] new_vtable;

    return index;
}

void VtableHook::autoCleanVtable(const void *obj)
{
    quintptr fun = objDestructFun.value(obj);

    if (!fun)
        return;

    typedef void(*Destruct)(const void*);
    Destruct destruct = *reinterpret_cast<Destruct*>(&fun);
    // call origin destruct function
    destruct(obj);

    if (hasVtable(obj)) {// 需要判断一下，有可能在执行析构函数时虚表已经被删除
        // clean
        clearGhostVtable(obj);
    }
}

bool VtableHook::ensureVtable(const void *obj, std::function<void ()> destoryObjFun)
{
    quintptr **_obj = (quintptr**)(obj);

    if (objToOriginalVfptr.contains(_obj)) {
        // 不知道什么原因, 此时obj对象的虚表已经被还原
        if (objToGhostVfptr.value((void*)obj) != *_obj) {
            clearGhostVtable((void*)obj);
        } else {
            return true;
        }
    }

    if (!copyVtable(_obj))
        return false;

    // 查找对象的析构函数
    int index = getDestructFunIndex(_obj, destoryObjFun);

    // 虚析构函数查找失败
    if (index < 0) {
        qWarning("Failed do override destruct function");
        qDebug() << "object:" << obj;
        abort();
    }

    quintptr *new_vtable = *_obj;
    // 保存对象真实的析构函数
    objDestructFun[(void*)obj] = new_vtable[index];

    // 覆盖析构函数, 用于在对象析构时自动清理虚表
    new_vtable[index] = reinterpret_cast<quintptr>(&autoCleanVtable);

    // TODO: 由于未知原因,有的虚表会自动还原，导致虚析构不能正常HOOK，无法释放new出来的新虚表数组
    // 这里在程序退出时进行统一释放。后面知道详细原因再进行修改。
    std::call_once(exitFlag, std::bind(atexit, clearAllGhostVtable));
    return true;
}

/*!
 * \brief VtableHook::hasVtable 对象的虚表已经被覆盖时返回true，否则返回false
 * \param obj
 * \return
 */
bool VtableHook::hasVtable(const void *obj)
{
    quintptr **_obj = (quintptr**)(obj);

    return objToGhostVfptr.contains(_obj);
}

void VtableHook::resetVtable(const void *obj)
{
    quintptr **_obj = (quintptr**)obj;
    int vtable_size = getVtableSize(_obj);
    // 获取obj对象原本虚表的入口
    quintptr *vfptr_t2 = (quintptr*)(*_obj)[vtable_size + 1];

    if (!vfptr_t2)
        return;

    if (!clearGhostVtable(obj))
        return;

    // 还原虚表
    *_obj = vfptr_t2;
}

/*!
 * \brief 将偏移量为functionOffset的虚函数还原到原本的实现
 * \param obj
 * \param functionIndex
 * \return 如果成功, 返回还原之前obj对象虚表中存储的函数指针, 否则返回0
 */
quintptr VtableHook::resetVfptrFun(const void *obj, quintptr functionOffset)
{
    quintptr *vfptr_t1 = *(quintptr**)obj;
    quintptr current_fun = *(vfptr_t1 + functionOffset / sizeof(quintptr));
    quintptr origin_fun = originalFun(obj, functionOffset);

    if (!origin_fun) {
        return 0;
    }

    // reset to original fun
    *(vfptr_t1 + functionOffset / sizeof(quintptr)) = origin_fun;

    return current_fun;
}

/*!
 * \brief 获取obj对象偏移量为functionOffset的虚函数原本的函数指针
 * \param obj
 * \param functionOffset
 * \return 如果obj对象虚表没有被覆盖, 或者函数偏移量正确, 将返回0
 */
quintptr VtableHook::originalFun(const void *obj, quintptr functionOffset)
{
    quintptr **_obj = (quintptr**)obj;
    int vtable_size = getVtableSize(_obj);
    // 获取obj对象原本虚表的入口
    quintptr *vfptr_t2 = (quintptr*)(*_obj)[vtable_size + 1];

    if (!vfptr_t2) {
        qWarning() << "Not override the object virtual table" << obj;

        return 0;
    }

    if (functionOffset > UINT_LEAST16_MAX) {
        qWarning() << "Is not a virtual function, function address: 0x" << hex << functionOffset;

        return 0;
    }

    return *(vfptr_t2 + functionOffset / sizeof(quintptr));
}

#if defined(Q_OS_LINUX)
static int readProtFromPsm(quintptr adr, size_t length)
{
    int prot = PROT_NONE;
    QString fname = "/proc/self/maps";
    QFile f(fname);
    if (!f.open(QIODevice::ReadOnly)) {
        qFatal("%s", f.errorString().toStdString().data());
        //return prot; // never be executed
    }

    QByteArray data = f.readAll();
    bool ok = false;
    quintptr startAddr = 0, endAddr = 0;
    QTextStream ts(data);
    while (Q_UNLIKELY(!ts.atEnd())) {
        const QString line = ts.readLine();
        const QStringList &maps = line.split(' ');
        if (Q_UNLIKELY(maps.size() < 3)) {
            continue;
        }

        //"00400000-00431000" "r--p"
        const QStringList addrs = maps.value(0).split('-');
        startAddr = addrs.value(0).toULongLong(&ok, 16);
        Q_ASSERT(ok);
        endAddr = addrs.value(1).toULongLong(&ok, 16);
        Q_ASSERT(ok);
        if (Q_LIKELY(adr >= endAddr)) {
            continue;
        }
        if (adr >= startAddr && adr + length <= endAddr) {
            QString ps = maps.value(1);
            //qDebug() << maps.value(0) << maps.value(1);
            for (QChar c : ps) {
                switch (c.toLatin1()) {
                case 'r':
                    prot |= PROT_READ;
                    break;
                case 'w':
                    prot |= PROT_WRITE;
                    break;
                case 'x':
                    prot |= PROT_EXEC;
                    break;
                default:
                    break; // '-' 'p' don't care
                }
            }
            break;
        } else if (adr < startAddr) {
            qFatal("%p not found in proc maps", reinterpret_cast<void *>(adr));
            //break; // 超出了地址不需要再去检查了
        }
    }

    return prot;
}
#endif

bool VtableHook::forceWriteMemory(void *adr, const void *data, size_t length)
{
#ifdef Q_OS_LINUX
    int page_size = sysconf(_SC_PAGESIZE);
    quintptr x = reinterpret_cast<quintptr>(adr);
    // 不减去一个pagesize防止跨越两个数据区域(对应/proc/self/maps两行数据)
    void *new_adr = reinterpret_cast<void *>((x /*- page_size - 1*/) & ~(page_size - 1));
    size_t override_data_length = length + x - reinterpret_cast<quintptr>(new_adr);

    int oldProt = readProtFromPsm(quintptr(new_adr), override_data_length);
    bool writeable = oldProt & PROT_WRITE;
    // 增加判断是否已经可写，不能写才调用。
    // 失败时直接放弃
    if (!writeable && mprotect(new_adr, override_data_length, PROT_READ | PROT_WRITE)) {
        qWarning() << "mprotect(change) failed" << strerror(errno);
        return false;
    }
#endif
    // 复制数据
    memcpy(adr, data, length);
#ifdef Q_OS_LINUX
    // 恢复内存标志位
    if (!writeable && mprotect(new_adr, override_data_length, oldProt)) {
        qWarning() << "mprotect(restore) failed" << strerror(errno);
        return false;
    }
#endif

    return true;
}

QFunctionPointer VtableHook::resolve(const char *symbol)
{
#ifdef Q_OS_LINUX
/**
  * ！！不要使用qt_linux_find_symbol_sys函数去获取符号
  *
  * 在龙芯平台上，qt_linux_find_symbol_sys 无法获取部分已加载动态库的符号，
  * 可能的原因是这个函数对 dlsym 的调用是在 libQt5Core 动态库中，这个库加载的比较早，
  * 有可能是因此导致无法获取比这个库加载更晚的库中的符号(仅为猜测)
  */
    return QFunctionPointer(dlsym(RTLD_DEFAULT, symbol));
#else
    // TODO
    return nullptr;
#endif
}

DPP_END_NAMESPACE
