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