//  Copyright © 2001 - 2002, RESON Inc. All Rights Reserved.
//
//  No part of this file may be reproduced or transmitted in any form or by
//  any means, electronic or mechanical, including photocopy, recording, or
//  information storage or retrieval system, without permission in writing
//  from RESON Inc.
//
//  Filename:   ThreadedSocket.h
//
//  Project:    6046.
//
//  Author(s):  W. Arcus
//
//  Purpose:    Class definition file for CThreadedSocket and helper templates.
//
//  Notes:      1)  When a complete network message is recieved (from possible fragmants) the installed
//                  callback routine will be invoked. At this time, the client should then handle this 
//                  data as quickly as possible (typically copy it into a working circular buffer etc).
//
//              2)  The callback occurs within the context of a thread running inside this DLL 
//                  and, as such, the client should use a critical section or similar to synchonize 
//                  access of the record data with its internal thread(s). It should also exit the callback
//                  as quickly as possible to avoid data loss.
//

#if !defined(AFX_THREADEDSOCKET_H__A9048515_BC0D_4E43_BA56_E3F8672B6A33__INCLUDED_)
#define AFX_THREADEDSOCKET_H__A9048515_BC0D_4E43_BA56_E3F8672B6A33__INCLUDED_

#if _MSC_VER > 1000
#pragma once
#endif // _MSC_VER > 1000

#include "DynamicBuffer.h"
#include "BaseThread.h"
#include "SystemTime.h"

#include <afxsock.h>

//////////////////////////////////////////////////////////////////////
// Socket related error codes.

typedef enum tagESOCKETERRORCODES
{
    socketErrorNone,
    socketErrorUnknownException,
    socketErrorBadPointer,
    socketErrorSocketAccept,
    socketErrorSocketMode,
    socketErrorRead,
    socketErrorCallbackError,
    socketErrorThreadError,
    socketErrorInvalidObject,
    socketErrorUnspecified,
    socketErrorTooManyConnections
}
ESOCKETERRORCODES;

typedef enum tagECALLBACKREASONS
{
    callbackConnect,
    callbackDisconnect,
    callbackData
}
ECALLBACKREASONS;

//////////////////////////////////////////////////////////////////////
// C compliant callback function definition when recieved data record is available.

extern "C"
{
typedef void ( * PFN_RECORD_READY )      (  BYTE           *pbyData,                // Data pointer where relevant - can be NULL.
                                            PVOID           lpvParam,               // Optional parameter.
                                            ULONG           ulNumBytes,             // Number of data bytes received - can be zero.
                                            INT             iSocketIndex,           // Socket index (relevant to server end).
                                            ULONG           ulTimeStamp,            // Time stamp.
                                            INT             iResonForCallback );    // Reason for the callback.
}

//////////////////////////////////////////////////////////////////////
// CThreadedSocket class definition.

template <class tProtocol>
class CThreadedSocket : public CBaseThread
{
public:
                            CThreadedSocket     (   PFN_RECORD_READY    pfnCallback,
                                                    void *              pvParam,
                                                    const char *        pszAddress,
                                                    unsigned int        uiPort,
                                                    unsigned long       ulMaxMessageLength,
                                                    int                 iSocketType );

    virtual                ~CThreadedSocket     (   void );

    bool                    IsInitialized       (   void ) const;

    bool                    Send                (   BYTE               *pbyStream,
                                                    unsigned long       ulNumBytes );

    CAsyncSocket &          GetSocket           (   void );
    
    bool                    SetServerSocketOptions( const bool         &rbNoDelay     = false,
                                                    const int          &riSocketIndex = -1 );

    bool                    Connect             (   bool                bNoDelay = false );

    unsigned int            Port                (   void ) const;

protected:

    // Attributes.

    bool                    m_bInitialized;

    BYTE                   *m_pbySocketData;
                        
    int                     m_iErrorCode;                                                   // Thread error code cf. eSocketErrorCodes above.
    int                     m_iSocketType;
    int                     m_iSocketIndex;

    unsigned int            m_uiPort;
    
    unsigned long           m_ulMessageNumber;
    unsigned long           m_ulRxBufferSize;

    PFN_RECORD_READY        m_pfnCallback;

    CString                 m_Address;
    CAsyncSocket            m_Socket;
    CDynamicBuffer<BYTE>    m_ReadFifo;

    CCritical               m_WriteCriticalSection;
    tProtocol               m_Protocol;

    // Methods.

    bool                    SendData            (   BYTE           *pbyStream, 
                                                    unsigned long   ulNumBytes );

    bool                    CallCallback        (   BYTE            *pbyData,               // Data pointer where relevant - can be NULL.
                                                    unsigned long   ulNumBytes,             // Number of data bytes received - can be zero.
                                                    int             iSocketIndex,           // Socket index (relevant to server end only - will be zero otherwise).
                                                    unsigned long   ulTimeStamp,
                                                    int             iReasonForCallback );

    unsigned int            HandleInboundData   (   void );

    int                     ReadSocketBytes     (   const int      &riBytesToRead );

    void                    WatchCycle          (   void );

private:

    // Attributes

    void                   *m_pvParam;                                                      // Optional callaback parameter.

    // Methods.

    // Make the following constructors and assignment operator private so that we get a compile error
    // if the compiler uses them behind our back.

                        CThreadedSocket   (   void );                                   // Not implemented.
                        CThreadedSocket   (   const   CThreadedSocket &rCopy );         // Not implemented.

    CThreadedSocket &     operator =      (   const   CThreadedSocket &rCopy );         // Not implemented.

};

//////////////////////////////////////////////////////////////////////
// Internal macros, definitions etc.

#define AbortIfThreadTerminateRequest_m()   if ( ! IsActive() ) \
                                            {                   \
                                                break;          \
                                            }                   \

#define ThrowSocketError_m()                if ( ( iSocketError = m_Socket.GetLastError() ) != WSAEWOULDBLOCK ) \
                                            {                                                                   \
                                                throw iSocketError;                                             \
                                            }                                                                   \

//////////////////////////////////////////////////////////////////////
// CThreadedSocket class implementation.

template <class tProtocol>
CThreadedSocket<tProtocol>::CThreadedSocket(    PFN_RECORD_READY    pfnCallback,
                                                void *              pvParam,
                                                const char *        pszAddress,
                                                unsigned int        uiPort,
                                                unsigned long       ulMaxMessageLength,
                                                int                 iSocketType )

                           :CBaseThread()
{
    m_bInitialized      = false;

    m_Address           = pszAddress;
    m_uiPort            = uiPort;

    m_ulRxBufferSize    = ulMaxMessageLength;

    m_pfnCallback       = pfnCallback;
    m_pvParam           = pvParam;

    // Must be a stream or datagram socket.

    ASSERT( ( iSocketType == SOCK_DGRAM ) || ( iSocketType == SOCK_STREAM ) );

    m_iSocketType       = iSocketType;
    m_iSocketIndex      = -1;

    m_ulMessageNumber   = 0UL;

    try
    {
        if ( m_Address.IsEmpty() )
        {
            // Don't generate an error since this may validly be NULL for a server.
            TRACE( _T( "Address is unspecified\n" ) );
        }

        if ( m_uiPort == 0 )
        {
            throw "Invalid input parameter: address or port number";
        }

        if ( ( m_pbySocketData = new BYTE [ m_ulRxBufferSize ] ) == NULL )
        {
            throw "Unable to allocate socket buffer";
        }

        m_bInitialized = true;
    }
    catch ( const char * pszMessage )
    {
        m_bInitialized = false;

        TRACE( "CThreadedSocket<tProtocol>::CThreadedSocket(), %s\n", pszMessage );
    }
    catch ( ... )
    {
        m_bInitialized = false;
        TRACE( "CThreadedSocket<tProtocol>::CThreadedSocket(), Unknown exception caught ...!\n" );
    }
}

template <class tProtocol>
CThreadedSocket<tProtocol>::CThreadedSocket( void )                                         // Not implemented.
{
    ASSERT( false );
    m_bInitialized = false;
}

template <class tProtocol>
CThreadedSocket<tProtocol>::CThreadedSocket( const CThreadedSocket &rCopy )                 // Not implemented.
{
    UNREFERENCED_PARAMETER( rCopy );

    ASSERT( false );

    m_bInitialized = false;
}

template <class tProtocol>
CThreadedSocket<tProtocol> & CThreadedSocket<tProtocol>::operator = ( const CThreadedSocket &rCopy )   // Not implemented.
{
    UNREFERENCED_PARAMETER( rCopy );

    ASSERT( false );

    m_bInitialized = false;

    return *this;
}

template <class tProtocol>
CThreadedSocket<tProtocol>::~CThreadedSocket( void )
{
    TRACE( "In CThreadedSocket<tProtocol>::~CThreadedSocket()\n" );

    // The following stratagy is used to terminate our read thread:
    // a) set the thread termination flag to false and then
    // b) close our non-lingering blocking socket which should any blocking read to abort.

    // Set the thread termination flag.

    SetActive( false );

    // Terminate the socket connection... should cause the read thread to unblock, if it's blocked.

    m_Socket.Close();

    WaitForThreadTermination( 2000 );

    // Deallocate our read holding buffer.

    if ( m_pbySocketData != NULL )
    {
        delete [] m_pbySocketData;
        m_pbySocketData = NULL;
    }
}

template <class tProtocol>
bool CThreadedSocket<tProtocol>::IsInitialized( void ) const
{
    return ( m_bInitialized && CBaseThread::IsInitialized() );
}

template <class tProtocol>
bool CThreadedSocket<tProtocol>::Connect( bool bNoDelay )
{
    bool    bSuccess          = false;

    int     iSocketError      = socketErrorNone,
            iSocketBufferSize = (int) ( 10 * m_ulRxBufferSize );

    try
    {
        BOOL    bDontLinger     = TRUE;
        DWORD   dwNonBlocking   = 0UL;

        if ( ! m_Socket.Create( 0, m_iSocketType, 0 ) )
        {
            ThrowSocketError_m();
        }

        // Here we set the socket to blocking mode. Note IOCtl alone may not be sufficient so issue 
        // an AsyncSelect(0) first.

        if ( ! m_Socket.AsyncSelect( 0 ) )
        {
            ThrowSocketError_m();
        }

        if ( ! m_Socket.IOCtl( FIONBIO, &dwNonBlocking ) )
        {
            ThrowSocketError_m();
        }

        // Next, set the socket Rx buffer size and the timeout state for the Close() operation.

        if ( ! m_Socket.SetSockOpt( SO_RCVBUF, &iSocketBufferSize, sizeof( iSocketBufferSize ) ) )
        {
            ThrowSocketError_m();
        }

        // If required, disable the Nagle algorithm to send packets imediately.

        if ( bNoDelay )
        {
            BOOL bDisableNagleAlgorithm = TRUE;

            if ( ! m_Socket.SetSockOpt( TCP_NODELAY, &bDisableNagleAlgorithm, sizeof( bDisableNagleAlgorithm ) ) )
            {
                ThrowSocketError_m();
            }
        }

        // Set the socket Tx buffer size.

        if ( ! m_Socket.SetSockOpt( SO_SNDBUF, &iSocketBufferSize, sizeof( iSocketBufferSize ) ) )
        {
            ThrowSocketError_m();
        }

        // Set non-lingering mode so that when we do a close on the socket, all the blocking calls
        // terminate. This is important so that our read thread operates and terminates correctly.

        if ( ! m_Socket.SetSockOpt( SO_DONTLINGER, &bDontLinger, sizeof( bDontLinger ) ) )
        {
            ThrowSocketError_m();
        }

        // Now, try and connect this client socket to our listening server.

        if ( ! m_Socket.Connect( m_Address, m_uiPort ) )
        {
            ThrowSocketError_m();
        }

        // Finally, release the suspended socket read thread.

        if ( ResumeThread() == 0xffffffff )
        {
            throw socketErrorThreadError;
        }

        bSuccess = true;
    }
    catch ( const ESOCKETERRORCODES &reSocketErrorCode )
    {
        bSuccess = false;
        m_iErrorCode = reSocketErrorCode;
        TRACE( "CThreadedSocket<tProtocol>::Connect(), Error code: %d ...!\n", reSocketErrorCode );
    }
    catch ( const int iErrorCode )
    {
        bSuccess = false;
        m_iErrorCode = iErrorCode;
        TRACE( "CThreadedSocket<tProtocol>::Connect(), Socket error %d ...!\n", iErrorCode );
    }
    catch (...)
    {
        bSuccess = false;
        m_iErrorCode = socketErrorUnknownException;
        TRACE( "CThreadedSocket<tProtocol>::Connect(), Unknown exception caught ...!\n" );
    }

    return bSuccess;
}

template <class tProtocol>
unsigned int CThreadedSocket<tProtocol>::Port( void ) const
{
    return m_uiPort;
}

template <class tProtocol>
bool CThreadedSocket<tProtocol>::Send( BYTE *pbyStream, unsigned long ulNumBytes )
{
    bool bSuccess = false;                              // Assume failure for now.

    if ( ulNumBytes == 0UL )
    {
        bSuccess = true;
    }
    else if ( pbyStream != NULL )
    {
        __TRY
        {
            m_WriteCriticalSection.Enter();

            m_Protocol.EncodePackets( pbyStream, ulNumBytes );

            BYTE           *pPacket;

            unsigned long   ulBytesInPacket,
                            ulNumPackets = m_Protocol.GetNumberOfPackets();

            for ( unsigned long ulPacket = 0UL; ulPacket < ulNumPackets; ulPacket++ )
            {
                if ( ( ( pPacket = m_Protocol.GetNextPacket( ulBytesInPacket ) ) != NULL ) &&
                     (   ulBytesInPacket > 0UL )                                            )
                {
                    if ( ( bSuccess = SendData( pPacket, ulBytesInPacket ) ) == FALSE )
                    {
                        break;
                    }
                }
            }
        }
        __FINALLY
        {
            m_WriteCriticalSection.Leave();
        }
        __ENDFINALLY
    }
    else
    {
        bSuccess = false;
    }

    return bSuccess;
}

// Private members.

template <class tProtocol>
bool CThreadedSocket<tProtocol>::SendData( BYTE *pbyStream, unsigned long ulNumBytes )
{
    bool            bSuccess = FALSE;
    unsigned long   ulSent;

    // Winsock doesn't guarantee all bytes will go at once, so iterate until they're gone... same for reading too.

    for ( unsigned long ulTotalSent = 0UL; ulTotalSent < ulNumBytes; )
    {
        bSuccess = ( ulSent = m_Socket.Send( &pbyStream[ ulTotalSent ], ulNumBytes - ulTotalSent ) )
                        != SOCKET_ERROR;
    
        if ( bSuccess )
        {
            ulTotalSent += ulSent;
        }
        else
        {
            break;
        }
    }

    if ( bSuccess )
    {
        ++m_ulMessageNumber;
    }

    return bSuccess;
}

template <class tProtocol>
bool CThreadedSocket<tProtocol>::CallCallback(  BYTE           *pbyData,                // Data pointer where relevant - can be NULL.
                                                unsigned long   ulNumBytes,             // Number of data bytes received - can be zero.
                                                int             iSocketIndex,           // Socket index (relevant to server end only - will be zero otherwise).
                                                unsigned long   ulTimeStamp,
                                                int             iReasonForCallback )
{
    ASSERT( m_pfnCallback != NULL );

    bool bSuccess = false;                  // Assume failure for now.

    if ( m_pfnCallback != NULL )
    {
        try
        {
            (m_pfnCallback) ( pbyData, m_pvParam, ulNumBytes, iSocketIndex, ulTimeStamp, iReasonForCallback );

            bSuccess = true;
        }
        catch ( ... )
        {
            TRACE( "CThreadedSocket<tProtocol>::CallCallback(), Callback m_pfnCallback generated an exception ... !!!\n" );
        }
    }

    return bSuccess;
}

template <class tProtocol>
int CThreadedSocket<tProtocol>::ReadSocketBytes( const int &riBytesToRead )
{
    // While the socket is open and connected and our thread is healthy execute the blocking call 
    // to retrieve a complete message.
    
    int iReceiveCode    = 0;
    int iTotalReceived  = 0;
    int iTrys           = 0;

    const int iMaxTrys  = 10;

    for ( iTotalReceived = 0; ( iTotalReceived < riBytesToRead ) && ( IsActive() ); )
    {
        // Do our receive on the data socket; this should be a blocking call on a non-lingering socket.
        //
        // NB: CAsyncSocket::Receive() is no longer used since it never returns (infinite loop) when a client 
        // disconnects without sending any bytes.
        //
        // CAsyncSocket::RecieveFrom() doesn't display this problem in that it returns immediatly but simply 
        // returns 0. In this case, we use a stratagy that when the maximum number of trys has been exceeded 
        // we assume a connection error and manually throw the socketErrorRead exception which will caught
        //
        //
        // The offending call seems to be the internally used recv() socket function; recvfrom() works as advertised.

        if ( ( iReceiveCode = m_Socket.ReceiveFrom( &m_pbySocketData[ iTotalReceived ],
                                                        riBytesToRead - iTotalReceived,
                                                            m_Address,
                                                                m_uiPort ) ) == SOCKET_ERROR )
        {
            // Note: this thread operates the data socket in blocking mode. Therefore we should 
            // not receive a WSAEWOULDBLOCK error condition, nevertheless, I've left it here in case
            // the thread is changed to non-blocking mode at some later point.

            int iErrorCode;

            if ( ( iErrorCode = m_Socket.GetLastError() ) != WSAEWOULDBLOCK )
            {
                // A terminal error condition has occured so report it and cause the thread to abort.
                // Note that this can be caused either from a true undesirable error or a deliberate 
                // abort on the thread by closing the socket.

                TRACE( _T( "CThreadedSocket<tProtocol>::ReadSocketBytes(), Receive socket error (%d) ...!\n" ), iErrorCode );
                iTotalReceived = 0;
                throw socketErrorRead;
            }
        }
        else if ( iReceiveCode == 0 )
        {
            if ( ++iTrys >= iMaxTrys )
            {
                TRACE( _T( "CThreadedSocket<tProtocol>::ReadSocketBytes(), Can't recieve data within %d trys; connection likely closed.\n" ), iMaxTrys );
                throw socketErrorRead;
            }
        }
        else
        {
            iTotalReceived += iReceiveCode;
        }
    }

    return iTotalReceived;
}

template <class tProtocol>
void CThreadedSocket<tProtocol>::WatchCycle( void )
{
    if ( ! CallCallback( NULL, 0UL, m_iSocketIndex, CSystemTime::GetTickCount(), callbackConnect ) )
    {
        TRACE( _T( "CThreadedSocket<tProtocol>::WatchCycle(), failed to make connection callback\n" ) );
    }

    try
    {
        int             iBytes                = 0;
        unsigned long   ulTimeStamp           = 0UL;

        unsigned long   ulTimeOfFirstFragment = 0UL;
        bool            bFirstFragment        = true;

        m_Protocol.ResetDecoder();

        while ( ( IsActive() ) && ( m_pbySocketData != NULL )  )
        {
            AbortIfThreadTerminateRequest_m();

            // Read the mandatory header.

            iBytes = m_Protocol.GetHeaderSize();

            if ( ReadSocketBytes( iBytes ) == iBytes )
            {
                ulTimeStamp = CSystemTime::GetTickCount();

                if ( m_Protocol.IsValidHeader( m_pbySocketData, iBytes ) )
                {
                    AbortIfThreadTerminateRequest_m();

                    if ( bFirstFragment )
                    {
                        bFirstFragment = false;
                        ulTimeOfFirstFragment = ulTimeStamp;
                    }

                    // Accumulate the header bytes.

                    m_Protocol.Add( &m_pbySocketData[ 0 ], iBytes, true );

                    // Read the optional message body, if relevant.

                    iBytes = m_Protocol.BytesFollowingHeader();

                    AbortIfThreadTerminateRequest_m();

                    if ( ( iBytes > 0 ) && ( ReadSocketBytes( iBytes ) == iBytes ) )
                    {
                        // Accumulate the packet's dynamic byte section.

                        m_Protocol.Add( &m_pbySocketData[ 0 ], iBytes, false );
                    }

                    AbortIfThreadTerminateRequest_m();

                    if ( m_Protocol.IsRecordComplete() )
                    {
                        CallCallback( m_Protocol.DecodedRecord(), m_Protocol.DecodedRecordBytes(), m_iSocketIndex, ulTimeOfFirstFragment, callbackData );
                        m_Protocol.ResetDecoder();
                        bFirstFragment = true;
                    }
                }
                else
                {
                    TRACE( _T( "Header is invalid, thread id: %lu\n" ), GetCurrentThreadId() );
                }
            }
        }
    }
    catch ( const ESOCKETERRORCODES &reCause )
    {
        TRACE( _T( "CThreadedSocket<tProtocol>::WatchCycle(), Error detected, cause : %d\n" ), static_cast<int>( reCause ) );
    }
    catch ( ... )
    {
        TRACE( "CThreadedSocket<tProtocol>::WatchCycle(), Unknown exception caught.\n" );
    }

    if ( ! CallCallback( NULL, 0UL, m_iSocketIndex, CSystemTime::GetTickCount(), callbackDisconnect ) )
    {
        TRACE( _T( "CThreadedSocket<tProtocol>::WatchCycle(), failed to make disconnection callback\n" ) );
    }

    TRACE( _T( "CThreadedSocket<tProtocol>::SocketReadThread(), Thread terminating...\n" ) );
}

template <class tProtocol>
CAsyncSocket & CThreadedSocket<tProtocol>::GetSocket( void )
{
    return m_Socket;
}

template <class tProtocol>
bool CThreadedSocket<tProtocol>::SetServerSocketOptions(    const bool  &rbNoDelay,
                                                            const int   &riSocketIndex )
{
    bool    bSuccess        = false;
    int     iSocketError    = socketErrorNone;

    m_iErrorCode = socketErrorSocketMode;

    try
    {
        BOOL    bDontLinger       = TRUE;
        DWORD   dwNonBlocking     = 0UL;
        int     iSocketBufferSize = 10 * m_ulRxBufferSize;

        m_iSocketIndex = riSocketIndex;

        // Here we set the socket to blocking mode. Note IOCtl alone may not be sufficient so issue 
        // a AsyncSelect(0) first.

        if ( ! m_Socket.AsyncSelect( 0 ) )
        {
            ThrowSocketError_m();
        }

        if ( ! m_Socket.IOCtl( FIONBIO, &dwNonBlocking ) )
        {
            ThrowSocketError_m();
        }

        // Next, set the socket Rx buffer size and the timeout state for the Close() operation.

        if ( ! m_Socket.SetSockOpt( SO_RCVBUF, &iSocketBufferSize, sizeof( iSocketBufferSize ) ) )
        {
            ThrowSocketError_m();
        }

        // If required, disable the Nagle algorithm to send packets imediately.

        if ( rbNoDelay )
        {
            BOOL bDisableNagleAlgorithm = TRUE;

            if ( ! m_Socket.SetSockOpt( TCP_NODELAY, &bDisableNagleAlgorithm, sizeof( bDisableNagleAlgorithm ) ) )
            {
                ThrowSocketError_m();
            }
        }

        // Set the socket Tx buffer size.

        if ( ! m_Socket.SetSockOpt( SO_SNDBUF, &iSocketBufferSize, sizeof( iSocketBufferSize ) ) )
        {
            ThrowSocketError_m();
        }

        if ( ! m_Socket.SetSockOpt( SO_DONTLINGER, &bDontLinger, sizeof( bDontLinger ) ) )
        {
            ThrowSocketError_m();
        }

        bSuccess = true;
    }
    catch ( const int iLastError )
    {
        bSuccess     = false;
        m_iErrorCode = iLastError;
        TRACE( "CThreadedSocket<tProtocol>::SetServerSocketOptions(), Socket error %d ...!\n", iLastError );
    }
    catch(...)
    {
        bSuccess     = false;
        m_iErrorCode = socketErrorUnknownException;
        TRACE( "CThreadedSocket<tProtocol>::SetServerSocketOptions(), Unknown exception caught ...!\n" );
    }

    return bSuccess;
}

#endif // !defined(AFX_THREADEDSOCKET_H__A9048515_BC0D_4E43_BA56_E3F8672B6A33__INCLUDED_)