/*

* (C) 2018,2019 Jack Lloyd

*

* Botan is released under the Simplified BSD License (see license.txt)

*/

#include <botan/internal/mem_pool.h>

#include <botan/mem_ops.h>

#include <algorithm>

#if defined(BOTAN_MEM_POOL_USE_MMU_PROTECTIONS)

   #include <botan/internal/os_utils.h>

#endif

namespace Botan {

/*

* Memory pool theory of operation

*

* This allocator is not useful for general purpose but works well within the

* context of allocating cryptographic keys. It makes several assumptions which

* don't work for implementing malloc but simplify and speed up the implementation:

*

* - There is some set of pages, which cannot be expanded later. These are pages

*   which were allocated, mlocked and passed to the Memory_Pool constructor.

*

* - The allocator is allowed to return null anytime it feels like not servicing

*   a request, in which case the request will be sent to calloc instead. In

*   particular, requests which are too small or too large are rejected.

*

* - Most allocations are powers of 2, the remainder are usually a multiple of 8

*

* - Free requests include the size of the allocation, so there is no need to

*   track this within the pool.

*

* - Alignment is important to the caller. For this allocator, any allocation of

*   size N is aligned evenly at N bytes.

*

* Initially each page is in the free page list. Each page is used for just one

* size of allocation, with requests bucketed into a small number of common

* sizes. If the allocation would be too big or too small it is rejected by the pool.

*

* The free list is maintained by a bitmap, one per page/Bucket. Since each

* Bucket only maintains objects of a single size, each bit set or clear

* indicates the status of one object.

*

* An allocation walks the list of buckets and asks each in turn if there is

* space. If a Bucket does not have any space, it sets a boolean flag m_is_full

* so that it does not need to rescan when asked again. The flag is cleared on

* first free from that bucket. If no bucket has space, but there are some free

* pages left, a free page is claimed as a new Bucket for that size. In this case

* it is pushed to the front of the list so it is first in line to service new

* requests.

*

* A deallocation also walks the list of buckets for the size and asks each

* Bucket in turn if it recognizes the pointer. When a Bucket becomes empty as a

* result of a deallocation, it is recycled back into the free pool. When this

* happens, the Buckets page goes to the end of the free list. All pages on the

* free list are marked in the MMU as noaccess, so anything touching them will

* immediately crash. They are only marked R/W once placed into a new bucket.

* Making the free list FIFO maximizes the time between the last free of a bucket

* and that page being writable again, maximizing chances of crashing after a

* use-after-free.

*

* Future work

* -------------

*

* The allocator is protected by a global lock. It would be good to break this

* up, since almost all of the work can actually be done in parallel especially

* when allocating objects of different sizes (which can't possibly share a

* bucket).

*

* It may be worthwhile to optimize deallocation by storing the Buckets in order

* (by pointer value) which would allow binary search to find the owning bucket.

*

* A useful addition would be to randomize the allocations. Memory_Pool would be

* changed to receive also a RandomNumberGenerator& object (presumably the system

* RNG, or maybe a ChaCha_RNG seeded with system RNG). Then the bucket to use and

* the offset within the bucket would be chosen randomly, instead of using first fit.

*

* Right now we don't make any provision for threading, so if two threads both

* allocate 32 byte values one after the other, the two allocations will likely

* share a cache line. Ensuring that distinct threads will (tend to) use distinct

* buckets would reduce this.

*

* Supporting a realloc-style API may be useful.

*/

namespace {

size_t choose_bucket(size_t n)

   {

   const size_t MINIMUM_ALLOCATION = 16;

   const size_t MAXIMUM_ALLOCATION = 256;

   if(n < MINIMUM_ALLOCATION || n > MAXIMUM_ALLOCATION)

      return 0;

   // Need to tune these

   const size_t buckets[] = {

      16, 24, 32, 48, 64, 80, 96, 112, 128, 160, 192, 256, 0,

   };

   for(size_t i = 0; buckets[i]; ++i)

      {

      if(n <= buckets[i])

         {

         return buckets[i];

         }

      }

   return 0;

   }

inline bool ptr_in_pool(const void* pool_ptr, size_t poolsize,

                        const void* buf_ptr, size_t bufsize)

   {

   const uintptr_t pool = reinterpret_cast<uintptr_t>(pool_ptr);

   const uintptr_t buf = reinterpret_cast<uintptr_t>(buf_ptr);

   return (buf >= pool) && (buf + bufsize <= pool + poolsize);

   }

// return index of first set bit

template<typename T>

size_t find_set_bit(T b)

   {

   size_t s = 8*sizeof(T) / 2;

   size_t bit = 0;

   // In this context we don't need to be const-time

   while(s > 0)

      {

      const T mask = (static_cast<T>(1) << s) - 1;

      if((b & mask) == 0)

         {

         bit += s;

         b >>= s;

         }

      s /= 2;

      }

   return bit;

   }

class BitMap final

   {

   public:

      BitMap(size_t bits) : m_len(bits)

         {

         m_bits.resize((bits + BITMASK_BITS - 1) / BITMASK_BITS);

         m_main_mask = static_cast<bitmask_type>(~0);

         m_last_mask = m_main_mask;

         if(bits % BITMASK_BITS != 0)

            m_last_mask = (static_cast<bitmask_type>(1) << (bits % BITMASK_BITS)) - 1;

         }

      bool find_free(size_t* bit);

      void free(size_t bit)

         {

         BOTAN_ASSERT_NOMSG(bit <= m_len);

         const size_t w = bit / BITMASK_BITS;

         BOTAN_ASSERT_NOMSG(w < m_bits.size());

         const bitmask_type mask = static_cast<bitmask_type>(1) << (bit % BITMASK_BITS);

         m_bits[w] = m_bits[w] & (~mask);

         }

      bool empty() const

         {

         for(size_t i = 0; i != m_bits.size(); ++i)

            {

            if(m_bits[i] != 0)

               {

               return false;

               }

            }

         return true;

         }

  private:

#if defined(BOTAN_ENABLE_DEBUG_ASSERTS)

      typedef uint8_t bitmask_type;

      enum { BITMASK_BITS = 8 };

#else

      typedef word bitmask_type;

      enum { BITMASK_BITS = BOTAN_MP_WORD_BITS };

#endif

      size_t m_len;

      bitmask_type m_main_mask;

      bitmask_type m_last_mask;

      std::vector<bitmask_type> m_bits;

   };

bool BitMap::find_free(size_t* bit)

   {

   for(size_t i = 0; i != m_bits.size(); ++i)

      {

      const bitmask_type mask = (i == m_bits.size() - 1) ? m_last_mask : m_main_mask;

      if((m_bits[i] & mask) != mask)

         {

         size_t free_bit = find_set_bit(~m_bits[i]);

         const bitmask_type bmask = static_cast<bitmask_type>(1) << (free_bit % BITMASK_BITS);

         BOTAN_ASSERT_NOMSG((m_bits[i] & bmask) == 0);

         m_bits[i] |= bmask;

         *bit = BITMASK_BITS*i + free_bit;

         return true;

         }

      }

   return false;

   }

}

class Bucket final

   {

   public:

      Bucket(uint8_t* mem, size_t mem_size, size_t item_size) :

         m_item_size(item_size),

         m_page_size(mem_size),

         m_range(mem),

         m_bitmap(mem_size / item_size),

         m_is_full(false)

         {

         }

      uint8_t* alloc()

         {

         if(m_is_full)

            {

            // I know I am full

            return nullptr;

            }

         size_t offset;

         if(!m_bitmap.find_free(&offset))

            {

            // I just found out I am full

            m_is_full = true;

            return nullptr;

            }

         BOTAN_ASSERT(offset * m_item_size < m_page_size, "Offset is in range");

         return m_range + m_item_size*offset;

         }

      bool free(void* p)

         {

         if(!in_this_bucket(p))

            return false;

         /*

         Zero also any trailing bytes, which should not have been written to,

         but maybe the user was bad and wrote past the end.

         */

         std::memset(p, 0, m_item_size);

         const size_t offset = (reinterpret_cast<uintptr_t>(p) - reinterpret_cast<uintptr_t>(m_range)) / m_item_size;

         m_bitmap.free(offset);

         m_is_full = false;

         return true;

         }

      bool in_this_bucket(void* p) const

         {

         return ptr_in_pool(m_range, m_page_size, p, m_item_size);

         }

      bool empty() const

         {

         return m_bitmap.empty();

         }

      uint8_t* ptr() const

         {

         return m_range;

         }

   private:

      size_t m_item_size;

      size_t m_page_size;

      uint8_t* m_range;

      BitMap m_bitmap;

      bool m_is_full;

   };

Memory_Pool::Memory_Pool(const std::vector<void*>& pages, size_t page_size) :

   m_page_size(page_size)

   {

   m_min_page_ptr = ~static_cast<uintptr_t>(0);

   m_max_page_ptr = 0;

   for(size_t i = 0; i != pages.size(); ++i)

      {

      const uintptr_t p = reinterpret_cast<uintptr_t>(pages[i]);

      m_min_page_ptr = std::min(p, m_min_page_ptr);

      m_max_page_ptr = std::max(p, m_max_page_ptr);

      clear_bytes(pages[i], m_page_size);

#if defined(BOTAN_MEM_POOL_USE_MMU_PROTECTIONS)

      OS::page_prohibit_access(pages[i]);

#endif

      m_free_pages.push_back(static_cast<uint8_t*>(pages[i]));

      }

   /*

   Right now this points to the start of the last page, adjust it to instead

   point to the first byte of the following page

   */

   m_max_page_ptr += page_size;

   }

Memory_Pool::~Memory_Pool()

   {

#if defined(BOTAN_MEM_POOL_USE_MMU_PROTECTIONS)

   for(size_t i = 0; i != m_free_pages.size(); ++i)

      {

      OS::page_allow_access(m_free_pages[i]);

      }

#endif

   }

void* Memory_Pool::allocate(size_t n)

   {

   if(n > m_page_size)

      return nullptr;

   const size_t n_bucket = choose_bucket(n);

   if(n_bucket > 0)

      {

      lock_guard_type<mutex_type> lock(m_mutex);

      std::deque<Bucket>& buckets = m_buckets_for[n_bucket];

      /*

      It would be optimal to pick the bucket with the most usage,

      since a bucket with say 1 item allocated out of it has a high

      chance of becoming later freed and then the whole page can be

      recycled.

      */

      for(auto& bucket : buckets)

         {

         if(uint8_t* p = bucket.alloc())

            return p;

         // If the bucket is full, maybe move it to the end of the list?

         // Otoh bucket search should be very fast

         }

      if(m_free_pages.size() > 0)

         {

         uint8_t* ptr = m_free_pages[0];

         m_free_pages.pop_front();

#if defined(BOTAN_MEM_POOL_USE_MMU_PROTECTIONS)

         OS::page_allow_access(ptr);

#endif

         buckets.push_front(Bucket(ptr, m_page_size, n_bucket));

         void* p = buckets[0].alloc();

         BOTAN_ASSERT_NOMSG(p != nullptr);

         return p;

         }

      }

   // out of room

   return nullptr;

   }

bool Memory_Pool::deallocate(void* p, size_t len) noexcept

   {

   // Do a fast range check first, before taking the lock

   const uintptr_t p_val = reinterpret_cast<uintptr_t>(p);

   if(p_val < m_min_page_ptr || p_val > m_max_page_ptr)

      return false;

   const size_t n_bucket = choose_bucket(len);

   if(n_bucket != 0)

      {

      try

         {

         lock_guard_type<mutex_type> lock(m_mutex);

         std::deque<Bucket>& buckets = m_buckets_for[n_bucket];

         for(size_t i = 0; i != buckets.size(); ++i)

            {

            Bucket& bucket = buckets[i];

            if(bucket.free(p))

               {

               if(bucket.empty())

                  {

#if defined(BOTAN_MEM_POOL_USE_MMU_PROTECTIONS)

                  OS::page_prohibit_access(bucket.ptr());

#endif

                  m_free_pages.push_back(bucket.ptr());

                  if(i != buckets.size() - 1)

                     std::swap(buckets.back(), buckets[i]);

                  buckets.pop_back();

                  }

               return true;

               }

            }

         }

      catch(...)

         {

         /*

         * The only exception throws that can occur in the above code are from

         * either the STL or BOTAN_ASSERT failures. In either case, such an

         * error indicates a logic error or data corruption in the memory

         * allocator such that it is no longer safe to continue executing.

         *

         * Since this function is noexcept, simply letting the exception escape

         * is sufficient for terminate to be called. However in this scenario

         * it is implementation defined if any stack unwinding is performed.

         * Since stack unwinding could cause further memory deallocations this

         * could result in further corruption in this allocator state. To prevent

         * this, call terminate directly.

         */

         std::terminate();

         }

      }

   return false;

   }

}