IOCP实现聊天服务

#include <stdio.h>

#include <stdlib.h>

#include <process.h>

#include <string.h>

#include <winsock2.h>

#pragma comment(lib, "ws2_32.lib")

#pragma warning(disable: 4996)

#define BUF_SIZE 1024

#define READ 1

#define WRITE 2

typedef struct // socket info

{

SOCKET hClntSock;

SOCKADDR_IN clntAddr;

} PER_HANDLE_DATA, *LPPER_HANDLE_DATA;

typedef struct // buffer info

{

OVERLAPPED overlapped;

WSABUF wsaBuf;

char buffer[BUF_SIZE];

int rwMode; // READ or WRITE

} PER_IO_DATA, *LPPER_IO_DATA;

typedef struct

{

LPPER_HANDLE_DATA lpSockInfoSet[BUF_SIZE];

DWORD numofsock;

} SOCK_DATA_SET, *LPSOCK_DATA_SET;

SOCK_DATA_SET sockSet;

BOOL AddSock(LPPER_HANDLE_DATA lpSockInfo);

BOOL DeleteSock(LPPER_HANDLE_DATA lpSockInfo);

BOOL SendToAll(LPPER_IO_DATA lpIOInfo, DWORD bytesTrans);

void ErrorHandling(const char* message);

DWORD WINAPI IOHandlingThread(LPVOID CompletionPortIO);

void CreateNewWriteInfoFromReadInfo(LPPER_IO_DATA* lpSrc, LPPER_IO_DATA* lpDes, DWORD bytesTrans);

int main(int argc, char* argv[])

{

memset(&sockSet, 0, sizeof(SOCK_DATA_SET));

WSADATA wsaData;

HANDLE hComPort;

LPPER_HANDLE_DATA lpSockInfo;

LPPER_IO_DATA lpIOInfo;

SYSTEM_INFO sysInfo;

SOCKET hListenSock, hClientSock;

SOCKADDR_IN listenAddr, clientAddr;

int addrSize;

DWORD recvBytes, flags = 0;

if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)

ErrorHandling("WSAStartup() Error");

if (argc != 2)

{

printf("Usage: %s <port> \n", argv[0]);

exit(1);

}

hComPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);

GetSystemInfo(&sysInfo);

for (int i = 0; i < sysInfo.dwNumberOfProcessors; i++)

_beginthreadex(NULL, 0, IOHandlingThread, (LPVOID)hComPort, 0, NULL);

hListenSock = WSASocket(PF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);

memset(&listenAddr, 0, sizeof(listenAddr));

listenAddr.sin_family = AF_INET;

listenAddr.sin_addr.S_un.S_addr = htonl(INADDR_ANY);

listenAddr.sin_port = htons(atoi(argv[1]));

if (bind(hListenSock, (SOCKADDR*)&listenAddr, sizeof(listenAddr)) != 0)

ErrorHandling("bind() Error");

if (listen(hListenSock, 5) != 0)

ErrorHandling("listen() Error");

while (1)

{

addrSize = sizeof(clientAddr);

hClientSock = accept(hListenSock, (SOCKADDR*)&clientAddr, &addrSize);

if (hClientSock == INVALID_SOCKET)

ErrorHandling("accept() error");

lpSockInfo = (LPPER_HANDLE_DATA)malloc(sizeof(PER_HANDLE_DATA));

lpSockInfo->hClntSock = hClientSock;

lpSockInfo->clntAddr = clientAddr;

CreateIoCompletionPort((HANDLE)hClientSock, hComPort, (DWORD)lpSockInfo, 0);

lpIOInfo = (LPPER_IO_DATA)malloc(sizeof(PER_IO_DATA));

memset(&(lpIOInfo->overlapped), 0, sizeof(PER_IO_DATA));

lpIOInfo->wsaBuf.len = BUF_SIZE;

lpIOInfo->wsaBuf.buf = lpIOInfo->buffer;

lpIOInfo->rwMode = READ;

WSARecv(lpSockInfo->hClntSock, &(lpIOInfo->wsaBuf), 1, &recvBytes, &flags, &(lpIOInfo->overlapped), NULL);

AddSock(lpSockInfo);

}

return 0;

}

BOOL AddSock(LPPER_HANDLE_DATA lpSockInfo)

{

int numofsock = sockSet.numofsock;

sockSet.lpSockInfoSet[numofsock] = lpSockInfo;

sockSet.numofsock++;

return TRUE;

}

BOOL DeleteSock(LPPER_HANDLE_DATA lpSockInfo)

{

int numofsock = sockSet.numofsock;

LPPER_HANDLE_DATA* sockInfoArr = sockSet.lpSockInfoSet;

for (int i = 0; i < numofsock; i++)

{

if (lpSockInfo == sockInfoArr[i])

{

for (int next = i + 1; next < numofsock; next++)

sockInfoArr[i] = sockInfoArr[next];

sockSet.numofsock--;

}

}

return TRUE;

}

void ErrorHandling(const char* message)

{

fprintf(stderr, "%s, Error Code %d \n", message, WSAGetLastError());

exit(1);

}

DWORD WINAPI IOHandlingThread(LPVOID CompletionPortIO)

{

HANDLE hComPort = (HANDLE)CompletionPortIO;

SOCKET sock;

DWORD bytesTrans;

LPPER_IO_DATA lpIOInfo;

LPPER_HANDLE_DATA lpSockInfo;

DWORD flags = 0;

while (1)

{

GetQueuedCompletionStatus(hComPort, &bytesTrans, (LPDWORD)& lpSockInfo, (LPOVERLAPPED*)& lpIOInfo, INFINITE);

sock = lpSockInfo->hClntSock;

if (lpIOInfo->rwMode == READ)

{

puts("message recvived");

if (bytesTrans == 0) // EOF

{

DeleteSock(lpSockInfo);

free(lpSockInfo);

free(lpIOInfo);

continue;

}

SendToAll(lpIOInfo, bytesTrans);

lpIOInfo = (LPPER_IO_DATA)malloc(sizeof(PER_IO_DATA));

memset(&(lpIOInfo->overlapped), 0, sizeof(PER_IO_DATA));

lpIOInfo->wsaBuf.len = BUF_SIZE;

lpIOInfo->wsaBuf.buf = lpIOInfo->buffer;

lpIOInfo->rwMode = READ;

WSARecv(sock, &(lpIOInfo->wsaBuf), 1, NULL, &flags, &(lpIOInfo->overlapped), NULL);

}

else

{

puts("message send!");

free(lpIOInfo);

}

}

return 0;

}

BOOL SendToAll(LPPER_IO_DATA lpIOInfo, DWORD bytesTrans)

{

LPPER_IO_DATA lpReadInfo = lpIOInfo;

LPPER_IO_DATA lpWriteInfo;

DWORD sendBytes;

for (int i = 0; i < sockSet.numofsock; i++)

{

SOCKET sock = sockSet.lpSockInfoSet[i]->hClntSock;

CreateNewWriteInfoFromReadInfo(&lpReadInfo, &lpWriteInfo, bytesTrans);

WSASend(sock, &(lpWriteInfo->wsaBuf), 1, &sendBytes, 0, &(lpWriteInfo->overlapped), NULL);

}

return TRUE;

}

void CreateNewWriteInfoFromReadInfo(LPPER_IO_DATA* lpSrc, LPPER_IO_DATA* lpDes, DWORD bytesTrans)

{

*lpDes = (LPPER_IO_DATA)malloc(sizeof(PER_IO_DATA));

LPPER_IO_DATA src = *lpSrc;

LPPER_IO_DATA des = *lpDes;

memcpy(des, src, sizeof(PER_IO_DATA));

des->wsaBuf.buf = des->buffer;

des->wsaBuf.len = bytesTrans;

des->rwMode = WRITE;

}

客户端实现

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <windows.h>

#include <process.h>

#pragma comment(lib, "ws2_32.lib")

#pragma warning(disable: 4996)

#define BUF_SIZE 100

#define NAME_SIZE 20

unsigned WINAPI SendMsg(void * arg);

unsigned WINAPI RecvMsg(void * arg);

void ErrorHandling(char * msg);

char name[NAME_SIZE] = "[DEFAULT]";

char msg[BUF_SIZE];

int main(int argc, char* argv[])

{

WSADATA wsaData;

SOCKET hSock;

SOCKADDR_IN servAddr;

HANDLE hSendThread, hRecvThread;

if (argc != 4)

{

printf("Usage: %s <IP> <port> <name> \n", argv[0]);

exit(1);

}

if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)

ErrorHandling("WSAStartup() error");

sprintf(name, "[%s]", argv[3]);

hSock = socket(PF_INET, SOCK_STREAM, 0);

memset(&servAddr, 0, sizeof(servAddr));

servAddr.sin_family = AF_INET;

servAddr.sin_addr.S_un.S_addr = inet_addr(argv[1]);

servAddr.sin_port = htons(atoi(argv[2]));

if (connect(hSock, (SOCKADDR*)& servAddr, sizeof(servAddr)) != 0)

ErrorHandling("connect() error");

hSendThread = (HANDLE)_beginthreadex(NULL, 0, SendMsg, (void*)&hSock, 0, NULL);

hRecvThread = (HANDLE)_beginthreadex(NULL, 0, RecvMsg, (void*)&hSock, 0, NULL);

WaitForSingleObject(hSendThread, INFINITE);

WaitForSingleObject(hRecvThread, INFINITE);

closesocket(hSock);

WSACleanup();

return 0;

}

unsigned WINAPI SendMsg(void* arg)

{

SOCKET hSock = *((SOCKET*)arg);

char nameMsg[NAME_SIZE+BUF_SIZE];

while (1)

{

fgets(msg, BUF_SIZE, stdin);

if (!strcmp(msg, "q\n") || !strcmp(msg, "Q\n"))

{

closesocket(hSock);

exit(0);

}

sprintf(nameMsg, "%s %s", name, msg);

send(hSock, nameMsg, strlen(nameMsg), 0);

}

return 0;

}

unsigned WINAPI RecvMsg(void* arg)

{

SOCKET hSock = *((SOCKET*)arg);

char nameMsg[NAME_SIZE+BUF_SIZE];

int strLen;

while (1)

{

strLen = recv(hSock, nameMsg, NAME_SIZE+BUF_SIZE-1, 0);

if (strLen == -1)

return -1;

nameMsg[strLen] = '\0';

fputs(nameMsg, stdout);

}

return 0;

}

void ErrorHandling(char* msg)

{

fprintf(stderr, "%s, Error Code %d \n", msg, WSAGetLastError());

exit(1);

}