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

#include "serialization_base.h"

#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include "OEMCryptoCENC.h"
#include "bump_allocator.h"
#include "marshaller_base.h"
#include "shared_memory_allocator.h"
#include "shared_memory_interface.h"

struct _Message {
  uint8_t* base;
  size_t capacity;
  size_t size;
  size_t read_offset;
  MessageStatus status;
};

typedef enum {
  TAG_INVALID,
  TAG_BOOL,
  TAG_SIZE_T,
  TAG_UINT8,
  TAG_UINT16,
  TAG_UINT32,
  TAG_UINT64,
  TAG_MEMORY,
  TAG_SHARED_MEMORY,
  TAG_EOM
} TagType;

static bool ValidMessage(Message* message) {
  if (message == NULL) {
    return false;
  }
  if (message->status != MESSAGE_STATUS_OK) {
    return false;
  }
  if (message->base == NULL) {
    message->status = MESSAGE_STATUS_NULL_POINTER_ERROR;
    return false;
  }
  if (message->read_offset > message->capacity ||
      message->size > message->capacity ||
      message->read_offset > message->size) {
    message->status = MESSAGE_STATUS_OVERFLOW_ERROR;
    return false;
  }
  return true;
}

static bool NullCheck(Message* message, const void* ptr) {
  if (!ptr) {
    message->status = MESSAGE_STATUS_NULL_POINTER_ERROR;
    return true;
  }
  return false;
}

/* Callers are expected to have validated the message already */
static void PackBytes(Message* message, const uint8_t* ptr, size_t count) {
  if (!NullCheck(message, ptr)) {
    if (count <= message->capacity - message->size) {
      memcpy(message->base + message->size, ptr, count);
      message->size += count;
    } else {
      message->status = MESSAGE_STATUS_OVERFLOW_ERROR;
    }
  }
}

static void PackTag(Message* message, TagType tag) {
  if (!ValidMessage(message)) return;
  uint8_t byte = (uint8_t)tag;
  PackBytes(message, &byte, sizeof(byte));
}

LengthType LengthFromUint32T(uint32_t length) {
  LengthType length_type = {.tag = LENGTH_TYPE_UINT32_T,
                            .type.uint32_t_value = length};
  return length_type;
}

LengthType LengthFromSizeT(size_t length) {
  LengthType length_type = {.tag = LENGTH_TYPE_SIZE_T,
                            .type.size_t_value = length};
  return length_type;
}

LengthType LengthFromUint32TPointer(const uint32_t* length) {
  LengthType length_type = {.tag = LENGTH_TYPE_UINT32_T_POINTER,
                            .type.uint32_t_pointer = length};
  return length_type;
}

LengthType LengthFromSizeTPointer(const size_t* length) {
  LengthType length_type = {.tag = LENGTH_TYPE_SIZE_T_POINTER,
                            .type.size_t_pointer = length};
  return length_type;
}

LengthType LengthFromUint32TDoublePointer(uint32_t** length) {
  LengthType length_type = {.tag = LENGTH_TYPE_UINT32_T_DOUBLE_POINTER,
                            .type.uint32_t_double_pointer = length};
  return length_type;
}

LengthType LengthFromSizeTDoublePointer(size_t** length) {
  LengthType length_type = {.tag = LENGTH_TYPE_SIZE_T_DOUBLE_POINTER,
                            .type.size_t_double_pointer = length};
  return length_type;
}

size_t LengthAsSizeT(LengthType length) {
  size_t result = 0;
  switch(length.tag) {
    case LENGTH_TYPE_SIZE_T:
      result = length.type.size_t_value;
      break;
    case LENGTH_TYPE_UINT32_T:
      result = (size_t)length.type.uint32_t_value;
      break;
    case LENGTH_TYPE_SIZE_T_POINTER:
      if (length.type.size_t_pointer) {
        result = *length.type.size_t_pointer;
      }
      break;
    case LENGTH_TYPE_UINT32_T_POINTER:
      if (length.type.uint32_t_pointer) {
        result = (size_t)*length.type.uint32_t_pointer;
      }
      break;
    case LENGTH_TYPE_SIZE_T_DOUBLE_POINTER:
      if (length.type.size_t_double_pointer && *length.type.size_t_double_pointer) {
        result = **length.type.size_t_double_pointer;
      }
      break;
    case LENGTH_TYPE_UINT32_T_DOUBLE_POINTER:
      if (length.type.uint32_t_double_pointer && *length.type.uint32_t_double_pointer) {
        result = (size_t)**length.type.uint32_t_double_pointer;
      }
      break;
  }
  return result;
}

bool LengthIsNull(LengthType length) {
  switch(length.tag) {
    case LENGTH_TYPE_SIZE_T:
    case LENGTH_TYPE_UINT32_T:
      return false;
    case LENGTH_TYPE_SIZE_T_POINTER:
      return length.type.size_t_pointer == NULL;
    case LENGTH_TYPE_UINT32_T_POINTER:
      return length.type.uint32_t_pointer == NULL;
    case LENGTH_TYPE_SIZE_T_DOUBLE_POINTER:
      return length.type.size_t_double_pointer == NULL ||
          *length.type.size_t_double_pointer == NULL;
    case LENGTH_TYPE_UINT32_T_DOUBLE_POINTER:
      return length.type.uint32_t_double_pointer == NULL ||
          *length.type.uint32_t_double_pointer == NULL;
  }
  assert(false);
  return true;
}

void ODK_Pack_bool(Message* message, const bool* value) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_BOOL);
  uint8_t b = (*value) ? 1 : 0;
  PackBytes(message, &b, sizeof(b));
}

void ODK_Pack_size_t(Message* message, const size_t* value) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_SIZE_T);
  uint64_t u64 = *value;
  uint8_t buf[sizeof(uint64_t)];
  buf[0] = (uint8_t)u64;
  buf[1] = (uint8_t)(u64 >> 8);
  buf[2] = (uint8_t)(u64 >> 16);
  buf[3] = (uint8_t)(u64 >> 24);
  buf[4] = (uint8_t)(u64 >> 32);
  buf[5] = (uint8_t)(u64 >> 40);
  buf[6] = (uint8_t)(u64 >> 48);
  buf[7] = (uint8_t)(u64 >> 56);

  PackBytes(message, buf, sizeof(buf));
}

void ODK_Pack_c_str(Message* message, const char* value) {
  if (!ValidMessage(message) || NullCheck(message, value)) return;
  size_t length = strlen(value) + 1;
  ODK_PackMemory(message, (const uint8_t*)value, LengthFromSizeT(length));
}

void ODK_Pack_uint8_t(Message* message, const uint8_t* value) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_UINT8);
  PackBytes(message, (const uint8_t*)value, sizeof(*value));
}

void ODK_Pack_uint16_t(Message* message, const uint16_t* value) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_UINT16);
  uint8_t buf[sizeof(uint16_t)];
  buf[0] = (uint8_t)(*value);
  buf[1] = (uint8_t)(*value >> 8);

  PackBytes(message, buf, sizeof(buf));
}

void ODK_Pack_uint32_t(Message* message, const uint32_t* value) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_UINT32);
  uint8_t buf[sizeof(uint32_t)];
  buf[0] = (uint8_t)(*value);
  buf[1] = (uint8_t)(*value >> 8);
  buf[2] = (uint8_t)(*value >> 16);
  buf[3] = (uint8_t)(*value >> 24);

  PackBytes(message, buf, sizeof(buf));
}

void ODK_Pack_uint64_t(Message* message, const uint64_t* value) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_UINT64);
  uint8_t buf[sizeof(uint64_t)];
  buf[0] = (uint8_t)(*value);
  buf[1] = (uint8_t)(*value >> 8);
  buf[2] = (uint8_t)(*value >> 16);
  buf[3] = (uint8_t)(*value >> 24);
  buf[4] = (uint8_t)(*value >> 32);
  buf[5] = (uint8_t)(*value >> 40);
  buf[6] = (uint8_t)(*value >> 48);
  buf[7] = (uint8_t)(*value >> 56);

  PackBytes(message, buf, sizeof(buf));
}

/*
 * Convenience function used by the serializer to
 * pack a bool without having to allocate a local
 * variable.
 */
bool ODK_PackBoolValue(Message* message, bool value) {
  ODK_Pack_bool(message, (const bool*)&value);
  return value;
}

/*
 * Pack a boolean indicating if the pointer value is NULL
 */
bool ODK_PackIsNull(Message* message, const void* pointer) {
  return ODK_PackBoolValue(message, pointer == NULL);
}

/*
 * Pack a boolean indicating if the address value is NULL.  It's
 * effectivey the same as PackIsNull but provided for symmetry with
 * UnpackAlloc.
 */
void ODK_PackAlloc(Message* message, const void* addr) {
  ODK_PackIsNull(message, addr);
}

/*
 * Pack a range of memory given by address and length.  First pack a
 * bool indicating if the address is NULL. If the address is non-null
 * and length is non-null then pack the length and |length| bytes from
 * memory beginning at |address|.
 */
void ODK_PackMemory(Message* message, const uint8_t* address,
                    LengthType length) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_MEMORY);
  ODK_PackBoolValue(message, address == NULL || LengthIsNull(length));
  if (address && !LengthIsNull(length)) {
    size_t count = LengthAsSizeT(length);
    ODK_Pack_size_t(message, &count);
    if (LengthAsSizeT(length) > 0) {
      PackBytes(message, address, LengthAsSizeT(length));
    }
  }
}

/*
 * Pack memory representing an array of fundamental types, given by
 * address and length.
 */
void ODK_PackArray(Message* message, const uint8_t* address, size_t length) {
  ODK_PackMemory(message, address, LengthFromSizeT(length));
}

/*
 * Pack a variable length array of objects at the base address given
 * by |address|, with |length| elements of size |size|. The ObjPacker is
 * a pack function able to pack elements of the array.
 */
void ODK_PackObjArray(Message* message, const uint8_t* objs, LengthType length,
                      size_t size, ObjPacker packer) {
  bool is_null = objs == NULL || LengthIsNull(length);
  if (!ODK_PackBoolValue(message, is_null)) {
    for (size_t i = 0; i < LengthAsSizeT(length); i++) {
      (*packer)(message, objs + i * size);
    }
  }
}

/*
 * On the REE side, if address and length are not null, map the shared
 * memory segment specified by index into our address space using the
 * shared memory bump allocator and pack fields into the message so
 * that the receiver can map the corresponding segment if needed. Copy
 * the input data into the shared segment.
 *
 * Parameters:
 *   message - The message to pack into
 *   index - the index that identifies which segment to map
 *   address - base address of the local memory buffer. If address is
 *     null, shared memory will not be allocated, which is indicated
 *     by a packed bool value so the receiver will also know not to
 *     allocate the segment.
 *   length - the length of the segment. A segment of this size will
 *     be mapped if length is non-null and *length is non-zero.
 *     otherwise the segment will not be mapped.
 */
void ODK_PackSharedInputBuffer(Message* message, uint16_t index,
                               const uint8_t* address, LengthType length) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_SHARED_MEMORY);
  ODK_PackBoolValue(message, address == NULL || LengthIsNull(length));
  if (address && !LengthIsNull(length)) {
    uint8_t* shared_address = SharedMemory_Allocate(index, LengthAsSizeT(length));
    if (shared_address) {
      memcpy(shared_address, address, LengthAsSizeT(length));
    }
  }
}

/*
 * Pass information to the receiver indicating whether shared memory
 * needs to be allocated for the buffer indicated by |address| and |length|.
 *
 * Parameters:
 *   message - The message to pack into
 *   address - base address of the local memory buffer. If null, a shared
 *     memory segment will not need to be mapped by the receiver.
 *   length - the length of the segment. If null, or if the length is 0,
 *     a shared memory segment will not need to be mapped by the receiver.
 *
 * Returns:
 *   void
 */
void ODK_PackSharedOutputBuffer(Message* message, const uint8_t* address,
                                LengthType length) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_SHARED_MEMORY);
  ODK_PackBoolValue(message, address == NULL || LengthIsNull(length));
}

void ODK_PackEOM(Message* message) {
  if (!ValidMessage(message)) return;
  PackTag(message, TAG_EOM);
}

/**************************** Unpack Functions *******************************/

static void UnpackBytes(Message* message, uint8_t* ptr, size_t count) {
  if (!ValidMessage(message)) return;
  if (count <= message->size - message->read_offset) {
    if (ptr) {
      memcpy(ptr, message->base + message->read_offset, count);
    }
    message->read_offset += count;
  } else {
    message->status = MESSAGE_STATUS_UNDERFLOW_ERROR;
  }
}

static bool CheckTag(Message* message, TagType tag) {
  uint8_t byte = TAG_INVALID;
  UnpackBytes(message, &byte, sizeof(byte));
  if (tag != (TagType)byte) {
    message->status = MESSAGE_STATUS_INVALID_TAG_ERROR;
    return false;
  }
  return true;
}

void ODK_Unpack_bool(Message* message, bool* value) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_BOOL)) return;
  uint8_t b = 0;
  UnpackBytes(message, &b, sizeof(b));
  *value = b ? true : false;
}

void ODK_Unpack_size_t(Message* message, size_t* value) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_SIZE_T)) return;
  uint8_t buf[sizeof(uint64_t)];
  UnpackBytes(message, buf, sizeof(buf));
  uint64_t u64 = *value;
  u64 = buf[0];
  u64 |= (uint64_t)buf[1] << 8;
  u64 |= (uint64_t)buf[2] << 16;
  u64 |= (uint64_t)buf[3] << 24;
  u64 |= (uint64_t)buf[4] << 32;
  u64 |= (uint64_t)buf[5] << 40;
  u64 |= (uint64_t)buf[6] << 48;
  u64 |= (uint64_t)buf[7] << 56;
  *value = u64;
}

void ODK_Unpack_c_str(Message* message, char** value) {
  if (!ValidMessage(message) || NullCheck(message, value)) return;
  *value = NULL;
  uint8_t* ptr = NULL;
  ODK_UnpackPointerToMemory(message, &ptr);
  if (ptr) {
    /* calculate length, be careful to stay within length of message */
    size_t length = 0;
    uint8_t* p = ptr;
    while (*p && p < message->base + message->size) {
      length++;
      p++;
    }
    void* buffer = BumpAllocate(length + 1);
    if (buffer) {
      memcpy(buffer, ptr, length);
      *value = buffer;
    }
  }
}

void ODK_Unpack_uint8_t(Message* message, uint8_t* value) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_UINT8)) return;
  UnpackBytes(message, (uint8_t*)value, sizeof(*value));
}

void ODK_Unpack_uint16_t(Message* message, uint16_t* value) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_UINT16)) return;
  uint8_t buf[sizeof(uint16_t)];
  UnpackBytes(message, buf, sizeof(buf));
  *value = buf[0];
  *value |= (uint16_t)buf[1] << 8;
}

void ODK_Unpack_uint32_t(Message* message, uint32_t* value) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_UINT32)) return;
  uint8_t buf[sizeof(uint32_t)];
  UnpackBytes(message, buf, sizeof(buf));
  *value = buf[0];
  *value |= (uint32_t)buf[1] << 8;
  *value |= (uint32_t)buf[2] << 16;
  *value |= (uint32_t)buf[3] << 24;
}

void ODK_Unpack_uint64_t(Message* message, uint64_t* value) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_UINT64)) return;
  uint8_t buf[sizeof(uint64_t)];
  UnpackBytes(message, buf, sizeof(buf));
  *value = buf[0];
  *value |= (uint64_t)buf[1] << 8;
  *value |= (uint64_t)buf[2] << 16;
  *value |= (uint64_t)buf[3] << 24;
  *value |= (uint64_t)buf[4] << 32;
  *value |= (uint64_t)buf[5] << 40;
  *value |= (uint64_t)buf[6] << 48;
  *value |= (uint64_t)buf[7] << 56;
}

bool ODK_UnpackBoolValue(Message* message) {
  bool value = false;
  ODK_Unpack_bool(message, &value);
  return value;
}

/* Return true if the pointer was packed as NULL */
bool ODK_UnpackIsNull(Message* message) { return ODK_UnpackBoolValue(message); }

/*
 * If the pointer was packed as NULL, return NULL. Otherwise
 * return the number of bytes of memory from the bump allocator
 * requested by |size|
 */
uint8_t* ODK_UnpackAlloc(Message* message, size_t size) {
  return ODK_UnpackIsNull(message) ? NULL : BumpAllocate(size);
}

/*
 * If the pointer was packed as NULL, return NULL. Otherwise
 * return the number of bytes of memory from the bump allocator
 * requested by |size|
 */
uint8_t* ODK_UnpackAllocBuffer(Message* message, LengthType length, size_t size) {
  uint8_t* buffer = NULL;
  if (!ODK_UnpackIsNull(message) && !LengthIsNull(length)) {
    buffer = BumpAllocate(LengthAsSizeT(length) * size);
  }
  return buffer;
}

/*
 * If the pointer was packed as NULL, return NULL. Otherwise unpack
 * the array of objects into memory allocated from the bump allocator.
 */
void ODK_UnpackObjArray(Message* message, uint8_t** address, LengthType count,
                        size_t size, ObjUnpacker unpacker) {
  if (address) {
    *address = NULL;
  }
  if (!ODK_UnpackIsNull(message)) {
    if (address && !LengthIsNull(count)) {
      *address = BumpAllocate(LengthAsSizeT(count) * size);
      if (*address) {
        for (size_t i = 0; i < LengthAsSizeT(count); i++) {
          (*unpacker)(message, (*address) + size * i);
        }
      }
    }
  }
}

/*
 * Unpack a range of memory representing an array of fundamental types
 * given a pointer to a buffer and a length. The caller is responsible
 * for ensuring that the buffer is large enough to handle the number
 * of bytes specified by |length|. First unpack a bool that indicates
 * if the pointer to the memory is NULL and if not unpack the
 * requested number of bytes into the address specified by |address|.
 */
void ODK_UnpackArray(Message* message, uint8_t* address, size_t length) {
  if (!ValidMessage(message) || NullCheck(message, address)) return;
  uint8_t* array = NULL;
  ODK_UnpackPointerToMemory(message, &array);
  if (array) {
    memcpy(address, array, length);
  }
}

/*
 * Unpack a pointer to memory within the received message and return
 * it. First unpack a bool that indicates if the address was passed in
 * as NULL. If so set *address to NULL and return. Otherwise set
 * *|address| to the current message read offset and unpack the length
 * of the memory range. Then increment the message read offset by that
 * amount.
 */
void ODK_UnpackPointerToMemory(Message* message, uint8_t** address) {
  if (!ValidMessage(message) || NullCheck(message, address)) return;
  if (!CheckTag(message, TAG_MEMORY)) return;
  bool is_null = true;
  ODK_Unpack_bool(message, &is_null);
  if (is_null) {
    *address = NULL;
    return;
  } else {
    size_t length;
    ODK_Unpack_size_t(message, &length);
    size_t new_offset;
    if (__builtin_add_overflow(message->read_offset, length, &new_offset) ||
        new_offset > message->size) {
      message->status = MESSAGE_STATUS_OVERFLOW_ERROR;
    } else {
      *address = message->base + message->read_offset;
      message->read_offset = new_offset;
    }
  }
}

/*
 * Unpack fields from |message| and use them to map a shared memory
 * segment using the shared memory bump allocator.  If the packed bool
 * indicates that the output buffer is non-null, allocate a shared
 * memory segment identified by |index| of the size indicated by
 * |length|.
 *
 * Parameters:
 *   message - the message to unpack from
 *   index - the index that identifies which segment to map
 *   length - the length of the segment
 * Returns:
 *   The address of the mapped segment
 */
uint8_t* ODK_UnpackSharedBuffer(Message* message, uint16_t index,
                                LengthType length) {
  if (!ValidMessage(message)) return NULL;
  if (!CheckTag(message, TAG_SHARED_MEMORY)) return NULL;
  if (!ODK_UnpackIsNull(message)) {
    return SharedMemory_Allocate(index, LengthAsSizeT(length));
  }
  return NULL;
}

/*
 * In the REE, unpack fields from |message| and use them to map a
 * shared memory segment using the shared memory bump allocator.  If
 * the packed bool indicates that the output buffer is non-null,
 * allocate a shared memory segment identified by |index| of the size
 * indicated by |length|. Copy the buffer from shared memory to
 * the OEMCrypto API parameter.
 *
 * Parameters:
 *   message - the message to unpack from
 *   address - the OEMCrypto API parameter for the output buffer
 *   index - the index that identifies which segment to map
 *   length - the length of the segment
 */
void ODK_UnpackSharedOutputBuffer(Message* message, uint16_t index,
                                  uint8_t** address, LengthType length) {
  if (address && !LengthIsNull(length)) {
    uint8_t* shared_address = ODK_UnpackSharedBuffer(message, index, length);
    if (*address && LengthAsSizeT(length) > 0) {
      memcpy(*address, shared_address, LengthAsSizeT(length));
    }
  }
}

/* end of message */
void ODK_UnpackEOM(Message* message) {
  if (!ValidMessage(message)) return;
  if (!CheckTag(message, TAG_EOM)) return;
  if (message->read_offset != message->size) {
    message->status = MESSAGE_STATUS_END_OF_MESSAGE_ERROR;
  }
}