sockethandle.cpp

利用UDP协议

///////////////////////////////////////////////////////////////////////////////
//  File:       SocketHandle.cpp
//  Version:    1.1
//
//  Author:     Ernest Laurentin
//  E-mail:     [e-mail]
//
//  Implementation of the CSocketHandle and associated classes.
//
//  This code may be used in compiled form in any way you desire. This
//  file may be redistributed unmodified by any means PROVIDING it is
//  not sold for profit without the authors written consent, and
//  providing that this notice and the authors name and all copyright
//  notices remains intact.
//
//  This file is provided "AS IS" with no expressed or implied warranty.
//  The author accepts no liability for any damage/loss of business for the use of
//  class, indirect and consequential damages, even if the Author has been advised
//  of the possibility of such damages.
//
//  Version history
//
//  1.1 - Initial release.
///////////////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include <crtdbg.h>
#include "SocketHandle.h"
#include <stdio.h> 

///////////////////////////////////////////////////////////////////////////////
// SockAddrIn Struct

///////////////////////////////////////////////////////////////////////////////
// Copy
SockAddrIn& SockAddrIn::Copy(const SockAddrIn& sin)
{
    memcpy(this, &sin, Size());
    return *this;
}

///////////////////////////////////////////////////////////////////////////////
// IsEqual
bool SockAddrIn::IsEqual(const SockAddrIn& sin) const
{
    // Is it Equal? - ignore 'sin_zero'
    return (memcmp(this, &sin, Size()-sizeof(sin_zero)) == 0);
}

///////////////////////////////////////////////////////////////////////////////
// CreateFrom
bool SockAddrIn::CreateFrom(LPCTSTR sAddr, LPCTSTR sService, int nFamily /*=AF_INET*/)
{
    Clear();
    sin_addr.s_addr = htonl( CSocketHandle::GetIPAddress(sAddr) );
    sin_port = htons( CSocketHandle::GetPortNumber( sService ) );
    sin_family = static_cast<short>(nFamily);

    return !IsNull();
}


///////////////////////////////////////////////////////////////////////////////
// CSocketHandle
CSocketHandle::CSocketHandle()
: m_hSocket(INVALID_SOCKET)
{
}

CSocketHandle::~CSocketHandle()
{
    Close();
}

///////////////////////////////////////////////////////////////////////////////
// GetSockName
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              Get current address
// PARAMETERS:
//  SockAddrIn& saddr_in: address structure when connected
///////////////////////////////////////////////////////////////////////////////
bool CSocketHandle::GetSockName(SockAddrIn& saddr_in) const
{
    if (IsOpen())
    {
        int namelen = (int)saddr_in.Size();
        return (SOCKET_ERROR != getsockname(GetSocket(), (LPSOCKADDR)saddr_in, &namelen));
    }
    return false;
}

///////////////////////////////////////////////////////////////////////////////
// GetPeerName
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              Get peer address (connected socket)
// PARAMETERS:
//  SockAddrIn& saddr_in: address structure when connected
///////////////////////////////////////////////////////////////////////////////
bool CSocketHandle::GetPeerName(SockAddrIn& saddr_in) const
{
    if (IsOpen())
    {
        int namelen = (int)saddr_in.Size();
        return (SOCKET_ERROR != getpeername(GetSocket(), (LPSOCKADDR)saddr_in, &namelen));  
    }
    return false;
}

///////////////////////////////////////////////////////////////////////////////
// Close
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              Close a socket connection
// PARAMETERS:
//  None
///////////////////////////////////////////////////////////////////////////////
void CSocketHandle::Close()
{
    if ( IsOpen() )
    {
        ShutdownConnection(m_hSocket);
        ::InterlockedExchange((LONG*)&m_hSocket, INVALID_SOCKET);
    }
}


///////////////////////////////////////////////////////////////////////////////
// AddMembership
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              Add membership to a multicast address
// PARAMETERS:
//  LPCTSTR pszIPAddr: IP address for membership
//  LPCTSTR pszNIC: interface IP
///////////////////////////////////////////////////////////////////////////////
bool CSocketHandle::AddMembership(LPCTSTR pszIPAddr, LPCTSTR pszNIC)
{
    if ( IsOpen() )
    {
        int nType = 0;
        int nOptLen = sizeof(int);
        if ( SOCKET_ERROR != getsockopt(m_hSocket, SOL_SOCKET, SO_TYPE, (char*)&nType, &nOptLen))
        {
            if ( nType == SOCK_DGRAM )
            {
                // Setup interface for multicast traffic
                ULONG ulNIC = htonl( CSocketHandle::GetIPAddress( pszNIC ) );
                if ( SOCKET_ERROR != setsockopt(m_hSocket, IPPROTO_IP, IP_MULTICAST_IF,(char *) &ulNIC, sizeof(ulNIC)))
                {
                    int nTTL = SOCKHANDLE_TTL;
                    if ( SOCKET_ERROR != setsockopt(m_hSocket, IPPROTO_IP, IP_MULTICAST_TTL, (const char*)&nTTL, sizeof(nTTL)))
                    {
                        ip_mreq mreq;
                        mreq.imr_multiaddr.s_addr = htonl( CSocketHandle::GetIPAddress( pszIPAddr ) );
                        mreq.imr_interface.s_addr = ulNIC;
                        return ( SOCKET_ERROR != setsockopt(m_hSocket, IPPROTO_IP, IP_ADD_MEMBERSHIP, (const char*)&mreq, sizeof(mreq)));
                    }
                }
            }
        }
        SetLastError( WSAGetLastError() );
    }
    return false;
}


///////////////////////////////////////////////////////////////////////////////
// DropMembership
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              Remove membership from a multicast address
// PARAMETERS:
//  LPCTSTR pszIPAddr: IP address for membership
//  LPCTSTR pszNIC: interface IP
///////////////////////////////////////////////////////////////////////////////
bool CSocketHandle::DropMembership(LPCTSTR pszIPAddr, LPCTSTR pszNIC)
{
    if ( IsOpen() )
    {
        int nType = 0;
        int nOptLen = sizeof(int);
        if ( SOCKET_ERROR != getsockopt(m_hSocket, SOL_SOCKET, SO_TYPE, (char*)&nType, &nOptLen))
        {
            if ( nType == SOCK_DGRAM )
            {
                ip_mreq mreq;
                mreq.imr_multiaddr.s_addr = htonl( CSocketHandle::GetIPAddress( pszIPAddr ) );
                mreq.imr_interface.s_addr = htonl( CSocketHandle::GetIPAddress( pszNIC ) );
                return ( SOCKET_ERROR != setsockopt(m_hSocket, IPPROTO_IP, IP_DROP_MEMBERSHIP, (const char*)&mreq, sizeof(mreq)));
            }
        }
        SetLastError( WSAGetLastError() );
    }
    return false;
}


///////////////////////////////////////////////////////////////////////////////
// CreateSocket
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              This function creates a new socket for connection (SOCK_STREAM)
//              or an connectionless socket (SOCK_DGRAM).  A connectionless
//              socket should not call "accept()" since it cannot receive new
//              connection.  This is used as SERVER socket
// PARAMETERS:
//  LPCTSTR pszHostName: hostname (or NULL to select default adapter)
//  LPCTSTR pszServiceName: Service name or port number
//  int nFamily: address family to use (set to AF_INET)
//  int nType: type of socket to create (SOCK_STREAM, SOCK_DGRAM)
//  UINT uOptions: other options to use
///////////////////////////////////////////////////////////////////////////////
bool CSocketHandle::CreateSocket(LPCTSTR pszHostName, LPCTSTR pszServiceName,
                                int nFamily, int nType, UINT uOptions /* = 0 */)
{
    // Socket is already opened
    if ( IsOpen() )
        return false;

    // Create a Socket that is bound to a specific service provider
    // nFamily: (AF_INET)
    // nType: (SOCK_STREAM, SOCK_DGRAM)
#ifdef SOCKHANDLE_USE_OVERLAPPED
    SOCKET sock = WSASocket(nFamily, nType, IPPROTO_IP, NULL, 0, WSA_FLAG_OVERLAPPED);
#else
    SOCKET sock = socket(nFamily, nType, IPPROTO_IP);
#endif
    if (INVALID_SOCKET != sock)
    {
        if (uOptions & SO_REUSEADDR)
        {
            // Inform Windows Sockets provider that a bind on a socket should not be disallowed
            // because the desired address is already in use by another socket
            BOOL optval = TRUE;
            if ( SOCKET_ERROR == setsockopt( sock, SOL_SOCKET, SO_REUSEADDR, (char *) &optval, sizeof( BOOL ) ) )
            {
                SetLastError( WSAGetLastError() );
                closesocket( sock );
                return false;
            }
        }

        if (nType == SOCK_DGRAM)
        {
            if (uOptions & SO_BROADCAST)
            {
                // Inform Windows Sockets provider that broadcast messages are allowed
                BOOL optval = TRUE;
                if ( SOCKET_ERROR == setsockopt( sock, SOL_SOCKET, SO_BROADCAST, (char *) &optval, sizeof( BOOL ) ) )
                {
                    SetLastError( WSAGetLastError() );
                    closesocket( sock );
                    return false;
                }
            }
        }

        // Associate a local address with the socket
        SockAddrIn sockAddr;
        sockAddr.CreateFrom(pszHostName, pszServiceName, nFamily);

        if ( SOCKET_ERROR == bind(sock, (LPSOCKADDR)sockAddr, (int)sockAddr.Size()))
        {
            SetLastError( WSAGetLastError() );
            closesocket( sock );
            return false;
        }

        // Listen to the socket, only valid for connection socket (TCP)
        if (SOCK_STREAM == nType)
        {
            if ( SOCKET_ERROR == listen(sock, SOMAXCONN))
            {
                SetLastError( WSAGetLastError() );
                closesocket( sock );
                return false;
            }
        }

        // Success, now we may save this socket
        m_hSocket = sock;
    }

    return (INVALID_SOCKET != sock);
}


///////////////////////////////////////////////////////////////////////////////
// ConnectTo
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//              Establish connection with a server service or port
// PARAMETERS:
//  LPCTSTR pszHostName: hostname or address to connect (in .dot format)
//  LPCTSTR pszServiceName: Service name or port number
//  int nFamily: address family to use (set to AF_INET)
//  int nType: type of socket to create (SOCK_STREAM, SOCK_DGRAM)
///////////////////////////////////////////////////////////////////////////////
bool CSocketHandle::ConnectTo(LPCTSTR pszHostName, LPCTSTR pszServiceName,
                             int nFamily, int nType)
{
    // Socket is already opened
    if ( IsOpen() )
        return false;

    // Create a Socket that is bound to a specific service provider
    // nFamily: (AF_INET)
    // nType: (SOCK_STREAM, SOCK_DGRAM)
#ifdef SOCKHANDLE_USE_OVERLAPPED
    SOCKET sock = WSASocket(nFamily, nType, IPPROTO_IP, NULL, 0, WSA_FLAG_OVERLAPPED);
#else
    SOCKET sock = socket(nFamily, nType, IPPROTO_IP);
#endif
    if (INVALID_SOCKET != sock)
    {
        // Associate a local address with the socket
        LPCTSTR pszHost = NULL;
        if ( pszHostName && _tcscmp(pszHostName, TEXT("localhost")) == 0 )
            pszHost = TEXT("localhost");
        SockAddrIn sockAddr;
        if (false == sockAddr.CreateFrom(pszHost, TEXT("0"), nFamily))
        {
            SetLastError( WSAGetLastError() );
            closesocket( sock );
            return false;
        }

        if ( SOCKET_ERROR == bind(sock, (LPSOCKADDR)sockAddr, sizeof(SOCKADDR_IN)))
        {
            SetLastError( WSAGetLastError() );
            closesocket( sock );
            return false;
        }

        // Now get destination address & port
        sockAddr.CreateFrom( pszHostName, pszServiceName );

        // try to connect - if fail, server not ready
        if (SOCKET_ERROR == connect( sock, (LPSOCKADDR)sockAddr, sizeof(SOCKADDR_IN)))
        {
            SetLastError( WSAGetLastError() );
            closesocket( sock );
            return false;
        }

        // Success, now we may save this socket
        m_hSocket = sock;
    }
    return (INVALID_SOCKET != sock);
}


///////////////////////////////////////////////////////////////////////////////
// Read
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//      Reads the Socket Communication
// PARAMETERS:
//      LPBYTE lpBuffer: buffer to place new data
//      DWORD dwSize: maximum size of buffer
//      DWORD dwTimeout: timeout to use in millisecond
///////////////////////////////////////////////////////////////////////////////
DWORD CSocketHandle::Read(LPBYTE lpBuffer, DWORD dwSize, LPSOCKADDR lpAddrIn,
                         DWORD dwTimeout)
{
    _ASSERTE( IsOpen() );
    _ASSERTE( lpBuffer != NULL );

    if (!IsOpen() || lpBuffer == NULL || dwSize < 1L)
        return (DWORD)-1L;

    fd_set  fdRead  = { 0 };
    TIMEVAL stTime;
    TIMEVAL *pstTime = NULL;

    if ( INFINITE != dwTimeout ) {
        stTime.tv_sec = dwTimeout/1000;
        stTime.tv_usec = (dwTimeout%1000)*1000;
        pstTime = &stTime;
    }

    SOCKET s = GetSocket();

    // Set Descriptor
    FD_SET( s, &fdRead );

    // Select function set read timeout
    DWORD dwBytesRead = 0L;
    int res = select((int)s, &fdRead, NULL, NULL, pstTime );
    if ( res > 0)
    {
        if (lpAddrIn)
        {
            // UDP
            int fromlen = sizeof(SOCKADDR_IN);
            res = recvfrom( s,(LPSTR) lpBuffer, dwSize, 0, lpAddrIn, &fromlen);
        }
        else
        {
            // TCP
            res = recv( s,(LPSTR) lpBuffer, dwSize, 0);
        }
    }
    if ( res == SOCKET_ERROR )
    {
        SetLastError( WSAGetLastError() );
    }
    dwBytesRead = (DWORD)((res >= 0)?(res) : (-1));

    return dwBytesRead;
}

///////////////////////////////////////////////////////////////////////////////
// ReadEx
///////////////////////////////////////////////////////////////////////////////
// DESCRIPTION:
//      Asynchronous Read from the Socket Communication
// PARAMETERS:
//      LPBYTE lpBuffer: buffer to place new data
//      DWORD dwSize: maximum size of buffer
//      LPWSAOVERLAPPED lpOverlapped: Overlapped structure