// ProtocolStack.cpp
// Implementation of class cProtocolStack

#include "stdafx.h"
#include "ProtocolStack.h"

cMessageStack::cProtocolMessage::cProtocolMessage( BYTE *pbData )
: m_pbData( NULL ),
m_pbReceived( NULL ),
m_bOwn( false )
{
   cMessageHeader *pHeader = reinterpret_cast< cMessageHeader * >( pbData );

   if( pHeader->m_wFragmentCount == 1 )
   {
      // This is the only fragment, we use a direct pointer for dispatching
      m_pbData = pbData;
      return;
   }

   // This message has multiple fragments, create a buffer
   m_pbData = new BYTE[ calcMessageLength( pbData ) ];
   m_pbReceived = new bool[ pHeader->m_wFragmentCount ];
   ::memset( m_pbReceived, 0, sizeof( bool ) * pHeader->m_wFragmentCount );
   m_bOwn = true;

   // Two initializers, with 
   // Copy the header into the new buffer
   ::memcpy( m_pbData, pbData, sizeof( cMessageHeader ) );
   insertFragment( pbData );
}

cMessageStack::cProtocolMessage::cProtocolMessage( const cProtocolMessage &msg )
: m_pbData( msg.m_pbData ),
m_bOwn( msg.m_bOwn ),
m_pbReceived( msg.m_pbReceived )
{
   //msg.m_pbData = NULL;
   msg.m_bOwn = false;
   //msg.m_pbReceived = NULL;
}

cMessageStack::cProtocolMessage::~cProtocolMessage()
{
   if( m_bOwn )
   {
      delete[] m_pbData;
	  delete[] m_pbReceived;
   }
}

cMessageStack::cProtocolMessage &cMessageStack::cProtocolMessage::operator= ( const cProtocolMessage &msg )
{
   if( &msg == this )
      // Do not copy to self
      return *this;

   if( m_bOwn )
   {
		delete[] m_pbData;
		delete[] m_pbReceived;
   }

   m_pbData = msg.m_pbData;
   m_pbReceived = msg.m_pbReceived;
   m_bOwn = msg.m_bOwn;

   //msg.m_pbData = NULL;
   msg.m_bOwn = false;
   //msg.m_pbReceived = NULL;

   return *this;
}

bool cMessageStack::cProtocolMessage::isComplete() const
{
   if( m_pbReceived == NULL )
      // This is a single fragment
      return true;

   // Walk through all of the sections and make sure that we've found all the fragments
   bool *pbEndReceived = m_pbReceived + getMessageHeader()->m_wFragmentCount;
   for( bool *i = m_pbReceived; i != pbEndReceived; ++ i )
   {
      if( !(*i) )
         return false;
   }

   // No unreceiced fragments found - it must be complete
   return true;
}

bool cMessageStack::cProtocolMessage::fragmentMatch( BYTE *pFragmentStart )
{
   cMessageHeader *pThis = getMessageHeader(),
      *pOther = reinterpret_cast< cMessageHeader * >( pFragmentStart );

   return ( pThis->m_dwObjectID == pOther->m_dwObjectID && pThis->m_dwSequence == pOther->m_dwSequence );
}

#define MESSAGE_BODY 448

void cMessageStack::cProtocolMessage::insertFragment( BYTE *pFragmentStart )
{
   cMessageHeader *pHeader = reinterpret_cast< cMessageHeader * >( pFragmentStart );

   if( m_pbReceived[ pHeader->m_wFragmentIndex ] )
      // This fragment has already been inserted
      return;

   BYTE *pInsertLocation = m_pbData + sizeof( cMessageHeader ) + ( MESSAGE_BODY * pHeader->m_wFragmentIndex );
   ::memcpy( pInsertLocation, pFragmentStart + sizeof( cMessageHeader ), pHeader->m_wFragmentLength - sizeof( cMessageHeader ) );

   // Mark this fragment as received
   m_pbReceived[ pHeader->m_wFragmentIndex ] = true;

   if( pHeader->m_wFragmentIndex == ( pHeader->m_wFragmentCount - 1 ) )
      getMessageHeader()->m_wFragmentLength = ( pHeader->m_wFragmentCount - 1 ) * MESSAGE_BODY + pHeader->m_wFragmentLength;
}

DWORD cMessageStack::cProtocolMessage::calcMessageLength( BYTE *pFragmentStart )
{
   cMessageHeader *pHeader = reinterpret_cast< cMessageHeader * >( pFragmentStart );

   return ( pHeader->m_wFragmentIndex == ( pHeader->m_wFragmentCount - 1 ) ) ? 
      ( pHeader->m_wFragmentCount - 1 ) * MESSAGE_BODY + pHeader->m_wFragmentLength : pHeader->m_wFragmentCount * MESSAGE_BODY + sizeof( cMessageHeader );
}

cMessageStack::cMessageStack()
: m_pCallback( NULL )
{
}

cMessageStack::~cMessageStack()
{
   // Make sure we clean up and intermediate steps
   stop();
}

void cMessageStack::start( ACMessageSink *pCallback )
{
   // Setting this enables messages
   m_pCallback = pCallback;
}

void cMessageStack::stop()
{
   // Clean out the message list since these will never be matched
   m_messages.clear();
   m_pCallback = NULL;
}

void cMessageStack::processPacket( DWORD dwLength, BYTE *pbPayload )
{
	// First filter based on our status and the message type
	if( m_pCallback == NULL )
		// We currently only support send or receive type messages
		return;

	// AC messages must be at least as big as a packet header.+ a message header
	if (dwLength < (sizeof (cPacketHeader) + sizeof (cMessageHeader)))
		return;

	// Filter non-application messages
	cPacketHeader *pHeader = reinterpret_cast< cPacketHeader * >( pbPayload );

	if((pHeader->m_wTotalSize) != (dwLength - sizeof(cPacketHeader)))
		return;

	DWORD dwOffset = 0;
	if (pHeader->m_dwFlags & 0x00100000)
		// 8 extra header bytes
		dwOffset += 8;

	if (pHeader->m_dwFlags & 0x00200000)
		// 6 extra header bytes
		dwOffset += 6;

   if (pHeader->m_dwFlags & 0x00800000)
		// 4 extra header bytes
		dwOffset += 4;

	if (dwLength < (dwOffset + sizeof (cMessageHeader)))
		return;

	splitPacket( dwLength - dwOffset, pbPayload + dwOffset );
}

void cMessageStack::processPacketOut( DWORD dwLength, const BYTE *pbPayload )
{
	// First filter based on our status and the message type
	if( m_pCallback == NULL )
		// We currently only support send or receive type messages
		return;

	// AC messages must be at least as big as a packet header.+ a message header
	if (dwLength < (sizeof (cPacketHeader) + sizeof (cMessageHeader)))
		return;

	// Filter non-application messages
	const cPacketHeader *pHeader = reinterpret_cast< const cPacketHeader * >( pbPayload );

	if((pHeader->m_wTotalSize) != (dwLength - sizeof(cPacketHeader)))
		return;

	DWORD dwOffset = 0;
	if (pHeader->m_dwFlags & 0x00100000)
		// 8 extra header bytes
		dwOffset += 8;

	if (pHeader->m_dwFlags & 0x00200000)
		// 6 extra header bytes
		dwOffset += 6;

   if (pHeader->m_dwFlags & 0x00800000)
		// 4 extra header bytes
		dwOffset += 4;

	if (dwLength < (dwOffset + sizeof (cMessageHeader)))
		return;

	splitPacketOut( dwLength - dwOffset, const_cast<const BYTE *>(pbPayload + dwOffset) );
}

void cMessageStack::splitPacket( DWORD dwLength, BYTE *pbPayload )
{
   DWORD dwFragLength = 0;
   BYTE *i_end_packet = pbPayload + dwLength;

   for( BYTE *iFrag = pbPayload + sizeof( cPacketHeader ); iFrag != i_end_packet; iFrag += dwFragLength )
   {
      dwFragLength = reinterpret_cast< cMessageHeader * >( iFrag )->m_wFragmentLength;

      if( dwFragLength == 0 || ( iFrag + dwFragLength ) > i_end_packet )
         // We are off the end somehow
         return;

      // First walk through our cached fragments and see if one will take this fragment
      for( cMessageList::iterator i = m_messages.begin(); i != m_messages.end(); ++ i )
      {
         if( i->fragmentMatch( iFrag ) )
         {
            i->insertFragment( iFrag );
            if( i->isComplete() )
            {
               // Remove it from the list and dispatch the message
               m_pCallback->onMessage( *i );
               m_messages.erase( i );
            }
            break;
         }
      }

      if( i != m_messages.end() )
         // The fragment was found in the list
         continue;

      // Ok, we have a new fragment on our hands - generate the fragment object
      cProtocolMessage msg( iFrag );

      if( msg.isComplete() )
         // It's only one piece, dispatch right away
         m_pCallback->onMessage( msg );
      else
         // This requires more parts add it into the queue to wait
         m_messages.push_back( msg );
   }
}

void cMessageStack::splitPacketOut( DWORD dwLength, const BYTE *pbPayload )
{
   DWORD dwFragLength = 0;
   BYTE *i_end_packet = const_cast< BYTE * >(pbPayload) + dwLength;

   for( BYTE *iFrag = const_cast< BYTE * >(pbPayload) + sizeof( cPacketHeader ); iFrag != i_end_packet; iFrag += dwFragLength )
   {
      dwFragLength = reinterpret_cast< const cMessageHeader * >( iFrag )->m_wFragmentLength;

      if( dwFragLength == 0 || ( iFrag + dwFragLength ) > i_end_packet )
         // We are off the end somehow
         return;

      // First walk through our cached fragments and see if one will take this fragment
      for( cMessageList::iterator i = m_messages.begin(); i != m_messages.end(); ++ i )
      {
         if( i->fragmentMatch( iFrag ) )
         {
            i->insertFragment( iFrag );
            if( i->isComplete() )
            {
               // Remove it from the list and dispatch the message
               m_pCallback->onMessageOut( *i );
               m_messages.erase( i );
            }
            break;
         }
      }

      if( i != m_messages.end() )
         // The fragment was found in the list
         continue;

      // Ok, we have a new fragment on our hands - generate the fragment object
      cProtocolMessage msg( iFrag );

      if( msg.isComplete() )
         // It's only one piece, dispatch right away
         m_pCallback->onMessageOut( msg );
      else
         // This requires more parts add it into the queue to wait
         m_messages.push_back( msg );
   }
}