/* Copyright (c) 2002-2018 Dovecot authors, see the included COPYING file */

/* @UNSAFE: whole file */

#include "lib.h"

#include "safe-memset.h"

#include "mempool.h"

/*

 * As the name implies, alloconly pools support only allocating memory.

 * Memory freeing is not supported, except as a special case - the pool's

 * last allocation can be freed.  Additionally, p_realloc() also tries to

 * grow an existing allocation if and only if it is the last allocation,

 * otherwise it just allocates a new memory area and copies the data there.

 *

 * Alloconly pools are commonly used for an object that builds its state

 * from many memory allocations, but doesn't change (much of) its state.

 * It is simpler to free such an object by destroying the entire memory

 * pool.

 *

 * Implementation

 * ==============

 *

 * Each alloconly pool contains a pool structure (struct alloconly_pool) to

 * keep track of alloconly-specific pool information and one or more blocks

 * (struct pool_block) that keep track of ranges of memory used to back the

 * allocations.  The blocks are kept in a linked list implementing a stack.

 * The block size decreases the further down the stack one goes.

 *

 * +-----------+

 * | alloconly |

 * |    pool   |

 * +-----+-----+

 *       |

 *       | block  +------------+ next  +------------+ next

 *       \------->| pool block |------>| pool block |------>...

 *                +------------+       +------------+

 *                |   <data>   |       |   <data>   |

 *                      .                    .

 *                      .                    .

 *                      .              |   <data>   |

 *                      .              +------------+

 *                |   <data>   |

 *                +------------+

 *

 * Creation

 * --------

 *

 * When an alloconly pool is created, one block is allocated.  This block is

 * large enough to hold the necessary internal structures (struct

 * alloconly_pool and struct pool_block) and still have enough space to

 * satisfy allocations for at least the amount of space requested by the

 * consumer via the size argument to pool_alloconly_create().

 *

 * Allocation

 * ----------

 *

 * Each allocation (via p_malloc()) checks the top-most block to see whether

 * or not it has enough space to satisfy the allocation.  If there is not

 * enough space, it allocates a new block (via block_alloc()) to serve as

 * the new top-most block.  This newly-allocated block is guaranteed to have

 * enough space for the allocation.  Then, regardless of whether or not a

 * new block was allocated, the allocation code reserves enough space in the

 * top-most block for the allocation and returns a pointer to it to the

 * caller.

 *

 * The free space tracking within each block is very simple.  In addition to

 * keeping track of the size of the block, the block header contains a

 * "pointer" to the beginning of free space.  A new allocation simply moves

 * this pointer by the number of bytes allocated.

 *

 * Reallocation

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

 *

 * If the passed in allocation is the last allocation in a block and there

 * is enough space after it, the allocation is resized.  Otherwise, a new

 * buffer is allocated (see Allocation above) and the contents are copied

 * over.

 *

 * Freeing

 * -------

 *

 * Freeing of the last allocation moves the "pointer" to free space back by

 * the size of the last allocation.

 *

 * Freeing of any other allocation is a no-op.

 *

 * Clearing

 * --------

 *

 * Clearing the pool is supposed to return the pool to the same state it was

 * in when it was first created.  To that end, the alloconly pool frees all

 * the blocks allocated since the pool's creation.  The remaining block

 * (allocated during creation) is reset to consider all the space for

 * allocations as available.

 *

 * In other words, the per-block free space tracking variables are set to

 * indicate that the full block is available and that there have been no

 * allocations.

 *

 * Finally, if the pool was created via pool_alloconly_create_clean(), all

 * blocks are safe_memset()/memset() to zero before being free()d.

 *

 * Destruction

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

 *

 * Destroying a pool first clears it (see above).  The clearing leaves the

 * pool in a minimal state with only one block allocated.  This remaining

 * block may be safe_memset() to zero if the pool was created with

 * pool_alloconly_create_clean().

 *

 * Since the pool structure itself is allocated from the first block, this

 * final call to free() will release the memory allocated for struct

 * alloconly_pool and struct pool.

 */

#ifndef DEBUG

#  define POOL_ALLOCONLY_MAX_EXTRA MEM_ALIGN(1)

#else

#  define POOL_ALLOCONLY_MAX_EXTRA \

  (MEM_ALIGN(sizeof(size_t)) + MEM_ALIGN(1) + MEM_ALIGN(SENTRY_COUNT))

#endif

struct alloconly_pool {

  struct pool pool;

  int refcount;

  struct pool_block *block;

#ifdef DEBUG

  const char *name;

  size_t base_size;

  bool disable_warning;

#endif

  bool clean_frees;

};

struct pool_block {

  struct pool_block *prev;

  size_t size;

  size_t left;

  size_t last_alloc_size;

  /* unsigned char data[]; */

};

#define SIZEOF_POOLBLOCK (MEM_ALIGN(sizeof(struct pool_block)))

#define POOL_BLOCK_DATA(block) \

  ((unsigned char *) (block) + SIZEOF_POOLBLOCK)

#define DEFAULT_BASE_SIZE MEM_ALIGN(sizeof(struct alloconly_pool))

#ifdef DEBUG

#  define CLEAR_CHR 0xde

#  define SENTRY_COUNT 8

#else

#  define SENTRY_COUNT 0

#  define CLEAR_CHR 0

#endif

static const char *pool_alloconly_get_name(pool_t pool);

static void pool_alloconly_ref(pool_t pool);

static void pool_alloconly_unref(pool_t *pool);

static void *pool_alloconly_malloc(pool_t pool, size_t size);

static void pool_alloconly_free(pool_t pool, void *mem);

static void *pool_alloconly_realloc(pool_t pool, void *mem,

            size_t old_size, size_t new_size);

static void pool_alloconly_clear(pool_t pool);

static size_t pool_alloconly_get_max_easy_alloc_size(pool_t pool);

static void block_alloc(struct alloconly_pool *pool, size_t size);

static const struct pool_vfuncs static_alloconly_pool_vfuncs = {

  pool_alloconly_get_name,

  pool_alloconly_ref,

  pool_alloconly_unref,

  pool_alloconly_malloc,

  pool_alloconly_free,

  pool_alloconly_realloc,

  pool_alloconly_clear,

  pool_alloconly_get_max_easy_alloc_size

};

static const struct pool static_alloconly_pool = {

  .v = &static_alloconly_pool_vfuncs,

  .alloconly_pool = TRUE,

  .datastack_pool = FALSE

};

#ifdef DEBUG

static void check_sentries(struct pool_block *block)

{

  const unsigned char *data = POOL_BLOCK_DATA(block);

  size_t i, max_pos, alloc_size, used_size;

  used_size = block->size - block->left;

  for (i = 0; i < used_size; ) {

    alloc_size = *(size_t *)(data + i);

    if (alloc_size == 0 || used_size - i < alloc_size)

      i_panic("mempool-alloconly: saved alloc size broken");

    i += MEM_ALIGN(sizeof(alloc_size));

    max_pos = i + MEM_ALIGN(alloc_size + SENTRY_COUNT);

    i += alloc_size;

    for (; i < max_pos; i++) {

      if (data[i] != CLEAR_CHR)

        i_panic("mempool-alloconly: buffer overflow");

    }

  }

  if (i != used_size)

    i_panic("mempool-alloconly: used_size wrong");

  /* The unused data must be NULs */

  for (; i < block->size; i++) {

    if (data[i] != '\0')

      i_unreached();

  }

  if (block->prev != NULL)

    check_sentries(block->prev);

}

#endif

pool_t pool_alloconly_create(const char *name ATTR_UNUSED, size_t size)

{

  struct alloconly_pool apool, *new_apool;

  size_t min_alloc = SIZEOF_POOLBLOCK +

    MEM_ALIGN(sizeof(struct alloconly_pool) + SENTRY_COUNT);

  (void) COMPILE_ERROR_IF_TRUE(POOL_ALLOCONLY_MAX_EXTRA >

             (SSIZE_T_MAX - POOL_MAX_ALLOC_SIZE));

#ifdef DEBUG

  min_alloc += MEM_ALIGN(strlen(name) + 1 + SENTRY_COUNT) +

    sizeof(size_t)*2;

#endif

  /* create a fake alloconly_pool so we can call block_alloc() */

  i_zero(&apool);

  apool.pool = static_alloconly_pool;

  apool.refcount = 1;

  if (size < min_alloc)

    size = nearest_power(size + min_alloc);

  block_alloc(&apool, size);

  /* now allocate the actual alloconly_pool from the created block */

  new_apool = p_new(&apool.pool, struct alloconly_pool, 1);

  *new_apool = apool;

#ifdef DEBUG

  if (str_begins(name, MEMPOOL_GROWING, &name) ||

      getenv("DEBUG_SILENT") != NULL)

    new_apool->disable_warning = TRUE;

  new_apool->name = p_strdup(&new_apool->pool, name);

  /* set base_size so p_clear() doesn't trash alloconly_pool structure. */

  new_apool->base_size = new_apool->block->size - new_apool->block->left;

  new_apool->block->last_alloc_size = 0;

#endif

  /* the first pool allocations must be from the first block */

  i_assert(new_apool->block->prev == NULL);

  return &new_apool->pool;

}

pool_t pool_alloconly_create_clean(const char *name, size_t size)

{

  struct alloconly_pool *apool;

  pool_t pool;

  pool = pool_alloconly_create(name, size);

  apool = container_of(pool, struct alloconly_pool, pool);

  apool->clean_frees = TRUE;

  return pool;

}

static void pool_alloconly_free_block(struct alloconly_pool *apool ATTR_UNUSED,

              struct pool_block *block)

{

#ifdef DEBUG

  safe_memset(block, CLEAR_CHR, SIZEOF_POOLBLOCK + block->size);

#else

  if (apool->clean_frees) {

    safe_memset(block, CLEAR_CHR,

          SIZEOF_POOLBLOCK + block->size);

  }

#endif

  free(block);

}

static void

pool_alloconly_free_blocks_until_last(struct alloconly_pool *apool)

{

  struct pool_block *block;

  /* destroy all blocks but the oldest, which contains the

     struct alloconly_pool allocation. */

  while (apool->block->prev != NULL) {

    block = apool->block;

    apool->block = block->prev;

    pool_alloconly_free_block(apool, block);

  }

}

static void pool_alloconly_destroy(struct alloconly_pool *apool)

{

  /* destroy all but the last block */

  pool_alloconly_free_blocks_until_last(apool);

  /* destroy the last block */

  pool_alloconly_free_block(apool, apool->block);

}

static const char *pool_alloconly_get_name(pool_t pool ATTR_UNUSED)

{

#ifdef DEBUG

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  return apool->name;

#else

  return "alloconly";

#endif

}

static void pool_alloconly_ref(pool_t pool)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  apool->refcount++;

}

static void pool_alloconly_unref(pool_t *pool)

{

  struct alloconly_pool *apool =

    container_of(*pool, struct alloconly_pool, pool);

  /* erase the pointer before freeing anything, as the pointer may

     exist inside the pool's memory area */

  *pool = NULL;

  if (--apool->refcount > 0)

    return;

  pool_external_refs_unref(&apool->pool);

  pool_alloconly_destroy(apool);

}

static void block_alloc(struct alloconly_pool *apool, size_t size)

{

  struct pool_block *block;

  i_assert(size > SIZEOF_POOLBLOCK);

  i_assert(size <= SSIZE_T_MAX);

  if (apool->block != NULL) {

    /* each block is at least twice the size of the previous one */

    if (size <= apool->block->size)

      size += apool->block->size;

    /* avoid crashing in nearest_power() if size is too large */

    size = I_MIN(size, SSIZE_T_MAX);

    size = nearest_power(size);

    /* nearest_power() could have grown size to SSIZE_T_MAX+1 */

    size = I_MIN(size, SSIZE_T_MAX);

#ifdef DEBUG

    if (!apool->disable_warning) {

      /* i_debug() overwrites unallocated data in data

         stack, so make sure everything is allocated before

         calling it. */

      t_buffer_alloc_last_full();

      i_debug("Growing pool '%s' with: %zu",

          apool->name, size);

    }

#endif

  }

  block = calloc(size, 1);

  if (unlikely(block == NULL)) {

    i_fatal_status(FATAL_OUTOFMEM, "block_alloc(%zu"

             "): Out of memory", size);

  }

  block->prev = apool->block;

  apool->block = block;

  block->size = size - SIZEOF_POOLBLOCK;

  block->left = block->size;

}

static void *pool_alloconly_malloc(pool_t pool, size_t size)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  void *mem;

  size_t alloc_size;

#ifndef DEBUG

  alloc_size = MEM_ALIGN(size);

#else

  alloc_size = MEM_ALIGN(sizeof(size)) + MEM_ALIGN(size + SENTRY_COUNT);

#endif

  if (apool->block->left < alloc_size) {

    /* we need a new block */

    block_alloc(apool, alloc_size + SIZEOF_POOLBLOCK);

  }

  mem = POOL_BLOCK_DATA(apool->block) +

    (apool->block->size - apool->block->left);

  apool->block->left -= alloc_size;

  apool->block->last_alloc_size = alloc_size;

#ifdef DEBUG

  memcpy(mem, &size, sizeof(size));

  mem = PTR_OFFSET(mem, MEM_ALIGN(sizeof(size)));

  /* write CLEAR_CHRs to sentry */

  memset(PTR_OFFSET(mem, size), CLEAR_CHR,

         MEM_ALIGN(size + SENTRY_COUNT) - size);

#endif

  return mem;

}

static void pool_alloconly_free(pool_t pool, void *mem)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  /* we can free only the last allocation */

  if (POOL_BLOCK_DATA(apool->block) +

      (apool->block->size - apool->block->left -

       apool->block->last_alloc_size) == mem) {

    memset(mem, 0, apool->block->last_alloc_size);

    apool->block->left += apool->block->last_alloc_size;

    apool->block->last_alloc_size = 0;

  }

}

static bool pool_alloconly_try_grow(struct alloconly_pool *apool, void *mem, size_t size)

{

  /* see if we want to grow the memory we allocated last */

  if (POOL_BLOCK_DATA(apool->block) +

      (apool->block->size - apool->block->left -

       apool->block->last_alloc_size) == mem) {

    /* yeah, see if we can grow */

    if (apool->block->left >= size-apool->block->last_alloc_size) {

      /* just shrink the available size */

      apool->block->left -=

        size - apool->block->last_alloc_size;

      apool->block->last_alloc_size = size;

      return TRUE;

    }

  }

  return FALSE;

}

static void *pool_alloconly_realloc(pool_t pool, void *mem,

            size_t old_size, size_t new_size)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  unsigned char *new_mem;

  i_assert(old_size < SIZE_MAX);

  if (new_size <= old_size)

    return mem;

  new_size = MEM_ALIGN(new_size);

  /* see if we can directly grow it */

  if (!pool_alloconly_try_grow(apool, mem, new_size)) {

    /* slow way - allocate + copy */

    new_mem = pool_alloconly_malloc(pool, new_size);

    if (old_size > 0)

      memcpy(new_mem, mem, old_size);

    mem = new_mem;

  }

  return mem;

}

static void pool_alloconly_clear(pool_t pool)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  size_t base_size, avail_size;

#ifdef DEBUG

  check_sentries(apool->block);

#endif

  pool_alloconly_free_blocks_until_last(apool);

  /* clear the first block */

#ifdef DEBUG

  base_size = apool->base_size;

#else

  base_size = DEFAULT_BASE_SIZE;

#endif

  avail_size = apool->block->size - base_size;

  memset(PTR_OFFSET(POOL_BLOCK_DATA(apool->block), base_size), 0,

         avail_size - apool->block->left);

  apool->block->left = avail_size;

  apool->block->last_alloc_size = 0;

}

static size_t pool_alloconly_get_max_easy_alloc_size(pool_t pool)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  return apool->block->left;

}

size_t pool_alloconly_get_total_used_size(pool_t pool)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  struct pool_block *block;

  size_t size = 0;

  i_assert(pool->v == &static_alloconly_pool_vfuncs);

  for (block = apool->block; block != NULL; block = block->prev)

    size += block->size - block->left;

  return size;

}

size_t pool_alloconly_get_total_alloc_size(pool_t pool)

{

  struct alloconly_pool *apool =

    container_of(pool, struct alloconly_pool, pool);

  struct pool_block *block;

  size_t size = 0;

  i_assert(pool->v == &static_alloconly_pool_vfuncs);

  for (block = apool->block; block != NULL; block = block->prev)

    size += block->size + SIZEOF_POOLBLOCK;

  return size;

}