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

#include <botan/bzip2.h>
#include <cstring>
#define BZ_NO_STDIO
#include <bzlib.h>

namespace Botan {

namespace {

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

/*************************************************
* Allocation Function for Bzlib                  *
*************************************************/
void bzlib_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 Bzip2 Stream                  *
*************************************************/
struct bz_stream_wrapper
   {
   bz_stream stream;
   bz_stream_wrapper()
      {
      std::memset(&stream, 0, sizeof(bz_stream));
      stream.bzalloc = bzlib_malloc;
      stream.bzfree = bzlib_free;
      stream.opaque = get_allocator();
      }
   ~bz_stream_wrapper() { std::memset(&stream, 0, sizeof(bz_stream)); }
   };

/*************************************************
* Start Compressing with Bzip                    *
*************************************************/
void Bzip_Compression::start_msg()
   {
   clear();
   bz = new bz_stream_wrapper;
   if(BZ2_bzCompressInit(&(bz->stream), level, 0, 0) != BZ_OK)
      throw Exception("Bzip_Compression: Memory allocation error");
   }

/*************************************************
* Compress Input with Bzip                       *
*************************************************/
void Bzip_Compression::write(const byte input[], u32bit length)
   {
   bz->stream.next_in = (char*)input;
   bz->stream.avail_in = length;

   while(bz->stream.avail_in != 0)
      {
      bz->stream.next_out = (char*)buffer.ptr();
      bz->stream.avail_out = buffer.size();
      BZ2_bzCompress(&(bz->stream), BZ_RUN);
      send(buffer, buffer.size() - bz->stream.avail_out);
      }
   }

/*************************************************
* Finish Compressing with Bzip                   *
*************************************************/
void Bzip_Compression::end_msg()
   {
   bz->stream.next_in = 0;
   bz->stream.avail_in = 0;

   int rc = BZ_OK;
   while(rc != BZ_STREAM_END)
      {
      bz->stream.next_out = (char*)buffer.ptr();
      bz->stream.avail_out = buffer.size();
      rc = BZ2_bzCompress(&(bz->stream), BZ_FINISH);
      send(buffer, buffer.size() - bz->stream.avail_out);
      }
   clear();
   }

/*************************************************
* Flush the Bzip Compressor                      *
*************************************************/
void Bzip_Compression::flush()
   {
   bz->stream.next_in = 0;
   bz->stream.avail_in = 0;

   int rc = BZ_OK;
   while(rc != BZ_RUN_OK)
      {
      bz->stream.next_out = (char*)buffer.ptr();
      bz->stream.avail_out = buffer.size();
      rc = BZ2_bzCompress(&(bz->stream), BZ_FLUSH);
      send(buffer, buffer.size() - bz->stream.avail_out);
      }
   }

/*************************************************
* Clean up Compression Context                   *
*************************************************/
void Bzip_Compression::clear()
   {
   if(!bz) return;
   BZ2_bzCompressEnd(&(bz->stream));
   delete bz;
   bz = 0;
   }

/*************************************************
* Decompress Input with Bzip                     *
*************************************************/
void Bzip_Decompression::write(const byte input[], u32bit length)
   {
   if(length) no_writes = false;
   bz->stream.next_in = (char*)input;
   bz->stream.avail_in = length;

   while(bz->stream.avail_in != 0)
      {
      u64bit total_bytes = ((u64bit)bz->stream.total_in_hi32 << 32) |
                           (bz->stream.total_in_lo32);

      bz->stream.next_out = (char*)buffer.ptr();
      bz->stream.avail_out = buffer.size();

      int rc = BZ2_bzDecompress(&(bz->stream));
      if(rc != BZ_OK && rc != BZ_STREAM_END)
         {
         clear();
         if(rc == BZ_DATA_ERROR)
            throw Decoding_Error("Bzip_Decompression: Data integrity error");
         if(rc == BZ_DATA_ERROR_MAGIC)
            throw Decoding_Error("Bzip_Decompression: Invalid input");
         if(rc == BZ_MEM_ERROR)
            throw Exception("Bzip_Decompression: Memory allocation error");
         throw Exception("Bzip_Decompression: Unknown decompress error");
         }
      send(buffer, buffer.size() - bz->stream.avail_out);
      if(rc == BZ_STREAM_END)
         {
         u32bit read_from_block = (((u64bit)bz->stream.total_in_hi32 << 32) |
                                  (bz->stream.total_in_lo32)) -
                                  total_bytes;
         start_msg();
         bz->stream.next_in = (char*)input + read_from_block;
         bz->stream.avail_in = length - read_from_block;
         input += read_from_block;
         length -= read_from_block;
         }
      }
   }

/*************************************************
* Start Decompressing with Bzip                  *
*************************************************/
void Bzip_Decompression::start_msg()
   {
   clear();
   bz = new bz_stream_wrapper;
   if(BZ2_bzDecompressInit(&(bz->stream), 0, small_mem) != BZ_OK)
      throw Exception("Bzip_Decompression: Memory allocation error");
   no_writes = true;
   }

/*************************************************
* Finish Decompressing with Bzip                 *
*************************************************/
void Bzip_Decompression::end_msg()
   {
   if(no_writes) return;
   bz->stream.next_in = 0;
   bz->stream.avail_in = 0;

   int rc = BZ_OK;
   while(rc != BZ_STREAM_END)
      {
      bz->stream.next_out = (char*)buffer.ptr();
      bz->stream.avail_out = buffer.size();
      rc = BZ2_bzDecompress(&(bz->stream));
      if(rc != BZ_OK && rc != BZ_STREAM_END)
         {
         clear();
         throw Exception("Bzip_Decompression: Error finalizing decompression");
         }
      send(buffer, buffer.size() - bz->stream.avail_out);
      }
   clear();
   }

/*************************************************
* Clean up Decompression Context                 *
*************************************************/
void Bzip_Decompression::clear()
   {
   if(!bz) return;
   BZ2_bzDecompressEnd(&(bz->stream));
   delete bz;
   bz = 0;
   }

}
