113 lines
3.8 KiB
C++
113 lines
3.8 KiB
C++
#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)
|
||
HOOK_SYSTEM_FUN(read);
|
||
HOOK_SYSTEM_FUN(write);
|
||
HOOK_SYSTEM_FUN(accept);
|
||
|
||
namespace tinyrpc {
|
||
|
||
|
||
|
||
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);
|
||
}
|
||
|
||
int accept_hook(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
|
||
logger() << "accept_hook is calling";
|
||
FdEvent fe(sockfd);
|
||
fe.addListenEvent(IOEvent::READ);
|
||
Coroutine* curCoro = Coroutine::getCurrCoroutine();
|
||
fe.setReadCallback([curCoro] () -> void{
|
||
curCoro->resume();
|
||
});
|
||
// logger() << "accept_hook fd = " << fe.getFd();
|
||
// fd 设置为 nonblock
|
||
fe.setNonblock();
|
||
// 尝试一下系统 accept 返回值大于 0 直接返回
|
||
int ret = g_sys_accept_fun(sockfd, addr, addrlen);
|
||
if(ret >= 0) return ret;
|
||
// fd 添加到 epoll 中
|
||
Reactor::getReactor()->addFdEvent(&fe);
|
||
logger() << "accept_hook cor yeild";
|
||
Coroutine::yeild(); // yeild
|
||
logger() << "accept_hook cor resume then call g_sys_accept_fun";
|
||
Reactor::getReactor()->delFdEvent(&fe);
|
||
// 调用系统 write 返回
|
||
return g_sys_accept_fun(sockfd, addr, addrlen);
|
||
}
|
||
|
||
}
|
||
|
||
ssize_t read(int fd, void *buf, size_t count) {
|
||
if (tinyrpc::isEnableHook == false) {
|
||
return 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 g_sys_write_fun(fd, buf, count); // 没有启用 hook， 直接转发到系统调用
|
||
}
|
||
return tinyrpc::write_hook(fd, buf, count);
|
||
}
|
||
|
||
int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen) {
|
||
if (tinyrpc::isEnableHook == false) {
|
||
return g_sys_accept_fun(sockfd, addr, addrlen); // 没有启用 hook， 直接转发到系统调用
|
||
}
|
||
return tinyrpc::accept_hook(sockfd, addr, addrlen);
|
||
} |