[C] 소켓 공부하기!!

#include 
#include 
#include 

#define BUFSIZE    512

// 소켓 정보 저장을 위한 구조체
struct SOCKETINFO
{
	SOCKET sock;
	char buf[BUFSIZE+1];
	int recvbytes;
	int sendbytes;
};
///////////////////////////////////////////////////////////////////////////////////////
struct THREADINFO
{
	DWORD ThreadId;
	int Total_Sockets;
	SOCKETINFO *SocketInfoArray[FD_SETSIZE];
	THREADINFO *next;
};

THREADINFO *ThreadInfoList;

BOOL Add_ThreadInfo(DWORD ThreadId)
{
	THREADINFO *ptr = new THREADINFO;
	if(ptr == NULL) return FALSE;

	ptr->ThreadId = ThreadId;
	ptr->Total_Sockets = 0;
	ptr->next = NULL;
	ThreadInfoList = ptr;

	return TRUE;
}

THREADINFO *GetThreadInfo(DWORD ThreadId)
{
	THREADINFO *ptr = ThreadInfoList;

	while(ptr){
		if(ptr->ThreadId == ThreadId)
			return ptr;
		ptr = ptr->next;
	}

	return NULL;
}

THREADINFO *WhichThreadAvailable()
{
	THREADINFO *ptr = ThreadInfoList;

	while(ptr){
		if(ptr->Total_Sockets < FD_SETSIZE)
			return ptr;
		ptr = ptr->next;
	}

	return NULL;
}

void RemoveThreadInfo(DWORD ThreadId)
{
	THREADINFO *curr = ThreadInfoList;
	THREADINFO *prev = NULL;

	while(curr){
		if(curr->ThreadId == ThreadId){
			if(prev)
				prev->next = curr->next;
			else
				ThreadInfoList = curr->next;
			delete curr;
			return;
		}
		prev = curr;
		curr = curr->next;
	}
}
///////////////////////////////////////////////////////////////////////////////////////

// 소켓 관리 함수
BOOL AddSocketInfo(SOCKET sock, THREADINFO *ThreadInfo);
void RemoveSocketInfo(int nIndex, THREADINFO *ThreadInfo);
// 오류 출력 함수
void err_quit(char *msg);
void err_display(char *msg);

DWORD WINAPI ProcessClient(LPVOID arg)
{
	THREADINFO *ThreadInfo = (THREADINFO *)arg;
	FD_SET rset;
	FD_SET wset;
	SOCKADDR_IN clientaddr;
	int addrlen;
	int retval;

	while(1){
		if(ThreadInfo->Total_Sockets == 0){
			RemoveThreadInfo(ThreadInfo->ThreadId);
			return 0;
		}

		// 소켓 셋 초기화
		FD_ZERO(&rset);
		FD_ZERO(&wset);
		for(int i=0; iTotal_Sockets; i++){
			if(ThreadInfo->SocketInfoArray[i]->recvbytes > 
				ThreadInfo->SocketInfoArray[i]->sendbytes)
				FD_SET(ThreadInfo->SocketInfoArray[i]->sock, &wset);
			else
				FD_SET(ThreadInfo->SocketInfoArray[i]->sock, &rset);
		}

		// select()
		TIMEVAL timeout;
		timeout.tv_sec = 0;
		timeout.tv_usec = 10000; // 0.01초마다 select()가 리턴함
		retval = select(0, &rset, &wset, NULL, &timeout);
		if(retval == SOCKET_ERROR) err_quit("select()");
		else if(retval == 0) continue;

		// 소켓 셋 검사: 데이터 통신
		for(int i=0; iTotal_Sockets; i++){
			SOCKETINFO *ptr = ThreadInfo->SocketInfoArray[i];
			if(FD_ISSET(ptr->sock, &rset)){
				// 데이터 받기
				retval = recv(ptr->sock, ptr->buf, BUFSIZE, 0);
				if(retval == SOCKET_ERROR){
					if(WSAGetLastError() != WSAEWOULDBLOCK){
						err_display("recv()");
						RemoveSocketInfo(i, ThreadInfo);
					}
					continue;
				}
				else if(retval == 0){
					RemoveSocketInfo(i, ThreadInfo);
					continue;
				}
				ptr->recvbytes = retval;
				// 받은 데이터 출력
				addrlen = sizeof(clientaddr);
				getpeername(ptr->sock, (SOCKADDR *)&clientaddr, &addrlen);
				ptr->buf[retval] = '\0';
				printf("[TCP/%s:%d] %s\n", inet_ntoa(clientaddr.sin_addr),
					ntohs(clientaddr.sin_port), ptr->buf);
			}
			if(FD_ISSET(ptr->sock, &wset)){
				// 데이터 보내기
				retval = send(ptr->sock, ptr->buf + ptr->sendbytes, 
					ptr->recvbytes - ptr->sendbytes, 0);
				if(retval == SOCKET_ERROR){
					if(WSAGetLastError() != WSAEWOULDBLOCK){
						err_display("send()");
						RemoveSocketInfo(i, ThreadInfo);
					}
					continue;
				}
				ptr->sendbytes += retval;
				if(ptr->recvbytes == ptr->sendbytes){
					ptr->recvbytes = ptr->sendbytes = 0;
				}
			}
		}
	}

	return 0;
}

int main(int argc, char* argv[])
{
	int retval;

	// 윈속 초기화
	WSADATA wsa;
	if(WSAStartup(MAKEWORD(2,2), &wsa) != 0)
		return -1;

	// socket()
	SOCKET listen_sock = socket(AF_INET, SOCK_STREAM, 0);
	if(listen_sock == INVALID_SOCKET) err_quit("socket()");

	// bind()
	SOCKADDR_IN serveraddr;
	ZeroMemory(&serveraddr, sizeof(serveraddr));
	serveraddr.sin_family = AF_INET;
	serveraddr.sin_port = htons(9000);
	serveraddr.sin_addr.s_addr = htonl(INADDR_ANY);
	retval = bind(listen_sock, (SOCKADDR *)&serveraddr, sizeof(serveraddr));
	if(retval == SOCKET_ERROR) err_quit("bind()");
	
	// listen()
	retval = listen(listen_sock, SOMAXCONN);
	if(retval == SOCKET_ERROR) err_quit("listen()");

	// 데이터 통신에 사용할 변수
	SOCKET client_sock;
	SOCKADDR_IN clientaddr;
	int addrlen;

	while(1){
		addrlen = sizeof(clientaddr);
		client_sock = accept(listen_sock, (SOCKADDR *)&clientaddr, &addrlen);
		if(client_sock == INVALID_SOCKET){
			err_display("accept()");
			continue;
		}
		printf("[TCP 서버] 클라이언트 접속: IP 주소=%s, 포트 번호=%d\n", 
			inet_ntoa(clientaddr.sin_addr), ntohs(clientaddr.sin_port));

		// 넌블로킹 소켓으로 전환
		u_long on = TRUE;
		retval = ioctlsocket(client_sock, FIONBIO, &on);
		if(retval == SOCKET_ERROR) err_display("ioctlsocket()");

		THREADINFO *ThreadInfo = WhichThreadAvailable();
		if(ThreadInfo == NULL){
			HANDLE hThread;
			DWORD ThreadId = -1;
			if(Add_ThreadInfo(ThreadId) == NULL) break;
			ThreadInfo = GetThreadInfo(ThreadId);
			hThread = CreateThread(NULL, 0, ProcessClient, 
				(LPVOID)ThreadInfo, CREATE_SUSPENDED, &ThreadId);
			if(hThread == NULL) break;
			ThreadInfo->ThreadId = ThreadId;

			if(AddSocketInfo(client_sock, ThreadInfo) == FALSE){
				printf("[TCP 서버] 클라이언트 접속을 해제합니다!\n");
				closesocket(client_sock);
			}
			
			ResumeThread(hThread);
			CloseHandle(hThread);
		}
		else{
			if(AddSocketInfo(client_sock, ThreadInfo) == FALSE){
				printf("[TCP 서버] 클라이언트 접속을 해제합니다!\n");
				closesocket(client_sock);
			}
		}
	}

	// 윈속 종료
	WSACleanup();
	return 0;
}

// 소켓 정보 추가
BOOL AddSocketInfo(SOCKET sock, THREADINFO *ThreadInfo)
{
	if(ThreadInfo->Total_Sockets >= FD_SETSIZE){
		printf("[오류] 소켓 정보를 추가할 수 없습니다!\n");
		return FALSE;
	}

	SOCKETINFO *ptr = new SOCKETINFO;
	if(ptr == NULL){
		printf("[오류] 메모리가 부족합니다!\n");
		return FALSE;
	}

	ptr->sock = sock;
	ptr->recvbytes = 0;
	ptr->sendbytes = 0;
	ThreadInfo->SocketInfoArray[ThreadInfo->Total_Sockets++] = ptr;

	return TRUE;
}

// 소켓 정보 삭제
void RemoveSocketInfo(int nIndex, THREADINFO *ThreadInfo)
{
	SOCKETINFO *ptr = ThreadInfo->SocketInfoArray[nIndex];

	// 클라이언트 정보 얻기
	SOCKADDR_IN clientaddr;
	int addrlen = sizeof(clientaddr);
	getpeername(ptr->sock, (SOCKADDR *)&clientaddr, &addrlen);
	printf("[TCP 서버] 클라이언트 종료: IP 주소=%s, 포트 번호=%d\n", 
		inet_ntoa(clientaddr.sin_addr), ntohs(clientaddr.sin_port));

	closesocket(ptr->sock);
	delete ptr;

	for(int i=nIndex; iTotal_Sockets; i++){
		ThreadInfo->SocketInfoArray[i] = ThreadInfo->SocketInfoArray[i+1];
	}
	ThreadInfo->Total_Sockets--;
}

// 소켓 함수 오류 출력 후 종료
void err_quit(char *msg)
{
	LPVOID lpMsgBuf;
	FormatMessage( 
		FORMAT_MESSAGE_ALLOCATE_BUFFER|
		FORMAT_MESSAGE_FROM_SYSTEM,
		NULL, WSAGetLastError(),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPTSTR)&lpMsgBuf, 0, NULL);
	MessageBox(NULL, (LPCTSTR)lpMsgBuf, msg, MB_ICONERROR);
	LocalFree(lpMsgBuf);
	exit(-1);
}

// 소켓 함수 오류 출력
void err_display(char *msg)
{
	LPVOID lpMsgBuf;
	FormatMessage( 
		FORMAT_MESSAGE_ALLOCATE_BUFFER|
		FORMAT_MESSAGE_FROM_SYSTEM,
		NULL, WSAGetLastError(),
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPTSTR)&lpMsgBuf, 0, NULL);
	printf("[%s] %s", msg, (LPCTSTR)lpMsgBuf);
	LocalFree(lpMsgBuf);
}