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;
}
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;
}