odkitee/serialization/shared_memory_allocator.c

/*
 * Copyright 2020 Google LLC. All Rights Reserved. This file and proprietary
 * source code may only be used and distributed under the Widevine Master
 * License Agreement.
 */

#include <string.h>

#include "shared_memory_allocator.h"
#include "shared_memory_interface.h"
#include "oemcrypto_overflow.h"

/*
 * This file implements the allocator for shared memory. Shared memory
 * is managed using two different schemes. There is a memory region,
 * called the "pool" that is allocated once at initialization and kept
 * open for the lifetime of OEMCrypto, and discrete allocations that
 * are made on each OEMCrypto call where additional shared memory is
 * required.
 *
 * Every shared memory allocation is identified by an index number,
 * which is provided by the caller. Allowing the caller to specify the
 * index allows the code generator to refer to specific regions
 * explicitly, which is needed to correlate segments across
 * serialization endpoints. The index numbers start from 0.  Given an
 * index number, the address and size of the allocation can be
 * returned. The code generator requires that the allocation algorithm
 * be deterministic and repeatable, i.e. that a specific sequence of
 * allocations when replayed after a reset will result in the same
 * shared memory regions being mapped to the same indexes.
 *
 * The memory pool is managed with a bump allocator. The pool is
 * allocated from the porting layer's ODK_SharedMemory interface. The
 * pool is divided into "segments" which are allocated consecutively
 * from the pool by AllocateSegment. The allocator is reset before
 * each api call, which removes all segments.
 *
 * Shared memory allocations that are too large to put in the pool are
 * allocated directly from ODK_SharedMemory and are released at the
 * end of each OEMCrypto call. A threshold is defined (e.g. pool
 * size/4) that determines whether allocations are from the pool or
 * from the discrete regions.
 */

/* Region ID for the shared memory pool */
#define SHARED_MEMORY_POOL_REGION_ID 0

/* The ID of the next discrete region to allocate */
static size_t next_discrete_region_id_ = SHARED_MEMORY_POOL_REGION_ID + 1;

/*
 * Descriptor for a segment allocated from the pool.
 */
typedef struct {
  uint32_t size;
  uint32_t offset;
} Segment;

/*
 * Descriptor for an allocation, which may be either from the pool
 * or from a discrete region.
 */
typedef enum {FREE, POOL_SEGMENT, DISCRETE_REGION} AllocationType;
typedef struct {
  AllocationType type;
  union {
    Segment segment;
    ODK_SharedHandle* handle;
  } u;
} Allocation;

/*
 * Maximum number of allocations allowed. Indexes go from 0 to
 * MAX_ALLOCATIONS - 1
 */
#define MAX_ALLOCATIONS 16
static Allocation allocations_[MAX_ALLOCATIONS];

/*
 * Shared memory handle for the bump allocator pool
 */
static ODK_SharedHandle *pool_handle_;

/*
 * The offset from pool_address of the next segment that will be
 * allocated.
 */
static size_t next_segment_offset_ = 0;

/*
 * Initial size of the memory pool
 */
static size_t pool_size_ = 0;

/*
 * Initialize shared memory state variables. Set the initial pool size
 * to the specified value.
 */
void SharedMemory_Initialize(size_t pool_size) {
  ODK_SharedMemory_Initialize();
  memset(&allocations_[0], 0, sizeof(allocations_));
  pool_size_ = pool_size;
  pool_handle_ = NULL;
  next_discrete_region_id_ = SHARED_MEMORY_POOL_REGION_ID + 1;
}

/*
 * Release all shared memory resources from the current process.
 */
void SharedMemory_Terminate(void) {
  SharedMemory_Reset();
  ODK_SharedMemory_Free(pool_handle_);
  pool_handle_ = NULL;
  ODK_SharedMemory_Terminate();
}


/*
 * Allocate shared memory as a discrete region
 */
static uint8_t* AllocateAsRegion(uint16_t index, size_t size) {
  ODK_SharedHandle* handle =
      ODK_SharedMemory_Allocate(next_discrete_region_id_, size);
  if (handle) {
    Allocation* alloc = &allocations_[index];
    alloc->type = DISCRETE_REGION;
    alloc->u.handle = handle;
    next_discrete_region_id_++;
  }
  return ODK_SharedMemory_GetAddress(handle);
}

/*
 * Allocate a segment of shared memory from the pool
 */
static uint8_t* AllocateFromPool(uint16_t index, size_t size) {
  if (!pool_handle_) {
    pool_handle_ = ODK_SharedMemory_Allocate(SHARED_MEMORY_POOL_REGION_ID,
                                             pool_size_);
    if (!pool_handle_) {
      return NULL;
    }
    next_segment_offset_ = 0;
  }
  size_t new_offset = 0;
  if (AddOverflowUX(next_segment_offset_, size, &new_offset)) {
    return NULL;
  }
  if (new_offset > ODK_SharedMemory_GetSize(pool_handle_)) {
    return AllocateAsRegion(index, size);
  }
  uint8_t* segment_address =
      ODK_SharedMemory_GetAddress(pool_handle_) + next_segment_offset_;

  Allocation* alloc = &allocations_[index];
  alloc->type = POOL_SEGMENT;
  Segment* segment = &alloc->u.segment;
  segment->size = size;
  segment->offset = next_segment_offset_;
  next_segment_offset_ = new_offset;
  return segment_address;
}

/*
 * Allocate shared memory from either the pool or a discrete region.
 * If the requested memory size is larger than 1/4 of the pool size, or
 * if the pool is full, allocate as a discrete region.
 *
 * Parameters:
 *  index - the index number of the allocation, assigned by the caller.
 *    indexes start at 0. The index cannot exceed MAX_ALLOCATIONS.
 *
 * Returns:
 *  The address of the allocation or NULL if the index has already
 *  been allocated or allocation fails.
 */
uint8_t* SharedMemory_Allocate(uint16_t index, size_t size) {
  if (index >= MAX_ALLOCATIONS || allocations_[index].type != FREE) {
    return NULL;
  }
  if (size <= pool_size_ / 4) {
    return AllocateFromPool(index, size);
  }  else {
    return AllocateAsRegion(index, size);
  }
}

/*
 * Get the size of a shared memory segment.
 *
 * Parameters:
 *   index - identifies the allocation. Valid values are
 *   0..MAX_ALLOCATIONS-1
 *
 *  Returns:
 *    The size of the allocation indicated by |index| or 0 if the
 *    index is invalid or no matching allocation was found
 */
size_t SharedMemory_GetSize(uint16_t index) {
  if (index >= MAX_ALLOCATIONS) {
    return 0;
  }
  Allocation* alloc = &allocations_[index];
  switch (alloc->type) {
    case DISCRETE_REGION:
      return ODK_SharedMemory_GetSize(alloc->u.handle);
    case POOL_SEGMENT:
      return alloc->u.segment.size;
    default:
      break;
  }
  return 0;
}

/*
 * Get the address of a shared memory segment.
 *
 * Parameters:
 *   index - identifies the allocation. Valid values are
 *   0..MAX_ALLOCATIONS-1
 *
 *  Returns:
 *    The address of the allocation indicated by |index| or NULL if the
 *    index is invalid or no matching allocation was found
 */
uint8_t* SharedMemory_GetAddress(uint16_t index) {
  if (index >= MAX_ALLOCATIONS) {
    return NULL;
  }
  Allocation* alloc = &allocations_[index];
  switch (alloc->type) {
    case DISCRETE_REGION:
      return ODK_SharedMemory_GetAddress(alloc->u.handle);
    case POOL_SEGMENT:
      if (!pool_handle_) {
        return NULL;
      }
      return ODK_SharedMemory_GetAddress(pool_handle_) +
          alloc->u.segment.offset;
    default:
      break;
  }
  return NULL;
}

/*
 * Set the allocator to its initial state, where there are no segments
 * allocated from the pool and there are no discrete allocations.
 */
void SharedMemory_Reset(void) {
  next_segment_offset_ = 0;
  for (size_t i = 0; i < MAX_ALLOCATIONS; i++) {
    Allocation* alloc = &allocations_[i];
    if (alloc->type == DISCRETE_REGION) {
      ODK_SharedMemory_Free(alloc->u.handle);
    }
  }
  memset(&allocations_[0], 0, sizeof(allocations_));
  next_discrete_region_id_ = SHARED_MEMORY_POOL_REGION_ID + 1;
}