iocpclient.cpp

这个是网络编程

//
// Sample: I/O Completion Port IPv4/IPv6 Client
//
// Files:
//      iocpclient.cpp    - this file
//      resolve.cpp       - Common name resolution routines
//      resolve.h         - Header file for name resolution routines
//
// Description:
//      This is a sample client which uses IO completion ports and
//      uses the ConnectEx and TransmitFile APIs. As such this client
//      only works on Windows XP+. This client will attempt the specified
//      number of connections for each server address resovled. That is
//      if the server's address is explicitly specified (via the -c NUM 
//      option) then the client will establish NUM connections. If the
//      server's hostname is specified and it resolves to 5 different
//      addresses then NUM connections are attemped for each address.
//
//      Again, this client is an echo client. It sends a specified amount
//      of data on the connection and then waits to receive the data
//      back. Normally, the client uses WSASend for data, but if the
//      -t SIZE flag is given then TransmitFile is used with a temporary
//      file of SIZE bytes. This temp file is created on the fly and
//      deleted upon client exit.
//
//      The -r command is used for basic rate limiting of data transfers
//      from the client to the server. For reasonable limits the mechanism
//      works; however, the higher the limit is the less reliable this method
//      becomes (as the sleep time approaches zero - then the time it takes
//      to post the sends becomes an issue).  This option is provided to
//      allow a method of simulating a steady connection rate (as the 
//      connection requests are shaped as well) in addition to smoothing the
//      data transmission to the server to prevent sudden spikes and/or
//      saturating the network bandwidth.
//
//      NOTE: This client only supports the TCP protocol.
// 
// Compile:
//      cl -o iocpclient.exe iocpclient.cpp resolve.cpp ws2_32.lib
//
// Usage:
//      iocpclient.exe [options]
//          -a 4|6     Address family, 4 = IPv4, 6 = IPv6 [default = IPv4]
//          -b size    Buffer size for send/recv (in bytes)
//          -c count   Number of connections to establish
//          -e port    Port number
//          -n server  Server address or name to connect to
//          -l addr    Local address to bind to [default INADDR_ANY for IPv4 or INADDR6_ANY for IPv6]
//          -r rate    Rate at which to send data
//          -t size    Use TransmitFile instead of sends (size of file to send)
//          -x count   Number of sends
//

#include <winsock2.h>
#include <ws2tcpip.h>
#include <mswsock.h>
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>

#include "resolve.h"

#define DEFAULT_BUFFER_SIZE         4096   // default buffer size
#define DEFAULT_OVERLAPPED_COUNT    1      // Number of overlapped recv per socket
#define DEFAULT_CLIENT_CONNECTIONS  10     // Number of connections to initiate
#define DEFAULT_FILE_SIZE           2000000// Default size of file for TransmitFile
#define DEFAULT_SEND_COUNT          100    // How many send/TransmitFiles to perform

int gAddressFamily = AF_UNSPEC,         // default to unspecified
    gSocketType    = SOCK_STREAM,       // default to TCP socket type
    gProtocol      = IPPROTO_TCP,       // default to TCP protocol
    gBufferSize    = DEFAULT_BUFFER_SIZE,
    gOverlappedCount = DEFAULT_OVERLAPPED_COUNT,
    gConnectionCount = DEFAULT_CLIENT_CONNECTIONS,
    gFileSize        = DEFAULT_FILE_SIZE,
    gSendCount       = DEFAULT_SEND_COUNT,
    gRateLimit       = -1,
    gTimeout         = 0;

USHORT gLocalPort = 0x0000FFFD;

BOOL gTransmitFile = FALSE;             // Use TransmitFile instead

HANDLE gTempFile   = INVALID_HANDLE_VALUE;

char *gBindAddr    = NULL,              // local interface to bind to
     *gServerAddr  = NULL,              // Server address to connect to
     *gBindPort    = "5150";            // local port to bind to

//
// This is our per I/O buffer. It contains a WSAOVERLAPPED structure as well
//    as other necessary information for handling an IO operation on a socket.
//
typedef struct _BUFFER_OBJ
{
    WSAOVERLAPPED        ol;

    HANDLE               hFile;         // Open file handle for TransmitFile

    char                *buf;           // Buffer for recv/send/AcceptEx
    int                  buflen;        // Length of the buffer

    int                  operation;     // Type of operation issued
#define OP_CONNECT      0                   // ConnectEx
#define OP_READ         1                   // WSARecv/WSARecvFrom
#define OP_WRITE        2                   // WSASend/WSASendTo
#define OP_TRANSMIT     3                   // TransmitFile

    SOCKADDR_STORAGE     addr;
    int                  addrlen;

    struct _BUFFER_OBJ  *next;

} BUFFER_OBJ;

//
// This is our per socket buffer. It contains information about the socket handle
//    which is returned from each GetQueuedCompletionStatus call.
//
typedef struct _SOCKET_OBJ
{
    SOCKET               s;              // Socket handle
    int                  af;             // Address family of socket (AF_INET, AF_INET6)

    volatile LONG        OutstandingOps; // Number of outstanding overlapped ops on 
                                         //    socket
    volatile LONG        SendCount;      // Number of sends to perform on connection

    BOOL                 bConnected;
    BOOL                 bClosing;       // Is socket closing

    // Pointers to Microsoft specific extensions.
    LPFN_CONNECTEX       lpfnConnectEx;
    LPFN_TRANSMITFILE    lpfnTransmitFile;

    BUFFER_OBJ          *Repost;         // Send buffer to repost (used with rate limit)

    CRITICAL_SECTION     SockCritSec;

    struct _SOCKET_OBJ  *next,
                        *prev;
} SOCKET_OBJ;

// List of open sockets connected to server
SOCKET_OBJ *gConnectionList=NULL;

//
// Statistics counters
//
volatile LONG gBytesRead=0,
              gBytesSent=0,
              gStartTime=0,
              gBytesReadLast=0,
              gBytesSentLast=0,
              gStartTimeLast=0,
              gTotalConnections=0,
              gCurrentConnections=0,
              gConnectionRefused=0;

//
// Function: usage
//
// Description:
//      Prints usage information and exits the process.
//
void usage(char *progname)
{
    fprintf(stderr, "usage: %s [-a 4|6] [-e port] [-l local-addr] [-p udp|tcp]\n",
            progname);
    fprintf(stderr, "  -a 4|6     Address family, 4 = IPv4, 6 = IPv6 [default = IPv4]\n"
                    "  -b size    Buffer size for send/recv [default = %d]\n"
                    "  -c count   Number of connections to establish\n"
                    "  -e port    Port number [default = %s]\n"
                    "  -n server  Server address or name to connect to\n"
                    "  -p port    Local port number to bind to\n"
                    "  -l addr    Local address to bind to [default INADDR_ANY for IPv4 or INADDR6_ANY for IPv6]\n"
                    "  -r rate    Use the QOS provider to limit send rate\n"
                    "  -t size    Use TransmitFile instead of sends (size of file to send)\n"
                    "  -x count   Number of sends\n",
                    gBufferSize,
                    gBindPort
                    );
    ExitProcess(-1);
}

//
// Function: GetBufferObj
// 
// Description:
//    Allocate a BUFFER_OBJ. Each receive posted allocates one of these. 
//    After the recv is successful, the BUFFER_OBJ is queued for
//    sending by the send thread. To increase performance, a lookaside lists 
//    should be used to cache free BUFFER_OBJ.
//
BUFFER_OBJ *GetBufferObj(int buflen)
{
    BUFFER_OBJ *newobj=NULL;

    // Allocate the object
    newobj = (BUFFER_OBJ *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BUFFER_OBJ));
    if (newobj == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    // Allocate the buffer
    newobj->buf = (char *)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(BYTE) *buflen);
    if (newobj->buf == NULL)
    {
        fprintf(stderr, "GetBufferObj: HeapAlloc failed: %d\n", GetLastError());
        ExitProcess(-1);
    }
    newobj->buflen = buflen;

    newobj->addrlen = sizeof(newobj->addr);

    return newobj;
}

//
// Function: FreeBufferObj
// 
// Description:
//    Free the buffer object. To increase performance, a lookaside list should be
//    implemented to cache BUFFER_OBJ when freed.
//
void FreeBufferObj(BUFFER_OBJ *obj)
{
    HeapFree(GetProcessHeap(), 0, obj->buf);
    HeapFree(GetProcessHeap(), 0, obj);
}

//
// Function: GetSocketObj
//
// Description:
//    Allocate a socket object and initialize its members. A socket object is
//    allocated for each socket created (either by socket or accept).
//    Again, a lookaside list can be implemented to cache freed SOCKET_OBJ to
//    improve performance.
//
SOCKET_OBJ *GetSocketObj(SOCKET s, int af)
{
    SOCKET_OBJ  *sockobj=NULL;

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

    InitializeCriticalSection(&sockobj->SockCritSec);

    // Initialize the members
    sockobj->s = s;
    sockobj->af = af;

    return sockobj;
}

//
// Function: FreeSocketObj
//
// Description:
//    Frees a socket object. If there are outstanding operations, the object
//    is not freed. 
//
void FreeSocketObj(SOCKET_OBJ *obj)
{
    BUFFER_OBJ  *ptr=NULL,
                *tmp=NULL;

    if (obj->OutstandingOps != 0)
    {
        // Still outstanding operations so just return
        return;
    }
    // Close the socket if it hasn't already been closed
    if (obj->s != INVALID_SOCKET)
    {
        closesocket(obj->s);
        obj->s = INVALID_SOCKET;
    }

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

//
// Function: InsertSocketObj
//
// Description:
//    Insert a SOCKET_OBJ into a list of socket objects. Insertions
//    are performed at the end of the list.
//
void InsertSocketObj(SOCKET_OBJ **head, SOCKET_OBJ *obj)
{
    SOCKET_OBJ *end=NULL,
               *ptr=NULL;

    // Find the end of the list
    ptr = *head;
    if (ptr)
    {
        while (ptr->next)
            ptr = ptr->next;
        end = ptr;
    }

    obj->next = NULL;
    obj->prev = end;

    if (end == NULL)
    {
        // List is empty
        *head = obj;
    }
    else
    {
        // Put new object at the end
        end->next = obj;
        obj->prev = end;
    }
}

//
// Function: RemoveSocketObj
//
// Description:
//    Removes the specified SOCKET_OBJ from the list of socket objects.
//
SOCKET_OBJ *RemoveSocketObj(SOCKET_OBJ **head, SOCKET_OBJ *buf)
{
    // Make sure list isn't empty
    if (*head != NULL)
    {
        // Fix up the next and prev pointers
        if (buf->prev)
            buf->prev->next = buf->next;
        if (buf->next)
            buf->next->prev = buf->prev;

        if (*head == buf)
            (*head) = buf->next;
    }
    return buf;
}

//
// Function: CreateTempFile
//
// Description:
//    Creates a temporary file which is used by all TransmitFile
//    operations. This file is created as temporary and will be
//    deleted upon handle closure.
//
HANDLE CreateTempFile(char *filename, DWORD size)
{
    OVERLAPPED ol;
    HANDLE     hFile;
    DWORD      bytes2write,
               offset,
               nLeft,
               written,
               buflen=1024;
    char       buf[1024];
    int        rc;

    // Create the file as temporary.
    hFile = CreateFile(
            filename,
            GENERIC_READ | GENERIC_WRITE,
            FILE_SHARE_READ,
            NULL,
            CREATE_ALWAYS,
            FILE_ATTRIBUTE_TEMPORARY | FILE_FLAG_OVERLAPPED | 
                FILE_FLAG_SEQUENTIAL_SCAN | FILE_FLAG_DELETE_ON_CLOSE,
            NULL
            );
    if (hFile == INVALID_HANDLE_VALUE)
    {
        fprintf(stderr, "CreateTempFile failed: %d\n", GetLastError());
        return hFile;
    }

    memset(buf, '$', buflen);

    memset(&ol, 0, sizeof(ol));
    offset = 0;

    ol.hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if (ol.hEvent == FALSE)
    {
        fprintf(stderr, "CreateTempFile: CreateEvent failed: %d\n",
                GetLastError());
        return INVALID_HANDLE_VALUE;
    }

    // Write the specified amount of data to the file
    nLeft = size;
    while (nLeft > 0)
    {
        bytes2write = ((nLeft < buflen) ? nLeft : buflen);
        
        ol.Offset = offset;

        rc = WriteFile(
                hFile,
                buf,
                bytes2write,
               &written,
               &ol
                );
        if (rc == 0)
        {
            if (GetLastError() != ERROR_IO_PENDING)
            {
                fprintf(stderr, "CreateTempFile: WriteFile failed: %d\n",
                        GetLastError());
                return INVALID_HANDLE_VALUE;
            }
            else
            {
                rc = GetOverlappedResult(
                        hFile,
                       &ol,
                       &written,
                        TRUE
                        );
                if (rc == 0)
                {
                    fprintf(stderr, "CreateTempFile: GetOverlappedResult failed: %d\n",
                            GetLastError());
                    return INVALID_HANDLE_VALUE;
                }
            }
        }
        ResetEvent(ol.hEvent);

        // Ajust the offset and bytes left to write
        offset += written;
        nLeft  -= written;
    }

    printf("Created temp file of size: %d\n", offset);

    CloseHandle(ol.hEvent);

    return hFile;
}

//
// Function: ValidateArgs
//
// Description:
//      Parses the command line arguments and sets up some global 
//      variables.
//
void ValidateArgs(int argc, char **argv)
{
    int     i;

    for(i=1; i < argc ;i++)
    {
        if (((argv[i][0] != '/') && (argv[i][0] != '-')) || (strlen(argv[i]) < 2))
            usage(argv[0]);
        else
        {
            switch (tolower(argv[i][1]))
            {
                case 'a':               // address family - IPv4 or IPv6
                    if (i+1 >= argc)
                        usage(argv[0]);
                    if (argv[i+1][0] == '4')
                        gAddressFamily = AF_INET;
                    else if (argv[i+1][0] == '6')
                        gAddressFamily = AF_INET6;
                    else
                        usage(argv[0]);
                    i++;
                    break;
                case 'b':               // buffer size for send/recv
                    if (i+1 >= argc)
                        usage(argv[0]);
                    gBufferSize = atol(argv[++i]);