IOCP 配合 AcceptEx

#include "SOCKET.h"
#include <Windows.h>
DWORD WINAPI ThreadProc(LPVOID pvParam);
#define PORT 8080
#define LISTEN_QUEUE 200
// AcceptEx 和 GetAcceptExSockaddrs 的函数指针，用于调用这两个扩展函数
LPFN_ACCEPTEX                lpfnAcceptEx;                
LPFN_GETACCEPTEXSOCKADDRS lpfnGetAcceptExSockAddrs; 
void PostAcceptEx(IOCPHandle_s & listenHandle)
{
IO_DATA_s * p_io_data = new IO_DATA_s;
p_io_data->socket = socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);
listenHandle.Push(p_io_data);
p_io_data->type = ACCEPT;
lpfnAcceptEx(listenHandle.socket, p_io_data->socket, &p_io_data->addr, 0, 0, sizeof(SOCKADDR_IN) + 16, &p_io_data->len, &p_io_data->ol);
}
int main()
{
SocketInit();
IOCPHandle_s listenHandle;
listenHandle.socket = socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);
listenHandle.addr.sin_family = AF_INET;
listenHandle.addr.sin_addr.S_un.S_addr = htonl(INADDR_ANY);
listenHandle.addr.sin_port = htons(PORT);
HANDLE IOCPhandle;
IOCPhandle = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0 );
::CreateThread(0, 0, ThreadProc, (void *)IOCPhandle, 0, 0);
::bind(listenHandle.socket, (SOCKADDR *)&listenHandle.addr, sizeof(SOCKADDR));
::listen(listenHandle.socket, LISTEN_QUEUE);
CreateIoCompletionPort((HANDLE)listenHandle.socket, IOCPhandle, (unsigned long)&listenHandle, 0);
// 使用AcceptEx函数，因为这个是属于WinSock2规范之外的微软另外提供的扩展函数
// 所以需要额外获取一下函数的指针，
// 获取AcceptEx函数指针
GUID GuidAcceptEx = WSAID_ACCEPTEX;  
GUID GuidGetAcceptExSockAddrs = WSAID_GETACCEPTEXSOCKADDRS; 
DWORD dwBytes = 0;  
WSAIoctl(
listenHandle.socket, 
SIO_GET_EXTENSION_FUNCTION_POINTER, 
&GuidAcceptEx, 
sizeof(GuidAcceptEx), 
&lpfnAcceptEx, 
sizeof(lpfnAcceptEx), 
&dwBytes, 
NULL, 
NULL);
// 获取GetAcceptExSockAddrs函数指针，也是同理
WSAIoctl(
listenHandle.socket, 
SIO_GET_EXTENSION_FUNCTION_POINTER, 
&GuidGetAcceptExSockAddrs,
sizeof(GuidGetAcceptExSockAddrs), 
&lpfnGetAcceptExSockAddrs, 
sizeof(lpfnGetAcceptExSockAddrs),   
&dwBytes, 
NULL, 
NULL);
for(i32 i = 0 ; i < 50 ; ++i)
PostAcceptEx(listenHandle);
printf("主线程阻塞\n");
Sleep(INFINITE);
SocketUnInit();
getchar();
return 0;
}
DWORD WINAPI ThreadProc(LPVOID pvParam)
{
HANDLE IOCPhandle = pvParam;
DWORD   dwBytesTransfered = 0;
IOCPHandle_s * p_IOCPhandle;
IO_DATA_s * p_IOdata;
BOOL bReturn;
while(true)
{
//实际操作字节数
bReturn = GetQueuedCompletionStatus(IOCPhandle,&dwBytesTransfered,(unsigned long *)&p_IOCPhandle,(OVERLAPPED **)&p_IOdata,INFINITE);
if(!bReturn)
{
printf("出错\n");
closesocket(p_IOCPhandle->socket);
delete p_IOCPhandle;
delete p_IOdata;
continue;
}
else if(dwBytesTransfered == 0 && (p_IOdata->type == RECV || p_IOdata->type == SEND))
{
printf("客户端关闭了连接!\n");
closesocket(p_IOCPhandle->socket);
delete p_IOCPhandle;
delete p_IOdata;
}
else
{
printf("处理IO服务\n");
switch(p_IOdata->type)
{
case RECV:
{
printf("RECV\n");
p_IOdata->buf.buf[dwBytesTransfered] = '\0';
printf("数据:%s 长度:%d\n",p_IOdata->buf.buf,dwBytesTransfered);
p_IOdata->SetZero();
p_IOdata->type = RECV;
p_IOdata->buf.buf = p_IOCPhandle->recv_buff;
p_IOdata->buf.len = BUFF_SIZE;
p_IOdata->flags = 0;
int nBytesRecv = WSARecv(p_IOCPhandle->socket, &p_IOdata->buf, 1, &p_IOdata->len, &p_IOdata->flags, &p_IOdata->ol, 0);
int ret = WSAGetLastError();
if ((SOCKET_ERROR == nBytesRecv) && (WSA_IO_PENDING != ret))
{
printf("投递RECV失败\n");
}
}
break;
case SEND:
{

}
break;
case ACCEPT:
{
printf("ACCEPT\n");
IOCPHandle_s * p_accepthandle = new IOCPHandle_s;
p_accepthandle->socket = p_IOdata->socket;
//赋值SOCKET的属性与LISTENSOCKET的属性一样
setsockopt( p_accepthandle->socket,
SOL_SOCKET,
SO_UPDATE_ACCEPT_CONTEXT,
(char*)&p_IOCPhandle->socket,
sizeof(p_IOCPhandle->socket) ) ;
//lpfnGetAcceptExSockAddrs(&p_IOdata->buf.buf,0, 0, sizeof(SOCKADDR_IN) + 16, 0, 0, (LPSOCKADDR*)&p_addr, &remoteLen);
//memcpy(&p_accepthandle->addr, &p_IOdata->addr, sizeof(SOCKADDR_IN) + 16);
//printf("IP:%s 端口:%d\n",inet_ntoa(p_accepthandle->addr.sin_addr),ntohs(p_accepthandle->addr.sin_port));
IO_DATA_s * p_iodata = new IO_DATA_s;
p_accepthandle->Push(p_iodata);
CreateIoCompletionPort((HANDLE)p_accepthandle->socket, IOCPhandle, (unsigned long)p_accepthandle, 0);
p_iodata->type = RECV;
p_iodata->buf.buf = p_accepthandle->recv_buff;
p_iodata->buf.len = BUFF_SIZE;
p_iodata->flags = 0;
int nBytesRecv = WSARecv(p_accepthandle->socket, &p_iodata->buf, 1, &p_iodata->len, &p_iodata->flags, &p_iodata->ol, 0);
if ((SOCKET_ERROR == nBytesRecv) && (WSA_IO_PENDING != WSAGetLastError()))
{
printf("投递RECV失败\n");
return false;
}


p_IOdata->SetZero();
p_IOdata->socket = socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);
p_IOdata->type = ACCEPT;
lpfnAcceptEx(p_IOCPhandle->socket, p_IOdata->socket, &p_IOdata->addr, 0, 0, sizeof(SOCKADDR_IN) + 16, &p_IOdata->len, &p_IOdata->ol);
}
break;
default:
break;
}
}
 
}
return 0;

}

/*******************************************************SOCKET.h************************************************/

#include <winsock2.h>
#include <ws2tcpip.h>
#include <fcntl.h>
#include <MSWSock.h>
#include <vector>
#include <string.h>
#pragma comment(lib ,"ws2_32")
void SocketInit()
{
WSADATA wsaData;// 库
WSAStartup(MAKEWORD(2,2),&wsaData); //初始化
}
void SocketUnInit()
{
WSACleanup();//关闭
}
#define BUFF_SIZE 120
enum TYPE {SEND,RECV,ACCEPT,INIT};
typedef int i32;
typedef char i8;
typedef unsigned int u32;
typedef unsigned char u8;


struct IO_DATA_s
{
IO_DATA_s()
{
type = INIT;
SetZero();
}
WSAOVERLAPPED ol;
SOCKET socket; //操作的SOCKET
SOCKADDR_IN addr; //地址
DWORD len; //操作长度
DWORD flags; //标志
TYPE type; //操作类型
WSABUF buf; //WindowBUF
void SetZero()
{
memset(this, 0, sizeof(IO_DATA_s));
}
};
struct IOCPHandle_s
{
SOCKET socket; //套接字
SOCKADDR_IN addr; //地址
std::vector<IO_DATA_s *> io_array; //IO操作的地址
i8 send_buff[BUFF_SIZE]; //发送缓存
DWORD send_len;
char recv_buff[BUFF_SIZE]; //接收缓存
DWORD recv_len;
void Push(IO_DATA_s * v)
{
io_array.push_back(v);
}
};