소켓의 입출력 모델(Select - STL사용)

#include <winsock2.h>
#include <stdlib.h>
#include <stdio.h>
#include <list>
#include <algorithm>

#define MAX_BUFFER_SIZE 512

 

// 소켓 정보 저장을 위한 구조체
typedef struct
{
    SOCKET sock;
    char buffer[MAX_BUFFER_SIZE + 1];
    int recvBytes;
    int sendBytes;
    bool bReserveDisconnect;
} SOCKET_INFO;

 

typedef std::list<SOCKET_INFO*> SOCKET_INFO_LIST;

SOCKET_INFO_LIST g_aFreeSocketInfo;
SOCKET_INFO_LIST g_aActiveSocketInfo;

 

void err_quit(char* msg);
void err_display(char* msg);
BOOL AddActiveSocketInfo(SOCKET clientSocket);

 

class IsReserveDisconnect : public std::unary_function<SOCKET_INFO*, bool>
{
public:
    bool operator() (SOCKET_INFO* pSocketInfo)
    {
        if(pSocketInfo->bReserveDisconnect)
        {
            SOCKADDR_IN socketAddr;
            int nAddrLength = sizeof(socketAddr);
            getpeername(pSocketInfo->sock, (SOCKADDR*)&socketAddr, &nAddrLength);
            printf("[TCP 서버] 클라이언트 종료: IP 주소 = %s, 포트번호 = %d\n",
                inet_ntoa(socketAddr.sin_addr), ntohs(socketAddr.sin_port));

            closesocket(pSocketInfo->sock);

            g_aFreeSocketInfo.push_back(pSocketInfo);
        }

        return pSocketInfo->bReserveDisconnect;
    }
};

 

class Functor_FD_SET
{
private:
    FD_SET* m_pRSet;
    FD_SET* m_pWSet;

public:
    Functor_FD_SET(FD_SET* pRSet, FD_SET* pWSet) : m_pRSet(pRSet), m_pWSet(pWSet) {};
    void operator () (SOCKET_INFO* pSocketInfo)
    {
        if(pSocketInfo->recvBytes > pSocketInfo->sendBytes)
            FD_SET(pSocketInfo->sock, m_pWSet);
        else
            FD_SET(pSocketInfo->sock, m_pRSet);
    }
};

 

class Functor_FD_READ_WRITE
{
private:
    FD_SET* m_pRSet;
    FD_SET* m_pWSet;

public:
    Functor_FD_READ_WRITE(FD_SET* pRSet, FD_SET* pWSet) : m_pRSet(pRSet), m_pWSet(pWSet) {};
    void operator() (SOCKET_INFO* pSocketInfo)
    {
        int retValue;
        if(FD_ISSET(pSocketInfo->sock, m_pRSet))
        {
            // 데이터받기
            retValue = recv(pSocketInfo->sock, pSocketInfo->buffer, MAX_BUFFER_SIZE, 0);
            if(retValue == SOCKET_ERROR)
            {
                if(WSAGetLastError() != WSAEWOULDBLOCK)
                {
                    err_display("recv()");
                    pSocketInfo->bReserveDisconnect = true;
                    return;
                }
            }
            else if(retValue == 0)
            {
                pSocketInfo->bReserveDisconnect = true;
                return;
            }
            else
            {
                pSocketInfo->recvBytes = retValue;

                // 받은 데이터 출력
                SOCKADDR_IN sockAddr;
                int nAddrLength = sizeof(sockAddr);
                getpeername(pSocketInfo->sock, (SOCKADDR*)&sockAddr, &nAddrLength);

                pSocketInfo->buffer[retValue] = '\0';
                printf("[TCP/%s:%d] %s\n",
                    inet_ntoa(sockAddr.sin_addr), ntohs(sockAddr.sin_port), pSocketInfo->buffer);
            }
        }

        if(FD_ISSET(pSocketInfo->sock, m_pWSet))
        {
            // 데이터보내기
            retValue = send(pSocketInfo->sock, pSocketInfo->buffer + pSocketInfo->sendBytes,
                pSocketInfo->recvBytes - pSocketInfo->sendBytes, 0);

            if(retValue == SOCKET_ERROR)
            {
                if(WSAGetLastError() != WSAEWOULDBLOCK)
                {
                    err_display("send()");
                    pSocketInfo->bReserveDisconnect = true;
                    return;
                }
            }

            pSocketInfo->sendBytes += retValue;
            if(pSocketInfo->recvBytes == pSocketInfo->sendBytes)
                pSocketInfo->recvBytes = pSocketInfo->sendBytes = 0;
        }
    }
};

 

int _tmain(int argc, _TCHAR* argv[])
{
    int retValue;

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

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

    // bind()
    SOCKADDR_IN serverAddr;
    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(5001);
    serverAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    retValue = bind(listenSocket, (SOCKADDR*)&serverAddr, sizeof(serverAddr));
    if(retValue == SOCKET_ERROR)
        err_quit("bind()");

    //listen
    retValue = listen(listenSocket, SOMAXCONN);
    if(retValue == SOCKET_ERROR)
        err_quit("listen()");

    // 넌블로킹 소켓으로 전환
    unsigned long on = TRUE;
    retValue = ioctlsocket(listenSocket, FIONBIO, &on);
    if(retValue == SOCKET_ERROR)
        err_display("ioctlsocket()");

    for(int i = 0; i < FD_SETSIZE; ++i)
    {
        SOCKET_INFO* pInfo = new SOCKET_INFO;
        memset(pInfo, 0, sizeof(SOCKET_INFO));
        g_aFreeSocketInfo.push_back(pInfo);
    }

    // 통신에 사용할 변수
    FD_SET rset;
    FD_SET wset;
    SOCKET clientSocket;
    SOCKADDR_IN clientAddr;
    int nAddrLength;

    while(1)
    {
        // 소켓 셋 초기화
        FD_ZERO(&rset);
        FD_ZERO(&wset);
        FD_SET(listenSocket, &rset);

        for_each(g_aActiveSocketInfo.begin(), g_aActiveSocketInfo.end(), Functor_FD_SET(&rset, &wset));
        /*
        for(SOCKET_INFO_LIST::iterator iter = g_aActiveSocketInfo.begin(); iter != g_aActiveSocketInfo.end(); ++iter)
        {
        if((*iter)->recvBytes > (*iter)->sendBytes)
        FD_SET((*iter)->sock, wset);
        else
        FD_SET((*iter)->sock, rset);
        }
        */

        // select
        retValue = select(0, &rset, &wset, NULL, NULL);
        if(retValue == SOCKET_ERROR)
            err_quit("select()");

        // 소켓 셋 검사 #1 : 클라이언트 접속
        if(FD_ISSET(listenSocket, &rset))
        {
            nAddrLength = sizeof(clientAddr);
            clientSocket = accept(listenSocket, (SOCKADDR*)&clientAddr, &nAddrLength);
            if(clientSocket == INVALID_SOCKET)
            {
                if(WSAGetLastError() != WSAEWOULDBLOCK)
                    err_display("accept()");
            }
            else
            {
                printf("[TCP 서버] 클라이언트 접속: IP 주소 = %s, 포트번호 = %d\n",
                    inet_ntoa(clientAddr.sin_addr), ntohs(clientAddr.sin_port));

                if(AddActiveSocketInfo(clientSocket) == FALSE)
                {
                    printf("[TCP 서버] 클라이언트 접속을 해제 합니다n");
                    closesocket(clientSocket);
                }
            }
        }

        // 소켓 셋 검사 #2 : 데이터 통신
        for_each(g_aActiveSocketInfo.begin(), g_aActiveSocketInfo.end(), Functor_FD_READ_WRITE(&rset, &wset));

        // 접속 해제 일괄 처리
        g_aActiveSocketInfo.remove_if(IsReserveDisconnect());
    }

    return 0;
}

BOOL AddActiveSocketInfo(SOCKET clientSocket)
{
    if(g_aFreeSocketInfo.empty())
    {
        printf("[오류] 소켓 정보를 추가할 수 없습니다.\n");
        return FALSE;
    }

    SOCKET_INFO* pSocketInfo = (*g_aFreeSocketInfo.begin());
    g_aFreeSocketInfo.pop_front();

    pSocketInfo->sock = clientSocket;
    pSocketInfo->recvBytes = 0;
    pSocketInfo->sendBytes = 0;
    pSocketInfo->bReserveDisconnect = false;

    g_aActiveSocketInfo.push_back(pSocketInfo);

    return TRUE;
}

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\n", msg, (LPCTSTR)lpMsgBuf);
    LocalFree(lpMsgBuf);
}

[출처] 소켓의 입출력 모델(Select - STL사용) (정문수 개인 카페) |작성자 유빈아빠