// ProtocolStack.cpp
// Implementation of class cProtocolStack

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

cMessageStack::cMessage::cMessage( BYTE *pbData )
: m_pbData( 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 ) ];

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

   // First calculate the number of fragments we should have
   m_FragmentMask.SetSize (pHeader->m_wFragmentCount);

   insertFragment( pbData );
}

cMessageStack::cMessage::cMessage( const cMessage &msg )
: m_pbData( msg.m_pbData ),
m_bOwn( msg.m_bOwn )
{
	m_FragmentMask = msg.m_FragmentMask;
   msg.m_bOwn = false;
}

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

cMessageStack::cMessage &cMessageStack::cMessage::operator= ( const cMessage &msg )
{
   if( m_bOwn )
      delete[] m_pbData;

   m_pbData = msg.m_pbData;
   m_FragmentMask = msg.m_FragmentMask;
   m_bOwn = msg.m_bOwn;

   msg.m_bOwn = false;

   return *this;
}

bool cMessageStack::cMessage::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::cMessage::insertFragment( BYTE *pFragmentStart )
{
   cMessageHeader *pHeader = reinterpret_cast< cMessageHeader * >( pFragmentStart );

   BYTE *pInsertLocation = m_pbData + sizeof( cMessageHeader ) + ( MESSAGE_BODY * pHeader->m_wFragmentIndex );

   ::memcpy( pInsertLocation, pFragmentStart + sizeof( cMessageHeader ), pHeader->m_wFragmentLength - sizeof( cMessageHeader ) );

    m_FragmentMask.GotFragment (pHeader->m_wFragmentIndex);

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

DWORD cMessageStack::cMessage::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( cCallback *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 & 0x00400000)
	{
		// 0 extra header bytes
	}
	if (pHeader->m_dwFlags & 0x00800000)
	{
		// 4 extra header bytes
		dwOffset += 4;
	}

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

	splitPacket( dwLength - dwOffset, 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
      cMessage 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 );
   }
}
