WinSock 重叠I/O 服务端模型

#include "stdafx.h"
#include <WinSock2.h>
#include <mswsock.h>
#include <malloc.h>
#pragma comment(lib,"Ws2_32.lib")
#define BUFFER_SIZE 1024

void InitSock()
{
 WORD wVersionRequested;
 WSADATA wsaData;
 int err;
 
 wVersionRequested = MAKEWORD( 2, 2 );
 
 err = WSAStartup( wVersionRequested, &wsaData );
 if ( err != 0 ) {
  /* Tell the user that we could not find a usable */
  /* WinSock DLL.                                  */
  return;
 }
 
 /* Confirm that the WinSock DLL supports 2.2.*/
 /* Note that if the DLL supports versions greater    */
 /* than 2.2 in addition to 2.2, it will still return */
 /* 2.2 in wVersion since that is the version we      */
 /* requested.                                        */
 
 if ( LOBYTE( wsaData.wVersion ) != 2 ||
   HIBYTE( wsaData.wVersion ) != 2 ) {
  /* Tell the user that we could not find a usable */
  /* WinSock DLL.                                  */
  WSACleanup( );
  return;
 }
}

//socket 对象
typedef struct _SOCKET_OBJ
{
 SOCKET s;    //套接字句柄
 int nOutstandingOpts;  //记录此套接字上的重叠I/O数量
 LPFN_ACCEPTEX lpfnAcceptEx; //扩展函数AcceptEx的指针(仅对监听套接字而言)
}SOCKET_OBJ,*pSOCKET_OBJ;

//申请套接字对象
pSOCKET_OBJ GetSocketObj(SOCKET s)
{
 pSOCKET_OBJ pSocketObj=(pSOCKET_OBJ)malloc(sizeof(SOCKET_OBJ));
 if(pSocketObj!=NULL)
 {
  pSocketObj->s=s;
  pSocketObj->nOutstandingOpts=0;
  pSocketObj->lpfnAcceptEx=NULL;
 }
 return pSocketObj;
}

//释放套接字对象
void FreeSocketObj(pSOCKET_OBJ pSocketObj)
{
 if(INVALID_SOCKET!=pSocketObj->s)
 {
  closesocket(pSocketObj->s);
 }
 free(pSocketObj);
 pSocketObj=NULL;
}

//缓冲区对象
typedef struct _BUFFER_OBJ
{
 OVERLAPPED ol;   //重叠结构
 char *buf;       //send,recv,accept 缓冲区地址
 int nLen;        //buf的长度
 pSOCKET_OBJ pSocket;  //套接字对象
 int nOperation;   //操作类型
#define OP_ACCEPT 1
#define OP_READ 2
#define OP_WRITE 3
 SOCKET sAccept;   //用来保存acceptex接受的客户套接字(仅对监听套接字而言)
 _BUFFER_OBJ *pNext;   //指向下一个缓冲区对象
}BUFFER_OBJ,*pBUFFER_OBJ;

WSAEVENT events[WSA_MAXIMUM_WAIT_EVENTS];  //I/O事件句柄数组
int g_nBufCount=0;    // 数组中有句柄数量
pBUFFER_OBJ g_pBufferHead,g_pBufferTail; //缓冲区列表头和尾指针

//申请一个缓冲区对象
pBUFFER_OBJ GetBufferObj(pSOCKET_OBJ pSocket)
{
 if(g_nBufCount>WSA_MAXIMUM_WAIT_EVENTS-1)
 {
  printf("too many connections/n");
  return NULL;
 }
 pBUFFER_OBJ pBuffer=(pBUFFER_OBJ)malloc(sizeof(BUFFER_OBJ));
 if(pBuffer!=NULL)
 {
  pBuffer->buf=(char *)malloc(BUFFER_SIZE);
  pBuffer->pSocket=pSocket;
  pBuffer->ol.hEvent=WSACreateEvent();
  pBuffer->pNext=NULL;
  pBuffer->sAccept=INVALID_SOCKET;

  //将pBuffer添加到缓冲区列表中
  if(g_pBufferHead==NULL)
  {
   g_pBufferHead=g_pBufferTail=pBuffer;
  }
  else
  {
   g_pBufferTail->pNext=pBuffer;
   g_pBufferTail=pBuffer;
  }
  events[++g_nBufCount]=pBuffer->ol.hEvent;
 }
 return pBuffer;
}

//释放缓冲区对象
void FreeBufferObj(pBUFFER_OBJ pBuffer)
{
 pBUFFER_OBJ pTemp=g_pBufferHead;
 BOOL bFind=false;
 if(pTemp==pBuffer)
 {
  if(g_pBufferHead==g_pBufferTail)
  {
   g_pBufferHead=g_pBufferTail=NULL;
  }
  else
  {
   g_pBufferHead=pBuffer->pNext;
  }
  bFind=true;
 }
 else
 {
  while(pTemp!=NULL&&pTemp->pNext!=pBuffer)
  {
   pTemp=pTemp->pNext;
  }
  if(pTemp!=NULL)
  {
   pTemp->pNext=pBuffer->pNext;
   if(pTemp->pNext==NULL)
    g_pBufferTail=pTemp;
   bFind=true;
  }
 }
 if(bFind)
 {
  g_nBufCount--;
  ::CloseHandle(pBuffer->ol.hEvent);
  free(pBuffer->buf);
  free(pBuffer);
 }
}

pBUFFER_OBJ FindBufferObj(WSAEVENT hEvent)
{
 pBUFFER_OBJ pBuffer=g_pBufferHead;
 while(pBuffer!=NULL)
 {
  if(pBuffer->ol.hEvent==hEvent)
  {
   break;
  }
  pBuffer=pBuffer->pNext;
 }
 return pBuffer;
}

void RebuildArray()
{
 pBUFFER_OBJ pBufferObj=g_pBufferHead;
 int i=1;
 while(pBufferObj!=NULL)
 {
  events[i++]=pBufferObj->ol.hEvent;
  pBufferObj=pBufferObj->pNext;
 }
}
BOOL PostAccept(pBUFFER_OBJ pBuffer)
{
 pSOCKET_OBJ pSocket=pBuffer->pSocket;
 if(pSocket->lpfnAcceptEx!=NULL)
 {
  pBuffer->nOperation=OP_ACCEPT;
  pSocket->nOutstandingOpts++;
  DWORD dwBytes;
  pBuffer->sAccept=WSASocket(AF_INET,SOCK_STREAM,0,NULL,0,WSA_FLAG_OVERLAPPED);
  BOOL b=pSocket->lpfnAcceptEx(pSocket->s,pBuffer->sAccept,pBuffer->buf,BUFFER_SIZE-(sizeof(sockaddr)+16)*2,
   sizeof(sockaddr)+16,sizeof(sockaddr)+16,&dwBytes,&pBuffer->ol);
  if(!b)
  {
   if(WSA_IO_PENDING!=WSAGetLastError())
   {
    return false;
   }
  }
  return true;
 }
 return false;
}

BOOL PostRecv(pBUFFER_OBJ pBuffer)
{
 pSOCKET_OBJ pSocket=pBuffer->pSocket;
 pSocket->nOutstandingOpts++;
 pBuffer->nOperation=OP_READ;
 DWORD dwBytes;
 DWORD dwFlags=0;
 WSABUF buf;
 buf.buf=pBuffer->buf;
 buf.len=BUFFER_SIZE;
 if(SOCKET_ERROR!=WSARecv(pSocket->s,&buf,1,&dwBytes,&dwFlags,&pBuffer->ol,NULL))
 {
  if(WSA_IO_PENDING!=WSAGetLastError())
  {
   return false;
  }
 }
 return true;
}

//BOOL PostSend(...){....}

BOOL HandleIO(pBUFFER_OBJ pBuffer)
{
 pSOCKET_OBJ pSocket=pBuffer->pSocket;
 pSocket->nOutstandingOpts--;
 DWORD dwTrans;
 DWORD dwFlags;
 BOOL bRet=::WSAGetOverlappedResult(pSocket->s,&pBuffer->ol,&dwTrans,false,&dwFlags);
 if(!bRet)
 {
  if(pSocket->s!=INVALID_SOCKET)
  {
   closesocket(pSocket->s);
   pSocket->s=INVALID_SOCKET;
  }
  if(pSocket->nOutstandingOpts==0)
  {
   FreeSocketObj(pSocket);
  }
  FreeBufferObj(pBuffer);
  return false;
 }
 pSOCKET_OBJ pClient;
 pBUFFER_OBJ pRecv;
 switch(pBuffer->nOperation)
 {
 case OP_ACCEPT:
  pClient=GetSocketObj(pBuffer->sAccept);
  pRecv=GetBufferObj(pClient);
  if(pRecv==NULL)
  {
   return false;
  }
  RebuildArray();
  if(!PostRecv(pRecv))
  {
   ::FreeSocketObj(pClient);
   ::FreeBufferObj(pRecv);
   return false;
  }
  printf("一新连接/n");
  pBuffer->buf[dwTrans]='/0';
  printf("收到数据:%s/n",pBuffer->buf);
  ::ZeroMemory(pBuffer->buf,BUFFER_SIZE);
  PostAccept(pBuffer);
  break;
 case OP_READ:
  if(dwTrans>0)
  {
   pBuffer->buf[dwTrans]='/0';
   printf("收到数据:%s/n",pBuffer->buf);
   ::ZeroMemory(pBuffer->buf,BUFFER_SIZE);
   if(!PostRecv(pBuffer))
   {
    printf("继开一连接/n");
    ::closesocket(pSocket->s);
    pSocket->s=INVALID_SOCKET;
    ::FreeBufferObj(pRecv);
    return false;
   }
  }
  else
  {
   //必须关闭socket 以便
   printf("继开一连接/n");
    
   if(pSocket->s!=INVALID_SOCKET)
   {
    ::closesocket(pSocket->s);
    pSocket->s=INVALID_SOCKET;
   }
   if(pSocket->nOutstandingOpts==0)
   {
    FreeSocketObj(pSocket);
   }
   ::FreeBufferObj(pBuffer);
   return false;
  }
  break;
 }
 return true;
}
int _tmain(int argc, _TCHAR* argv[])
{
 InitSock();
 SOCKET sListen=WSASocket(AF_INET,SOCK_STREAM,0,NULL,0,WSA_FLAG_OVERLAPPED);
 sockaddr_in addr;
 addr.sin_addr.S_un.S_addr=INADDR_ANY;
 addr.sin_family=AF_INET;
 addr.sin_port=htons(1456);
 if(SOCKET_ERROR==bind(sListen,(sockaddr *)&addr,sizeof(sockaddr)))
 {
  printf("绑定失败/n");
  ::closesocket(sListen);
  ::WSACleanup();
  return 0;
 }
 listen(sListen,15);
 pSOCKET_OBJ pListen=GetSocketObj(sListen);
 GUID guid=WSAID_ACCEPTEX;
 DWORD dwBytes;
 ::WSAIoctl(pListen->s,SIO_GET_EXTENSION_FUNCTION_POINTER,&guid,sizeof(guid),&pListen->lpfnAcceptEx,sizeof(pListen->lpfnAcceptEx),
  &dwBytes,NULL,NULL);
 events[0]=WSACreateEvent();
 for(int i=0;i<5;i++)
 {
  BOOL B=PostAccept(GetBufferObj(pListen));
 }
 ::WSASetEvent(events[0]);
 while(true)
 {
  int nIndex=::WSAWaitForMultipleEvents(g_nBufCount+1,events,false,WSA_INFINITE,false);
  if(nIndex==WSA_WAIT_FAILED)
  {
   printf("wait failed/n");
   break;
  }
  nIndex=nIndex-WSA_WAIT_EVENT_0;
  for(int i=nIndex;i<g_nBufCount+1;i++)
  {
   int nRet=::WSAWaitForMultipleEvents(1,&events[i],true,0,false);
   if(nRet==WSA_WAIT_TIMEOUT||nRet==WSA_WAIT_FAILED)
   {
    continue;
   }
   else
   {
    ::WSAResetEvent(events[i]);
    if(i==0)
    {
     ::RebuildArray();
     continue;
    }
    pBUFFER_OBJ pBuffer=::FindBufferObj(events[i]);
    if(pBuffer!=NULL)
    {
     if(!HandleIO(pBuffer))
     {
      ::RebuildArray();
     }
    }
   }
  }
 }
 return 0;
}