Windows Socket I/O模型之 WSAAsyncSelect模式

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

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

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

#define WM_SOCKET (WM_USER + 10)

HWND gWorkerWindow=NULL;
CRITICAL_SECTION gSocketCritSec; // Syncrhonize access to socket list


SockObjHeader sockobjhead;

typedef struct _SOCKET_OBJ
{
    SOCKET s; // Socket handle

    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;

    CRITICAL_SECTION SockCritSec;
} SOCKET_OBJ;

SOCKET_OBJ* GetSocketObj(SOCKET s) {
    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->addrlen = sizeof(sockobj->addr);

    InitializeCriticalSection(&sockobj->SockCritSec);
    InitializeCriticalSection(&sockobj->buff.SendRecvQueueCritSec);

    return sockobj;
}

SOCKET_OBJ *FindSocketObj(SOCKET s) {
    EnterCriticalSection(&gSocketCritSec);

    SOCKET_OBJ *obj = NULL;
    SOCKET_OBJ *sobj = NULL;
    SockObj *sptr = (SockObj *)GotoNextDoubleList(&sockobjhead, &(sockobjhead.head));

    while (sptr) {
        obj = (SOCKET_OBJ *)container_of(SOCKET_OBJ, entry, sptr);

        if (obj->s == s) {
            sobj = obj;
            break;
        }

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

    LeaveCriticalSection(&gSocketCritSec);

    return sobj;
}

void RemoveSocketObj(SOCKET_OBJ *sock) {
    EnterCriticalSection(&gSocketCritSec);

    RemoveDoubleList(&sockobjhead, &sock->entry);

    LeaveCriticalSection(&gSocketCritSec);
}

void RemoveSocketObjByHandle(SOCKET s) {
    SOCKET_OBJ *obj = NULL;

    obj = FindSocketObj(s);
    if (obj) {
        RemoveSocketObj(obj);
    }

    return;
}

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

    buffobj = GetBufferObj(gBufferSize);

    ret = 0;

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

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

        printf("Closing\n");
        sockobj->closing = TRUE;

        if (sockobj->buff.head == NULL)
        {
            // If no sends are pending, close the socket for good

            closesocket(sockobj->s);

            ret = -1;
        }
        else
        {
            ret = WSAEWOULDBLOCK;
        }
    } else {
        printf("recv: %d.\n", rc);
        buffobj->buflen = rc;
        EnqueueSingleList(&sockobj->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;
}

LRESULT CALLBACK WindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
    SOCKET_OBJ *sockobj=NULL,
                *newsock=NULL;
    int rc;

    if (uMsg == WM_SOCKET) {
        if (WSAGETSELECTERROR(lParam)) {
            fprintf(stderr, "Socket failed.\n");
            closesocket(wParam);
            RemoveSocketObjByHandle(wParam);
        } else {
            sockobj = FindSocketObj(wParam);
            if (sockobj == NULL)
                return 0;

            switch (WSAGETSELECTEVENT(lParam)) {
                case FD_ACCEPT:
                    newsock = (SOCKET_OBJ*)GetSocketObj(INVALID_SOCKET);

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

                    EnqueueDoubleList(&sockobjhead, &newsock->entry);

                    rc = WSAAsyncSelect(newsock->s, hwnd, WM_SOCKET, FD_READ | FD_WRITE | FD_CLOSE);
                    if (rc == SOCKET_ERROR)
                    {
                        fprintf(stderr, "accept WSAAsyncSelect failed: %d\n", WSAGetLastError());
                        return -1;
                    }

                    break;
                case FD_READ:
                    rc = ReceivePendingData(sockobj);

                    if (rc == -1) {
                        RemoveSocketObj(sockobj);
                        break;
                    } else if (rc != WSAEWOULDBLOCK) {
                        PostMessage(hwnd, WM_SOCKET, wParam, FD_READ);
                    }
                case FD_WRITE:
                    rc = SendPendingData(sockobj);
                    if (rc == -1)
                    {
                        RemoveSocketObj(sockobj);
                    }
                    break;
                case FD_CLOSE:
                    sockobj->closing = TRUE;

                    PostMessage(hwnd, WM_SOCKET, wParam, FD_READ);
                    break;
                default:
                    printf("unknown message.\n");
                    ExitProcess(-1);
            }
        }
    }
    return DefWindowProc(hwnd, uMsg, wParam, lParam);
}

HWND MakeWorkerWindow(void) {
    WNDCLASS wndclass;
    CHAR *ProviderClass = "AsyncSelect";
    HWND windows;

    wndclass.style = CS_HREDRAW | CS_VREDRAW;
    wndclass.lpfnWndProc = (WNDPROC)WindowProc;
    wndclass.cbClsExtra = 0;
    wndclass.cbWndExtra = 0;
    wndclass.hInstance = NULL;
    wndclass.hIcon = LoadIcon(NULL, IDI_APPLICATION);
    wndclass.hCursor = LoadCursor(NULL, IDC_ARROW);
    wndclass.hbrBackground = (HBRUSH)GetStockObject(WHITE_BRUSH);
    wndclass.lpszMenuName = NULL;
    wndclass.lpszClassName = ProviderClass;

    if (RegisterClass(&wndclass) == 0)
    {
        fprintf(stderr, "RegisterClass() failed with error %d\n", GetLastError());
        ExitProcess(-1);
    }

    if ((windows = CreateWindow(
                    ProviderClass,
                    "",
                    WS_OVERLAPPEDWINDOW,
                    CW_USEDEFAULT,
                    CW_USEDEFAULT,
                    CW_USEDEFAULT,
                    CW_USEDEFAULT,
                    NULL,
                    NULL,
                    NULL,
                    NULL)) == NULL)
    {
        fprintf(stderr, "CreateWindow() failed with error %d\n", GetLastError());
        ExitProcess(-1);
    }

    return windows;
}

int _tmain(int argc, _TCHAR* argv[])
{
    WSADATA wsd;
    struct addrinfo *res=NULL,
                    *ptr=NULL;
    SOCKET_OBJ *sockobj=NULL,
                    *sptr=NULL,
                    *tmp=NULL;
    int rc;
    MSG msg;

    InitializeCriticalSection(&gSocketCritSec);

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

    gWorkerWindow = MakeWorkerWindow();

    res = ResolveAddress(gSrvAddr, gPort, gAddressFamily, gSocketType, gProtocol);
    if (res == NULL)
    {
        fprintf(stderr, "ResolveAddress failed to return any addresses!\n");
        return -1;
    }

    InitializeDoubleHead(&sockobjhead);

    ptr = res;
    while (ptr) {
        sockobj = GetSocketObj(INVALID_SOCKET);

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

        EnqueueDoubleListHead(&sockobjhead, &sockobj->entry);

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

        if (gProtocol == IPPROTO_TCP) {
            rc = listen(sockobj->s, 200);
            if (rc == SOCKET_ERROR) {
                fprintf(stderr, "bind failed.\n");
                ExitProcess(-1);
            }

            rc = WSAAsyncSelect(sockobj->s, gWorkerWindow, WM_SOCKET, FD_ACCEPT | FD_CLOSE);
            if (rc == SOCKET_ERROR) {
                fprintf(stderr, "WSAAsyncSelect failed.\n");
                ExitProcess(-1);
            }

        } else {
            ExitProcess(-1);
        }

        ptr = ptr->ai_next;
    }
    freeaddrinfo(res);

    while (rc = GetMessage(&msg, NULL, 0, 0)) {
        if (rc == -1) {
            fprintf(stderr, "GetMessage failed.\n");
            return -1;
        }    

        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    WSACleanup();
    DeleteCriticalSection(&gSocketCritSec);

    return 0;
}