Windows Socket I/O模型之 Select模式

#include <winsock2.h>
#include <ws2tcpip.h>

#include "resolve.h"
#include "public.h"

typedef SINGLE_LIST_HEADER BuffHeader;
typedef SINGLE_LIST BuffObj;
typedef DOUBLE_LIST_HEADER SockObjHeader;
typedef DOUBLE_LIST SockObj;

typedef struct _SOCKET_OBJ
{
    SOCKET s; // Socket handle

    int listening; // Socket is a listening socket (TCP)

    int closing; // Indicates whether the connection is closing


    SOCKADDR_STORAGE addr; // Used for client's remote address

    int addrlen; // Length of the address


    BuffHeader buff;

    DOUBLE_LIST entry;
} SOCKET_OBJ;

SockObjHeader sockobjhead;

SOCKET_OBJ* GetSocketObj(SOCKET s, int listening) {
    SOCKET_OBJ *sockobj = NULL;

    sockobj = (SOCKET_OBJ*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(SOCKET_OBJ));
    if (sockobj == NULL) {
        fprintf(stderr, "HeapAlloc failed.\n");
        ExitProcess(-1);
    }

    sockobj->s = s;
    sockobj->listening = listening;
    sockobj->addrlen = sizeof(sockobj->addr);

    InitializeCriticalSection(&sockobj->buff.SendRecvQueueCritSec);

    return sockobj;
}

void FreeSocketObj(SOCKET_OBJ *obj) {
    BuffObj *ptr = NULL;
    BUFFER_OBJ *blk = NULL;

    while (true) {
        ptr = DequeueSingleList(&obj->buff);
        if (ptr == NULL)
            break;

        blk = (BUFFER_OBJ *)container_of(BUFFER_OBJ, next, ptr);
        FreeBufferObj(blk);
    }

    HeapFree(GetProcessHeap(), 0, obj);
}

int ReceivePendingData(SOCKET_OBJ *obj) {
    BUFFER_OBJ *buffobj = NULL;
    int rc, ret;

    buffobj = GetBufferObj(gBufferSize);
    ret = 0;

    if (gProtocol == IPPROTO_TCP) {
        rc = recv(obj->s, buffobj->buf, buffobj->buflen, 0);
    } else {
        ExitProcess(-1);
    }

    if (rc == SOCKET_ERROR) {
        ExitProcess(-1);
    } else if (rc == 0) {
        FreeBufferObj(buffobj);
        obj->closing = TRUE;

        if (obj->buff.head == NULL) {
            closesocket(obj->s);
            ret = -1;
        }
    } else {
        buffobj->buflen = rc;
        EnqueueSingleList(&obj->buff, &buffobj->next);
    }
    return ret;
}

int SendPendingData(SOCKET_OBJ *sock) {
    BUFFER_OBJ *bufobj = NULL;
    BuffObj *entry = NULL;
    int nleft = 0,
                idx = 0,
                ret = 0,
                rc = 0;

    while (entry = DequeueSingleList(&sock->buff)) {
        bufobj = (BUFFER_OBJ *)container_of(BUFFER_OBJ, next, entry);

        if (gProtocol == IPPROTO_TCP) {
            nleft = bufobj->buflen;
            idx = 0;

            while (nleft > 0) {
                rc = send(sock->s, &(bufobj->buf[idx]), nleft, 0);
                if (rc == SOCKET_ERROR) {
                    ExitProcess(-1);
                } else {
                    idx += rc;
                    nleft -= rc;
                }
            }

            FreeBufferObj(bufobj);
        } else {
            ExitProcess(-1);
        }
    }

    if (sock->closing == TRUE) {
        closesocket(sock->s);
        ret = -1;

        printf("Closing Connection.\n");
    }

    return ret;
}

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

    SOCKET s;

    struct addrinfo *res = NULL;
    struct addrinfo *ptr = NULL;

    SOCKET_OBJ *sockobj = NULL;

    int rc;

    struct fd_set fdread, fdwrite, fdexcept;
    struct timeval timeout;

    if (WSAStartup(MAKEWORD(2, 2), &wsd) != 0) {
        fprintf(stderr, "load winsock2 failed.\n");
        return 0;
    }

    res = ResolveAddress(gSrvAddr, gPort, gAddressFamily, gSocketType, gProtocol);
    if (res == NULL) {
        fprintf(stderr, "resolve addr failed.\n");
        return -1;
    }

    InitializeDoubleHead(&sockobjhead);

    ptr = res;
    while (ptr) {
        s = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (s == INVALID_SOCKET) {
            fprintf(stderr, "create socket failed.\n");
            break;
        }

        sockobj = GetSocketObj(s, (gProtocol == IPPROTO_TCP) ? TRUE : FALSE);
        EnqueueDoubleList(&sockobjhead, &sockobj->entry);

        rc = bind(sockobj->s, ptr->ai_addr, ptr->ai_addrlen);
        if (rc == SOCKET_ERROR) {
            fprintf(stderr, "bind failed.\n");
            return 0;
        }

        if (gProtocol == IPPROTO_TCP)
        {
            rc = listen(sockobj->s, 200);
            if (rc == SOCKET_ERROR)
            {
                fprintf(stderr, "listen failed: %d\n", WSAGetLastError());
                return -1;
            }
        }

        ptr = ptr->ai_next;
    }

    freeaddrinfo(res);

    while (true) {
        FD_ZERO(&fdread);
        FD_ZERO(&fdwrite);
        FD_ZERO(&fdexcept);

        SockObj *sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, &(sockobjhead.head));
        SOCKET_OBJ *obj = NULL;
        while (sptr) {
            obj = (SOCKET_OBJ *)container_of(SOCKET_OBJ, entry, sptr);

            FD_SET(obj->s, &fdread);
            FD_SET(obj->s, &fdwrite);
            FD_SET(obj->s, &fdexcept);

            sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, sptr);
        }

        timeout.tv_sec = 5;
        timeout.tv_usec = 0;

        rc = select(0, &fdread, &fdwrite, &fdexcept, &timeout);
        if (rc == SOCKET_ERROR) {
            fprintf(stderr, "select failed.\n");
            return -1;
        } else if (rc == 0) {
            printf("Wait Time out.\n");
            continue;
        }

        sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, &(sockobjhead.head));
        SockObj *tmp = NULL;
        obj = NULL;
        while (sptr) {
            obj = (SOCKET_OBJ *)container_of(SOCKET_OBJ, entry, sptr);

            if (FD_ISSET(obj->s, &fdread)) {
                if (obj->listening) {
                    sockobj = (SOCKET_OBJ *)GetSocketObj(INVALID_SOCKET, FALSE);

                    s = accept(obj->s, (SOCKADDR *)&sockobj->addr, &sockobj->addrlen);
                    if (s == INVALID_SOCKET) {
                        fprintf(stderr, "accept failed.\n");
                        return -1;
                    }
                    sockobj->s = s;

                    EnqueueDoubleListHead(&sockobjhead, &sockobj->entry);
                } else {
                    if (ReceivePendingData(obj) != 0) {
                        printf("ReceivePendingData.\n");

                        tmp = sptr;
                        sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, sptr);
                        RemoveDoubleList(&sockobjhead, tmp);
                        FreeSocketObj(obj);

                        continue;
                    }

                    if (SendPendingData(obj) != 0) {
                        tmp = sptr;
                        sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, sptr);
                        RemoveDoubleList(&sockobjhead, tmp);
                        FreeSocketObj(obj);

                        continue;
                    }
                }
            }
            if (FD_ISSET(obj->s, &fdwrite)) {
                if (SendPendingData(obj) != 0) {
                    tmp = sptr;
                    sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, sptr);
                    RemoveDoubleList(&sockobjhead, tmp);
                    FreeSocketObj(obj);

                    continue;
                }
            }
            if (FD_ISSET(obj->s, &fdexcept)) {
                tmp = sptr;
                sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, sptr);
                RemoveDoubleList(&sockobjhead, tmp);
                FreeSocketObj(obj);
            }

            sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, sptr);
        }
    }

    WSACleanup();

    return 0;
}