首页
社区
课程
招聘
问答 CTF 排行榜 知识库 工具下载 峰会 看雪商城 证书查询
  1. 社区
  2. 编程技术
    3. 发新帖
[原创]NtSocket的稳定实现，Client与Server的简单封装，以及SocketAsyncSelect的一种APC实现
发表于: 2022-9-8 23:46 17126

[原创]NtSocket的稳定实现，Client与Server的简单封装，以及SocketAsyncSelect的一种APC实现

ANormalUser 活跃值
2022-9-8 23:46
17126

#源码：

SOCKET WSPSocket(
int AddressFamily,
int SocketType,
int Protocol) {
/// <summary>
/// 类似于Socket函数，可以创建一个Socket文件句柄
/// </summary>
/// <param name="AddressFamily">Address family(Support IPv6)</param>
/// <param name="SocketType">Socket Type</param>
/// <param name="Protocol">Protocol type</param>
/// <returns>如果失败返回INVALID_SOCKET，成功返回Socket文件句柄</returns>
if (AddressFamily == AF_UNSPEC && SocketType == 0 && Protocol == 0) {
   return INVALID_SOCKET;
}
//进行基础数据设置
if (AddressFamily == AF_UNSPEC) {
   AddressFamily = AF_INET;
}
if (SocketType == 0)
{
   switch (Protocol)
   {
   case IPPROTO_TCP:
       SocketType = SOCK_STREAM;
       break;
   case IPPROTO_UDP:
       SocketType = SOCK_DGRAM;
       break;
   case IPPROTO_RAW:
       SocketType = SOCK_RAW;
       break;
   default:
       SocketType = SOCK_STREAM;
       break;
   }
}
if (Protocol == 0)
{
   switch (SocketType)
   {
   case SOCK_STREAM:
       Protocol = IPPROTO_TCP;
       break;
   case SOCK_DGRAM:
       Protocol = IPPROTO_UDP;
       break;
   case SOCK_RAW:
       Protocol = IPPROTO_RAW;
       break;
   default:
       Protocol = IPPROTO_TCP;
       break;
   }
}
byte EaBuffer[] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x1E, 0x00,
   0x41, 0x66, 0x64, 0x4F, 0x70, 0x65, 0x6E, 0x50,
   0x61, 0x63, 0x6B, 0x65, 0x74, 0x58, 0x58, 0x00,
   0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
   0x02, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
   0x08, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
   0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
memmove((PVOID)((__int64)EaBuffer + 32), &AddressFamily, 0x4);
memmove((PVOID)((__int64)EaBuffer + 36), &SocketType, 0x4);
memmove((PVOID)((__int64)EaBuffer + 40), &Protocol, 0x4);
if (Protocol == IPPROTO_UDP)
{
   memmove((PVOID)((__int64)EaBuffer + 24), &Protocol, 0x4);
}
//初始化UNICODE_STRING：
UNICODE_STRING AfdName;
AfdName.Buffer = L"\\Device\\Afd\\Endpoint";
AfdName.Length = 2 * wcslen(AfdName.Buffer);
AfdName.MaximumLength = AfdName.Length + 2;
OBJECT_ATTRIBUTES  Object;
IO_STATUS_BLOCK IOSB;
//初始化OBJECT_ATTRIBUTES
InitializeObjectAttributes(&Object,
   &AfdName,
   OBJ_CASE_INSENSITIVE | OBJ_INHERIT,
   0,
   0);
HANDLE MySock;
NTSTATUS Status;
//创建AfdSocket：
Status = ((NtCreateFile)MyNtCreateFile)(&MySock,
   GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
   &Object,
   &IOSB,
   NULL,
   0,
   FILE_SHARE_READ | FILE_SHARE_WRITE,
   FILE_OPEN_IF,
   0,
   EaBuffer,
   sizeof(EaBuffer));
if (Status != STATUS_SUCCESS) {
   return INVALID_SOCKET;
}
else {
   return (SOCKET)MySock;
}
}
SOCKET WSPSocket(
int AddressFamily,
int SocketType,
int Protocol) {
/// <summary>
/// 类似于Socket函数，可以创建一个Socket文件句柄
/// </summary>
/// <param name="AddressFamily">Address family(Support IPv6)</param>
/// <param name="SocketType">Socket Type</param>
/// <param name="Protocol">Protocol type</param>
/// <returns>如果失败返回INVALID_SOCKET，成功返回Socket文件句柄</returns>
if (AddressFamily == AF_UNSPEC && SocketType == 0 && Protocol == 0) {
   return INVALID_SOCKET;
}
//进行基础数据设置
if (AddressFamily == AF_UNSPEC) {
   AddressFamily = AF_INET;
}
if (SocketType == 0)
{
   switch (Protocol)
   {
   case IPPROTO_TCP:
       SocketType = SOCK_STREAM;
       break;
   case IPPROTO_UDP:
       SocketType = SOCK_DGRAM;
       break;
   case IPPROTO_RAW:
       SocketType = SOCK_RAW;
       break;
   default:
       SocketType = SOCK_STREAM;
       break;
   }
}
if (Protocol == 0)
{
   switch (SocketType)
   {
   case SOCK_STREAM:
       Protocol = IPPROTO_TCP;
       break;
   case SOCK_DGRAM:
       Protocol = IPPROTO_UDP;
       break;
   case SOCK_RAW:
       Protocol = IPPROTO_RAW;
       break;
   default:
       Protocol = IPPROTO_TCP;
       break;
   }
}
byte EaBuffer[] = {0x00, 0x00, 0x00, 0x00, 0x00, 0x0F, 0x1E, 0x00,
   0x41, 0x66, 0x64, 0x4F, 0x70, 0x65, 0x6E, 0x50,
   0x61, 0x63, 0x6B, 0x65, 0x74, 0x58, 0x58, 0x00,
   0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
   0x02, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
   0x08, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
   0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
memmove((PVOID)((__int64)EaBuffer + 32), &AddressFamily, 0x4);
memmove((PVOID)((__int64)EaBuffer + 36), &SocketType, 0x4);
memmove((PVOID)((__int64)EaBuffer + 40), &Protocol, 0x4);
if (Protocol == IPPROTO_UDP)
{
   memmove((PVOID)((__int64)EaBuffer + 24), &Protocol, 0x4);
}
//初始化UNICODE_STRING：
UNICODE_STRING AfdName;
AfdName.Buffer = L"\\Device\\Afd\\Endpoint";
AfdName.Length = 2 * wcslen(AfdName.Buffer);
AfdName.MaximumLength = AfdName.Length + 2;
OBJECT_ATTRIBUTES  Object;
IO_STATUS_BLOCK IOSB;
//初始化OBJECT_ATTRIBUTES
InitializeObjectAttributes(&Object,
   &AfdName,
   OBJ_CASE_INSENSITIVE | OBJ_INHERIT,
   0,
   0);
HANDLE MySock;
NTSTATUS Status;
//创建AfdSocket：
Status = ((NtCreateFile)MyNtCreateFile)(&MySock,
   GENERIC_READ | GENERIC_WRITE | SYNCHRONIZE,
   &Object,
   &IOSB,
   NULL,
   0,
   FILE_SHARE_READ | FILE_SHARE_WRITE,
   FILE_OPEN_IF,
   0,
   EaBuffer,
   sizeof(EaBuffer));
if (Status != STATUS_SUCCESS) {
   return INVALID_SOCKET;
}
else {
   return (SOCKET)MySock;
}
}
NTSTATUS WSPProcessAsyncSelect(
SOCKET Handle,
PVOID ApcRoutine,
ULONG lNetworkEvents,
PVOID UserContext
) {
/// <summary>
/// WSPProcessAsyncSelect是本程序最大的亮点，利用异步化的IOCTL_AFD_SELECT
/// </summary>
/// <param name="Handle"></param>
/// <param name="ApcRoutine"></param>
/// <param name="lNetworkEvents"></param>
/// <returns></returns>
AFD_AsyncData* AsyncData = (AFD_AsyncData*)malloc(sizeof(AFD_AsyncData));
if (AsyncData == NULL)
{
   return -1;
}
memset(AsyncData, 0, sizeof(AFD_AsyncData));
AsyncData->NowSocket = Handle;
AsyncData->PollInfo.Timeout.HighPart = 0x7FFFFFFF;
AsyncData->PollInfo.Timeout.LowPart = 0xFFFFFFFF;
AsyncData->PollInfo.HandleCount = 1;
AsyncData->PollInfo.Handle = Handle;
AsyncData->PollInfo.Events = lNetworkEvents;
AsyncData->UserContext = UserContext;
NTSTATUS Status;
Status = ((NtDeviceIoControlFile)(MyNtDeviceIoControlFile))((HANDLE)Handle,
   NULL,
   ApcRoutine,
   AsyncData,
   &AsyncData->IOSB,
   IOCTL_AFD_SELECT,
   &AsyncData->PollInfo,
   sizeof(AFD_PollInfo),
   &AsyncData->PollInfo,
   sizeof(AFD_PollInfo));
return Status;
}
NTSTATUS WSPProcessAsyncSelect(
SOCKET Handle,
PVOID ApcRoutine,
ULONG lNetworkEvents,
PVOID UserContext
) {
/// <summary>
/// WSPProcessAsyncSelect是本程序最大的亮点，利用异步化的IOCTL_AFD_SELECT
/// </summary>
/// <param name="Handle"></param>
/// <param name="ApcRoutine"></param>
/// <param name="lNetworkEvents"></param>
/// <returns></returns>
AFD_AsyncData* AsyncData = (AFD_AsyncData*)malloc(sizeof(AFD_AsyncData));
if (AsyncData == NULL)
{
   return -1;
}
memset(AsyncData, 0, sizeof(AFD_AsyncData));
AsyncData->NowSocket = Handle;
AsyncData->PollInfo.Timeout.HighPart = 0x7FFFFFFF;
AsyncData->PollInfo.Timeout.LowPart = 0xFFFFFFFF;
AsyncData->PollInfo.HandleCount = 1;
AsyncData->PollInfo.Handle = Handle;
AsyncData->PollInfo.Events = lNetworkEvents;
AsyncData->UserContext = UserContext;
NTSTATUS Status;
Status = ((NtDeviceIoControlFile)(MyNtDeviceIoControlFile))((HANDLE)Handle,
   NULL,
   ApcRoutine,
   AsyncData,
   &AsyncData->IOSB,
   IOCTL_AFD_SELECT,
   &AsyncData->PollInfo,
   sizeof(AFD_PollInfo),
   &AsyncData->PollInfo,
   sizeof(AFD_PollInfo));
return Status;
}
VOID __stdcall internal_APCRoutine(
    PVOID ApcContext,
    PIO_STATUS_BLOCK IoStatusBlock,
    PVOID Reserved)
{
    /// <summary>
    /// 这是一个内部函数，也是本程序最大的亮点：APC异步select
    /// Client和服务器socket的select情况会全部调用这个函数，本函数用于分发回调事件。
    /// </summary>
    /// <param name="ApcContext"></param>
    /// <param name="IoStatusBlock"></param>
    /// <param name="Reserved"></param>
    /// <returns></returns>
    AFD_AsyncData* AsyncData = (AFD_AsyncData*)ApcContext;
    //读取出ApcContext，保存的是AsyncData，AsyncData的详细使用情况请看WSPProcessAsyncSelect函数的实现。
    WSPServer* self = (WSPServer*)AsyncData->UserContext;
    //读取出WSPServer对象，这个主要是读取出回调函数地址。
    //然后调用回调函数，参数是socket句柄和事件信息
    ((WSPServerCallBack)self->m_CallBack)(AsyncData->NowSocket, AsyncData->PollInfo.Events);
    //开启下一次APC异步select
    WSPProcessAsyncSelect(AsyncData->NowSocket, internal_APCRoutine, self->m_EnableEvent, (PVOID)self);
    //释放掉原来的AsyncData
    free(AsyncData);
}
 
 
VOID WSPServer::internal_APCThread(WSPServer* Server) {
    /// <summary>
    /// 本函数是APC异步select的线程函数，用于开启每个socket的APC异步select
    /// </summary>
    /// <param name="Server">WSPServer对象</param>
    int i = 0;
    while (1) {
        EnterCriticalSection(&Server->m_CriticalSection);
        if (!Server->IsRun) {
            //已经通知本线程退出
            LeaveCriticalSection(&Server->m_CriticalSection);
            break;
        }
        i = Server->m_NeedAPCSocket.size();
        i--;
        for (i; i >= 0; i--) {
            //将m_NeedAPCSocket每一socket读取出来，开启APC异步select
            WSPProcessAsyncSelect(
                Server->m_NeedAPCSocket[i],
                internal_APCRoutine,
                Server->m_EnableEvent,
                (PVOID)Server
            );
            //后面的APC异步select过程将由函数internal_APCRoutine完成
            Server->m_NeedAPCSocket.pop_back();
            //删除m_NeedAPCSocket对应的socket
        }
        LeaveCriticalSection(&Server->m_CriticalSection);
        SleepEx(1, true);
    }
}
 
HANDLE WSPServer::APCAsyncSelect(
    WSPServerCallBack*    ApcCallBack,
    int lNetworkEvents
) {
    /// <summary>
    /// 开启服务器的APC异步select模式，只能初始化一次
    /// </summary>
    /// <param name="ApcCallBack">回调函数</param>
    /// <param name="lNetworkEvents">需要异步处理的事件</param>
    /// <returns>返回异步处理线程的句柄</returns>
    if (this->m_CallBack != NULL) {
        return INVALID_HANDLE_VALUE;
    }
    EnterCriticalSection(&this->m_CriticalSection);
    this->m_EnableEvent = lNetworkEvents;
    this->m_CallBack = ApcCallBack;
    //下面将把m_AllClientSocket已有的句柄拷贝到m_NeedAPCSocket
    std::copy(m_AllClientSocket.begin(), m_AllClientSocket.end(), m_NeedAPCSocket.begin());
    //将服务器socket加入m_NeedAPCSocket
    m_NeedAPCSocket.push_back(this->m_socket);
    //启动线程
    m_ThreadHandle = CreateThread(NULL, 0, (PTHREAD_START_ROUTINE)internal_APCThread, this, 0, NULL);
    LeaveCriticalSection(&this->m_CriticalSection);
    return m_ThreadHandle;
}
VOID __stdcall internal_APCRoutine(
    PVOID ApcContext,
    PIO_STATUS_BLOCK IoStatusBlock,
    PVOID Reserved)
{

[培训]内核驱动高级班，冲击BAT一流互联网大厂工作，每周日13:00-18:00直播授课

最后于 2022-9-9 08:50 被ANormalUser编辑 ，原因： githun图片有时不能正常显示
上传的附件：
收藏
免费 ・6
支持
分享
最新回复 (16)
wangren
雪    币： 325
活跃值： (665)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
2 楼
卖1楼
2022-9-10 17:37
0
iamasbcx
雪    币： 3868
活跃值： (3643)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
3 楼
大佬牛逼 但是不知道有什么用
2022-9-10 19:39
0
咖啡_741298
雪    币： 6
活跃值： (3290)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
4 楼
其实不如IOCP简洁
2022-9-10 21:00
0
我skyddr
雪    币： 9
能力值： ( LV1，RANK：0 )
在线值：
发帖
回帖
粉丝
私信
5 楼
没看懂，这个有什么作用，老大的意思是说比应用层调用的函数更快？效率更高？能负载1万连接不卡吗？还是怎么样，性能和iocp比较呢？
2022-9-10 23:00
0
我skyddr
雪    币： 9
能力值： ( LV1，RANK：0 )
在线值：
发帖
回帖
粉丝
私信
6 楼
@ANormalUser
2022-9-10 23:00
0
ANormalUser
雪    币： 3870
活跃值： (3122)
能力值： ( LV3，RANK：20 )
在线值：
发帖
回帖
粉丝
私信
7 楼
我skyddr 没看懂，这个有什么作用，老大的意思是说比应用层调用的函数更快？效率更高？能负载1万连接不卡吗？还是怎么样，性能和iocp比较呢？
怎么说呢，IOCP其实也是用的AsyncSelect实现的，只是稍微封装了一下。
例如：当你使用IOCP绑定了socket之后，AcceptEx在调用NtDeviceIoControlFile时其实也是了APC，这个APC会通知到一个内部函数，然后内部函数再将异步请求结果插入IOCP的队列。
所以，我觉得如果知道了这一点后，这个异步过程完全可以自己实现，不必走系统的臃肿代码，于是就有了上面的WSPProcessAsyncSelect过程。
比直接效率理论上优于IOCP，但是比较IOCP是异步的，它的队列处理端可以写得更好。我这里的WSPProcessAsyncSelect的异步过程完全用一个单线程处理，过大处理可能会造成队列阻塞，不过想改进这部分也并不困难。
由于程序基于APC异步实现，所以不会出现轮询模型一类的超多连接后就会效率下降的问题。
我这里测试1万Client还不错，不过具体问题还得具体分析吧。
最后：系统的socket进行的很多处理（例如：储存每个socket对应信息的结构体，通知AFD等不必要操作等），而本程序可以说是最基础的一个socket功能，什么多余的部分也没有，相比原版socket命令的效率会高很多，光send函数对比就看得出来（但是问题是：由于太简单了，有些信息需要自己保存，略微麻烦）
2022-9-10 23:57
0
XieAPP
雪    币： 41
活跃值： (262)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
8 楼
以稳定性,兼容性换速度?
2022-9-12 21:01
0
ANormalUser
雪    币： 3870
活跃值： (3122)
能力值： ( LV3，RANK：20 )
在线值：
发帖
回帖
粉丝
私信
9 楼
XieAPP 以稳定性,兼容性换速度?
也不完全是，这些因素基本由使用情况决定。
2022-9-12 21:53
0
木志本柯
雪    币： 4714
活跃值： (4250)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
10 楼
收藏了以后慢慢抄 啊哈哈哈
2022-9-22 11:01
0
zhenwo
雪    币： 1555
活跃值： (3103)
能力值： ( LV11，RANK：180 )
在线值：
发帖
回帖
粉丝
私信
11 楼
2022-9-24 22:13
0
大鲤鱼
雪    币： 6307
活跃值： (3837)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
12 楼
我在客户进入的地方用WSPGetPeerName， 一直返回0啊……取不到客户端的IP地址。
然后，如果客户端意外退出，比如关机蓝屏之类的。服务端识别不到客户端退出。
2022-11-24 11:14
0
大鲤鱼
雪    币： 6307
活跃值： (3837)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
13 楼
sockaddr_in4 addr = { 0 };
SOCKET news = MyServer.AcceptClient(&addr);
我在这里传参，但是返回的IP地址错误。IP地址前面几个字节是乱码。
2022-11-24 14:37
0
我skyddr
雪    币： 9
能力值： ( LV1，RANK：0 )
在线值：
发帖
回帖
粉丝
私信
14 楼
@ANormalUser
你好楼主，我学习iocp很久了，有很重要的事告诉你，请私信方式联系我。我发不了消息 
2023-2-2 23:12
0
我skyddr
雪    币： 9
能力值： ( LV1，RANK：0 )
在线值：
发帖
回帖
粉丝
私信
15 楼
@ANormalUser
2023-2-5 14:43
0
我skyddr
雪    币： 9
能力值： ( LV1，RANK：0 )
在线值：
发帖
回帖
粉丝
私信
16 楼
@ANormalUser
在吗
2023-2-6 10:59
0
wsgaoshou
雪    币： 195
活跃值： (108)
能力值： ( LV2，RANK：10 )
在线值：
发帖
回帖
粉丝
私信
17 楼
应该是以速度，换兼容性
2023-4-30 12:30
0
游客
登录 | 注册 方可回帖
回帖 表情 雪币赚取及消费
高级回复
返回
//