Heap corruption using zlib inflate

Hi

I've written a class called ByteStream that will support encryption, decryption, compression and decompression. But I've ran into a problem with the decryption.

Here is the code (feel free to use it if you like to):

Code:

---

#ifndef __BYTESTREAM__H__
#define __BYTESTREAM__H__

#include <malloc.h>
#include <memory.h>

#include "../zlib.h"

#define CHUNK (256)

/*-----------------------------------------------------
ByteStream - by Daniel Lindén
-----------------------------------------------------*/
class ByteStream
{
public:

        //Constructor and destructor
        ByteStream();
        ~ByteStream() { free(m_ByteList); }

        //Write data to the ByteStream, type of your choice :)
        template <typename _T> void Write(const _T &data);
        //Read data, and you can choose the type here too
        template <typename _T> _T Read() const;

        //Get the length of the stream
        unsigned int GetLength() const { return m_DataLength; }
        //Fetch the stored data, this one gets a null terminator
        void GetData(unsigned char* &dest) const;
        //Fetch the stored data, no null terminator here
        void GetData(unsigned char* &dest, unsigned int *len) const;
        //Set the ByteStream data, WARNING: Delets all current data!!!
        void SetData(const unsigned char * data, unsigned int length);

        //Compress the stream with the zlib deflate function
        bool Deflate(int level = -1);
        //Uncompress the stream with the zlib inflate function
        bool Inflate();

        //Encrypt the data using the Tiny Encryption Algorithm
        void Encrypt(const unsigned long const key[4]);
        //Decrypt the data using the Tiny Encryption Algorithm
        void Decrypt(const unsigned long const key[4]);

        //Set where you want to read
        void SetReadHead(unsigned int location);

private:

        //Dont you dare touching those variables!
        unsigned char* m_ByteList;
        unsigned int m_ReadLocation;
        unsigned int m_DataLength;

};

extern "C"
{

        //encipher and decipher code is taken from this page: http://www.simonshepherd.supanet.com/source.htm
        void encipher(const unsigned long *const v,unsigned long *const w,
                const unsigned long * const k)
        {
                register unsigned long      y=v[0],z=v[1],sum=0,delta=0x9E3779B9,n=32;

                while(n-->0)
                {
                        y += (z << 4 ^ z >> 5) + z ^ sum + k[sum&3];
                        sum += delta;
                        z += (y << 4 ^ y >> 5) + y ^ sum + k[sum>>11 & 3];
                }

                w[0]=y; w[1]=z;
        }

        void decipher(const unsigned long *const v,unsigned long *const w,
                const unsigned long * const k)
        {
                register unsigned long      y=v[0],z=v[1],sum=0xC6EF3720,
                        delta=0x9E3779B9,n=32;

                /* sum = delta<<5, in general sum = delta * n */

                while(n-->0)
                {
                        z -= (y << 4 ^ y >> 5) + y ^ sum + k[sum>>11 & 3];
                        sum -= delta;
                        y -= (z << 4 ^ z >> 5) + z ^ sum + k[sum&3];
                }

                w[0]=y; w[1]=z;
        }

}
ByteStream::ByteStream()
{
        m_ByteList = 0;
        m_ReadLocation = 0;
        m_DataLength = 0;
}

template <typename _T> void ByteStream::Write(const _T &data)
{
        unsigned int size = (unsigned int)(sizeof(_T));

        if (!m_ByteList)
        {
                m_ByteList = (unsigned char*)malloc(size);
        } else
        {
                m_ByteList = (unsigned char*)realloc(m_ByteList, m_DataLength + size);
        }

        memcpy(m_ByteList + m_DataLength, &data, size);

        m_DataLength += size;
}

template <typename _T> _T ByteStream::Read() const
{
        _T retVal;
        unsigned int size = reinterpret_cast<unsigned int>(sizeof(_T));

        if (size > (m_DataLength - m_ReadLocation))
                return reinterpret_cast<_T>(NULL);

        memcpy(&retVal, m_ByteList + m_ReadLocation, size);

        m_ReadLocation += size;

        return retVal;
}

void ByteStream::SetReadHead(unsigned int location)
{
        if (location < 0)
                location = 0;

        if (location > m_DataLength)
                location = m_DataLength;

        m_ReadLocation = location;
}

void ByteStream::GetData(unsigned char* &dest) const
{
        dest = (unsigned char*)malloc(m_DataLength + 1);
        memcpy(dest, m_ByteList, m_DataLength);
        dest[m_DataLength] = '\0';
}

void ByteStream::GetData(unsigned char* &dest, unsigned int *length) const
{
        GetData(dest);

        if (length)
                *length = m_DataLength;
}

void ByteStream::SetData(const unsigned char *data, unsigned int length)
{
        if (!length || !data)
                return;

        if (m_ByteList)
        {
                m_ByteList = (unsigned char*)realloc(m_ByteList, length);
        } else
        {
                m_ByteList = (unsigned char*)malloc(length);
        }

        m_DataLength = length;

        memcpy(m_ByteList, data, length);
}
bool ByteStream::Deflate(int level)
{
        z_stream stream;
        stream.zalloc = Z_NULL;
        stream.zfree = Z_NULL;
        stream.opaque = Z_NULL;

        int ret = deflateInit(&stream, Z_DEFAULT_COMPRESSION);

        if (ret != Z_OK)
                return false;

        stream.avail_in = m_DataLength;

        stream.next_in = m_ByteList;


        int outSize = m_DataLength;

        unsigned char * out = reinterpret_cast<unsigned char*>(malloc(outSize));

        if (!out)
        {       
                deflateEnd(&stream);
                return false;
        }


        int availSpace = outSize;

        int writtenSpace = 0;

        stream.next_out = out;

        do
        {

                stream.avail_out = (unsigned int)availSpace;

                ret = deflate(&stream, Z_FINISH);

                writtenSpace += availSpace - stream.avail_out;

                availSpace = stream.avail_out;

                switch(ret)
                {
                case Z_NEED_DICT:
                case Z_DATA_ERROR:
                case Z_MEM_ERROR:
                case Z_VERSION_ERROR:
                case Z_STREAM_ERROR:
                        deflateEnd(&stream);
                        free(out);
                        return false;
                        break;
                case Z_OK:
                        if (!availSpace)
                        {
                                availSpace += CHUNK;
                                out = (unsigned char*)realloc(out, writtenSpace + availSpace);
                        }
                        break;
                case Z_BUF_ERROR:
                        if (stream.avail_in != 0)
                        {
                                availSpace += CHUNK;
                                out = (unsigned char*)realloc(out, writtenSpace + availSpace);
                        } else
                        {
                                return false;
                        }
                        break;
                }

        } while (ret != Z_STREAM_END);

        m_ByteList = (unsigned char*)realloc(m_ByteList, writtenSpace);
        memcpy(m_ByteList, out, writtenSpace);
        m_DataLength = writtenSpace;
        m_ReadLocation = 0;

        deflateEnd(&stream);
        free(out);

        return true;
}

/*********************
// The interesting function
*********************/
bool ByteStream::Inflate()
{
        z_stream stream;
        stream.zalloc = Z_NULL;
        stream.zfree = Z_NULL;
        stream.opaque = Z_NULL;
        stream.next_in = m_ByteList;
        stream.avail_in = m_DataLength;

        int tmp = inflateInit(&stream);

        if (tmp != Z_OK)
                return false;

        unsigned char* out;
        unsigned int availOut = m_DataLength;
        unsigned int dataWritten = 0;

        out = (unsigned char*)malloc(m_DataLength);

        stream.next_out = out;

        do
        {
                stream.avail_out = availOut;

                tmp = inflate(&stream, Z_SYNC_FLUSH);

                dataWritten += availOut - stream.avail_out;

                availOut = stream.avail_out;

                switch(tmp)
                {
                case Z_NEED_DICT:
                case Z_DATA_ERROR:
                case Z_MEM_ERROR:
                case Z_VERSION_ERROR:
                        inflateEnd(&stream);
                        free(out);
                        return false;
                        break;
                case Z_OK:
                        if (!availOut)
                        {
                                availOut = CHUNK + stream.avail_out;
                                out = (unsigned char*)realloc(out, dataWritten + availOut);
                        }
                        break;
                case Z_BUF_ERROR:
                        if (stream.avail_in != 0)
                        {
                                availOut += CHUNK;
                                out = (unsigned char*)realloc(out, dataWritten + availOut);
                        } else
                        {
                                return false;
                        }
                        break;
                }

        } while (tmp != Z_STREAM_END);


        m_ByteList = (unsigned char*)realloc(m_ByteList, dataWritten);
        memcpy(m_ByteList, out, dataWritten);
        m_DataLength = dataWritten;
        m_ReadLocation = 0;

        inflateEnd(&stream); //This is where it "crashes"
        free(out);

        return true;
}

void ByteStream::Encrypt(const unsigned long key[4])
{
        m_DataLength = m_DataLength + 8 - (m_DataLength % 8);

        m_ByteList = (unsigned char*)realloc(m_ByteList, m_DataLength);

        for (int i = 0; i < m_DataLength; i += 8)
        {
                encipher((unsigned long*)m_ByteList + i, (unsigned long*)m_ByteList + i, key);
        }
}

void ByteStream::Decrypt(const unsigned long key[4])
{
        if (m_DataLength % 8 != 0)
                return;

        for (int i = 0; i < m_DataLength; i += 8)
        {
                encipher((unsigned long*)m_ByteList + i, (unsigned long*)m_ByteList + i, key);
        }
}
#endif

---

The interesing function is Inflate(), close to the end. The problem is that windows triggers a breakpoint about heap corruption when I call inflateEnd(), this is the message from the output window:

Code:

---

HEAP[NetworkTestSDL.exe]: HEAP: Free Heap block bb0090 modified at bb00c4 after it was freed
Windows has triggered a breakpoint in NetworkTestSDL.exe.

---

I suspect that I might do something wrong with the memory management, but I have rewritten the function twice and still got the same problem. The reason is that if I've got data in the ByteStream that can't be compressed so much, this breakpoint doesn't occur.

Example of main function where the problem show:

Code:

---

int main()
{
        ByteStream bstr;
        bstr.SetData((unsigned char*)"bla bla bla bla", strlen("bla bla bla bla")+1);
        bstr.Deflate(9);
        bstr.Inflate();
        return 0;
}

---

Example where it doesn't:

Code:

---

int main()
{
        ByteStream bstr;
        bstr.SetData((unsigned char*)"dfi798d7jadsfd2347fgj", strlen("dfi798d7jadsfd2347fgj")+1);
        bstr.Deflate(9);
        bstr.Inflate();
        return 0;
}

---

The program can keep on running after this breakpoint but I am going to use this in a server application that I want to be as stable as possible.

I would appreciate some help as I'm comptely stuck :)

Thanks in advance

Daniel