#include "coroutine_hook.hpp" #include "coroutine.hpp" #include "logger.hpp" #include "fd_event.hpp" #include "reactor.hpp" #include #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); }