/*************************************************
* Zlib Compressor Source File                    *
* (C) 1999-2003 The Botan Project                *
*************************************************/

#include <botan/zlib.h>
#include <cstring>
#include <zlib.h>

namespace Botan {

namespace {

/*************************************************
* Allocation Function for Zlib                   *
*************************************************/
void* zlib_malloc(void* alloc_ptr, unsigned int n, unsigned int size)
   {
   SecureAllocator* alloc = static_cast<SecureAllocator*>(alloc_ptr);
   return alloc->allocate(n * size);
   }

/*************************************************
* Allocation Function for Zlib                   *
*************************************************/
void zlib_free(void* alloc_ptr, void* ptr)
   {
   SecureAllocator* alloc = static_cast<SecureAllocator*>(alloc_ptr);
   u32bit size = alloc->find_size(ptr);
   alloc->deallocate(ptr, size);
   }

}

/*************************************************
* Wrapper Type for Zlib Stream                   *
*************************************************/
struct zlib_stream_wrapper
   {
   z_stream stream;
   zlib_stream_wrapper()
      {
      std::memset(&stream, 0, sizeof(z_stream));
      stream.zalloc = zlib_malloc;
      stream.zfree = zlib_free;
      stream.opaque = get_allocator();
      }
   ~zlib_stream_wrapper() { std::memset(&stream, 0, sizeof(z_stream)); }
   };

/*************************************************
* Start Compressing with Zlib                    *
*************************************************/
void Zlib_Compression::start_msg()
   {
   clear();
   zlib = new zlib_stream_wrapper;
   if(deflateInit(&(zlib->stream), level) != Z_OK)
      throw Exception("Zlib_Compression: Memory allocation error");
   }

/*************************************************
* Compress Input with Zlib                       *
*************************************************/
void Zlib_Compression::write(const byte input[], u32bit length)
   {
   zlib->stream.next_in = (Bytef*)input;
   zlib->stream.avail_in = length;

   while(zlib->stream.avail_in != 0)
      {
      zlib->stream.next_out = (Bytef*)buffer.ptr();
      zlib->stream.avail_out = buffer.size();
      deflate(&(zlib->stream), Z_NO_FLUSH);
      send(buffer, buffer.size() - zlib->stream.avail_out);
      }
   }

/*************************************************
* Finish Compressing with Zlib                   *
*************************************************/
void Zlib_Compression::end_msg()
   {
   zlib->stream.next_in = 0;
   zlib->stream.avail_in = 0;

   int rc = Z_OK;
   while(rc != Z_STREAM_END)
      {
      zlib->stream.next_out = (Bytef*)buffer.ptr();
      zlib->stream.avail_out = buffer.size();
      rc = deflate(&(zlib->stream), Z_FINISH);
      send(buffer, buffer.size() - zlib->stream.avail_out);
      }
   clear();
   }

/*************************************************
* Flush the Zlib Compressor                      *
*************************************************/
void Zlib_Compression::flush()
   {
   zlib->stream.next_in = 0;
   zlib->stream.avail_in = 0;

   while(true)
      {
      zlib->stream.next_out = (Bytef*)buffer.ptr();
      zlib->stream.avail_out = buffer.size();
      deflate(&(zlib->stream), Z_FULL_FLUSH);
      send(buffer, buffer.size() - zlib->stream.avail_out);
      if(zlib->stream.avail_out == buffer.size()) break;
      }
   }

/*************************************************
* Clean up Compression Context                   *
*************************************************/
void Zlib_Compression::clear()
   {
   if(zlib)
      {
      deflateEnd(&(zlib->stream));
      delete zlib;
      zlib = 0;
      }
   }

/*************************************************
* Start Decompressing with Zlib                  *
*************************************************/
void Zlib_Decompression::start_msg()
   {
   clear();
   zlib = new zlib_stream_wrapper;
   if(inflateInit(&(zlib->stream)) != Z_OK)
      throw Exception("Zlib_Decompression: Memory allocation error");
   }

/*************************************************
* Decompress Input with Zlib                     *
*************************************************/
void Zlib_Decompression::write(const byte input[], u32bit length)
   {
   if(length) no_writes = false;

   zlib->stream.next_in = (Bytef*)input;
   zlib->stream.avail_in = length;

   while(zlib->stream.avail_in != 0)
      {
      u32bit total_bytes = zlib->stream.total_in;

      zlib->stream.next_out = (Bytef*)buffer.ptr();
      zlib->stream.avail_out = buffer.size();

      int rc = inflate(&(zlib->stream), Z_SYNC_FLUSH);
      if(rc != Z_OK && rc != Z_STREAM_END)
         {
         clear();
         if(rc == Z_DATA_ERROR)
            throw Decoding_Error("Zlib_Decompression: Data integrity error");
         if(rc == Z_NEED_DICT)
            throw Decoding_Error("Zlib_Decompression: Need preset dictionary");
         if(rc == Z_MEM_ERROR)
            throw Exception("Zlib_Decompression: Memory allocation error");
         throw Exception("Zlib_Decompression: Unknown decompress error");
         }
      send(buffer, buffer.size() - zlib->stream.avail_out);
      if(rc == Z_STREAM_END)
         {
         u32bit read_from_block = zlib->stream.total_in - total_bytes;
         start_msg();
         zlib->stream.next_in = (Bytef*)input + read_from_block;
         zlib->stream.avail_in = length - read_from_block;
         input += read_from_block;
         length -= read_from_block;
         }
      }
   }

/*************************************************
* Finish Decompressing with Zlib                 *
*************************************************/
void Zlib_Decompression::end_msg()
   {
   if(no_writes) return;
   zlib->stream.next_in = 0;
   zlib->stream.avail_in = 0;

   int rc = Z_OK;
   while(rc != Z_STREAM_END)
      {
      zlib->stream.next_out = (Bytef*)buffer.ptr();
      zlib->stream.avail_out = buffer.size();
      rc = inflate(&(zlib->stream), Z_SYNC_FLUSH);
      if(rc != Z_OK && rc != Z_STREAM_END)
         {
         clear();
         throw Exception("Zlib_Decompression: Error finalizing decompression");
         }
      send(buffer, buffer.size() - zlib->stream.avail_out);
      }
   clear();
   }

/*************************************************
* Clean up Decompression Context                 *
*************************************************/
void Zlib_Decompression::clear()
   {
   no_writes = true;

   if(zlib)
      {
      inflateEnd(&(zlib->stream));
      delete zlib;
      zlib = 0;
      }
   }

}
