/*************************************************
* Secure Allocator Source File                   *
* (C) 1999-2003 The Botan Project                *
*************************************************/

#include <botan/secalloc.h>
#include <botan/conf.h>
#include <botan/util.h>
#include <algorithm>

namespace Botan {

namespace {

/*************************************************
* Buffer Comparison Operators                    *
*************************************************/
bool Buffer_Cmp(const SecureAllocator::Buffer& a,
                const SecureAllocator::Buffer& b)
   { return ((byte*)a.buf < (byte*)b.buf); }

bool Empty_Buffer(const SecureAllocator::Buffer& block)
   { return (block.length == 0); }

}

/*************************************************
* SecureAllocator Constructor                    *
*************************************************/
SecureAllocator::SecureAllocator() :
   PREF_SIZE(Config::get_u32bit("base/memory_chunk"))
   {
   if(!PREF_SIZE)
      throw Invalid_Argument("The base/memory_chunk option is unset");
   lock = get_mutex();
   destroyed = false;
   }

/*************************************************
* SecureAllocator Destructor                     *
*************************************************/
SecureAllocator::~SecureAllocator()
   {
   delete lock;
   if(!destroyed)
      throw Invalid_State("SecureAllocator: Never released memory");
   }

/*************************************************
* Free all remaining memory                      *
*************************************************/
void SecureAllocator::destroy()
   {
   if(destroyed)
      throw Invalid_State("SecureAllocator::destroy(): Already destroyed");

   destroyed = true;
   for(u32bit j = 0; j != allocated.size(); j++)
      dealloc_block(allocated[j].buf, allocated[j].length);
   }

/*************************************************
* Return true if these buffers are contiguous    *
*************************************************/
bool SecureAllocator::are_contiguous(const Buffer& a, const Buffer& b)
   {
   if((byte*)a.buf + a.length == (byte*)b.buf)
      return true;
   return false;
   }

/*************************************************
* See if two free blocks are from the same block *
*************************************************/
bool SecureAllocator::same_buffer(Buffer& a, Buffer& b) const
   {
   return (find_block(a.buf) == find_block(b.buf));
   }

/*************************************************
* Find the block containing this memory          *
*************************************************/
u32bit SecureAllocator::find_block(void* addr) const
   {
   for(u32bit j = 0; j != allocated.size(); j++)
      {
      byte* buf_addr = (byte*)allocated[j].buf;
      if(buf_addr <= addr && addr < buf_addr + allocated[j].length)
         return j;
      }
   throw Internal_Error("SecureAllocator::find_block: no buffer found");
   }

/*************************************************
* Remove empty buffers from list                 *
*************************************************/
void SecureAllocator::remove_empty_buffers(std::vector<Buffer>& list) const
   {
   std::vector<Buffer>::iterator empty;

   empty = std::find_if(list.begin(), list.end(), Empty_Buffer);
   while(empty != list.end())
      {
      list.erase(empty);
      empty = std::find_if(list.begin(), list.end(), Empty_Buffer);
      }
   }

/*************************************************
* Allocation                                     *
*************************************************/
void* SecureAllocator::allocate(u32bit n) const
   {
   if(n == 0) return 0;
   n = round_up(n, ALIGN_TO);

   Mutex_Holder holder(lock);

   void* memory = find_free_block(n);
   if(memory)
      return memory;

   Buffer block;
   block.length = ((PREF_SIZE > n) ? PREF_SIZE : n);
   block.buf = get_block(block.length);
   if(!block.buf)
      throw Memory_Exhaustion();
   free_list.push_back(block);

   memory = find_free_block(n);
   if(memory)
      return memory;

   throw Memory_Exhaustion();
   }

/*************************************************
* Deallocation                                   *
*************************************************/
void SecureAllocator::deallocate(void* ptr, u32bit n) const
   {
   static u32bit defrag_counter = 0;
   const u32bit RUNS_TO_DEFRAGS = 4;

   if(ptr == 0 || n == 0) return;

   n = round_up(n, ALIGN_TO);

   std::memset(ptr, 0, n);

   Mutex_Holder holder(lock);

   free_list.push_back(Buffer(ptr, n));
   if(free_list.size() >= 2)
      std::inplace_merge(free_list.begin(), free_list.end() - 1,
                         free_list.end(), Buffer_Cmp);

   if(defrag_counter == 0)
      defrag_free_list();
   defrag_counter = (defrag_counter + 1) % RUNS_TO_DEFRAGS;

   for(u32bit j = 0; j != free_list.size(); j++)
      {
      bool erase = false;
      if(free_list[j].buf == 0) continue;

      for(u32bit k = 0; k != allocated.size(); k++)
         if(free_list[j].buf == allocated[k].buf &&
            free_list[j].length == allocated[k].length)
            erase = true;

      if(erase)
         {
         u32bit buf = find_block(free_list[j].buf);

         free_block(allocated[buf].buf, allocated[buf].length);

         free_list[j].buf = 0;
         free_list[j].length = 0;
         }
      }

   remove_empty_buffers(free_list);
   }

/*************************************************
* Handle Allocating New Memory                   *
*************************************************/
void* SecureAllocator::get_block(u32bit n) const
   {
   for(u32bit j = 0; j != allocated.size(); j++)
      if(!allocated[j].in_use && allocated[j].length == n)
         {
         allocated[j].in_use = true;
         return allocated[j].buf;
         }

   void* ptr = alloc_block(n);
   allocated.push_back(Buffer(ptr, n, true));
   return ptr;
   }

/*************************************************
* Handle Deallocating Memory                     *
*************************************************/
void SecureAllocator::free_block(void* ptr, u32bit n) const
   {
   if(!ptr) return;

   u32bit free_blocks = 0;
   for(u32bit j = 0; j != allocated.size(); j++)
      if(!allocated[j].in_use) free_blocks++;

   bool free_this_block = (free_blocks >= 2 || n != PREF_SIZE) ? true : false;

   for(u32bit j = 0; j != allocated.size(); j++)
      if(allocated[j].buf == ptr)
         {
         if(!allocated[j].in_use || allocated[j].length != n)
            throw Internal_Error("SecureAllocator: Size mismatch in free");

         if(free_this_block)
            {
            dealloc_block(allocated[j].buf, allocated[j].length);
            allocated[j].buf = 0;
            allocated[j].length = 0;
            remove_empty_buffers(allocated);
            }
         else
            allocated[j].in_use = false;

         return;
         }

   throw Internal_Error("SecureAllocator: Unknown pointer was freed");
   }

/*************************************************
* Defragment the free list                       *
*************************************************/
void SecureAllocator::defrag_free_list() const
   {
   if(free_list.size() < 2) return;

   for(u32bit j = 0; j != free_list.size(); j++)
      {
      if(free_list[j].length == 0)
         continue;

      if(j > 0 &&
         are_contiguous(free_list[j-1], free_list[j]) &&
         same_buffer(free_list[j-1], free_list[j]))
         {
         free_list[j].buf = free_list[j-1].buf;
         free_list[j].length += free_list[j-1].length;
         free_list[j-1].length = 0;
         }

      if(j < free_list.size() - 1 &&
         are_contiguous(free_list[j], free_list[j+1]) &&
         same_buffer(free_list[j], free_list[j+1]))
         {
         free_list[j+1].buf = free_list[j].buf;
         free_list[j+1].length += free_list[j].length;
         free_list[j].length = 0;
         }
      }
   remove_empty_buffers(free_list);
   }

/*************************************************
* Find a block on the free list                  *
*************************************************/
void* SecureAllocator::find_free_block(u32bit n) const
   {
   void* retval = 0;

   for(u32bit j = 0; j != free_list.size(); j++)
      if(free_list[j].length >= n)
         {
         retval = free_list[j].buf;

         if(free_list[j].length == n)
            free_list.erase(free_list.begin() + j);
         else if(free_list[j].length > n)
            {
            free_list[j].length -= n;
            free_list[j].buf = ((byte*)free_list[j].buf) + n;
            }
         break;
         }

   return retval;
   }

}
