tinyrpc/src/coroutine/coroutine_hook.cc

#include "coroutine_hook.hpp"
#include "coroutine.hpp"
#include "logger.hpp"
#include "fd_event.hpp"
#include "reactor.hpp"
#include <dlfcn.h>


#define HOOK_SYSTEM_FUN(name) name##_fun_ptr_t g_sys_##name##_fun = (name##_fun_ptr_t)dlsym(RTLD_NEXT, #name)


namespace tinyrpc {
    
    HOOK_SYSTEM_FUN(read);
    HOOK_SYSTEM_FUN(write);

    static bool isEnableHook = false;
    void enableHook() {
        isEnableHook = true;
    }
    void disableHook() {
        isEnableHook = false;
    }

    ssize_t read_hook(int fd, void *buf, size_t count) {
        logger() << "read_hook is calling";
        FdEvent fe(fd);
        fe.addListenEvent(IOEvent::READ);
        Coroutine* curCoro = Coroutine::getCurrCoroutine();
        fe.setReadCallback([curCoro] () -> void{
            curCoro->resume();
        });
        //  fd 设置为 nonblock
        fe.setNonblock();
        //  尝试一下系统read， 返回值大于0直接返回
        int ret = g_sys_read_fun(fd, buf, count);
        if(ret > 0) return ret;
        //  fd 添加到 epoll 中
        Reactor::getReactor()->addFdEvent(&fe);
        Coroutine::yeild(); //  yeild
        Reactor::getReactor()->delFdEvent(&fe);
        //  调用系统 read 返回
        return g_sys_read_fun(fd, buf, count);
    }

    ssize_t write_hook(int fd, const void *buf, size_t count) {
        logger() << "write_hook is calling";
        FdEvent fe(fd);
        fe.addListenEvent(IOEvent::WRITE);
        Coroutine* curCoro = Coroutine::getCurrCoroutine();
        fe.setWriteCallback([curCoro] () -> void{
            curCoro->resume();
        });
        //  fd 设置为 nonblock
        fe.setNonblock();
        //  尝试一下系统 write 返回值大于0直接返回
        int ret = g_sys_write_fun(fd, buf, count);
        if(ret > 0) return ret;
        //  fd 添加到 epoll 中
        Reactor::getReactor()->addFdEvent(&fe);
        Coroutine::yeild(); //  yeild
        Reactor::getReactor()->delFdEvent(&fe);
        //  调用系统 write 返回
        return g_sys_write_fun(fd, buf, count);
    }

}

ssize_t read(int fd, void *buf, size_t count) {
    if (tinyrpc::isEnableHook == false) {
        return tinyrpc::g_sys_read_fun(fd, buf, count);  //  没有启用 hook， 直接转发到系统调用
    }
    return tinyrpc::read_hook(fd, buf, count);
}

ssize_t write(int fd, const void *buf, size_t count) {
    if (tinyrpc::isEnableHook == false) {
        return tinyrpc::g_sys_write_fun(fd, buf, count);  //  没有启用 hook， 直接转发到系统调用
    }
    return tinyrpc::write_hook(fd, buf, count);
}