/*
 * 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 "special_cases.h"

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

#include "deserializer.h"
#include "OEMCryptoCENC.h"
#include "serialization_base.h"
#include "shared_memory_allocator.h"
#include "shared_memory_interface.h"
#include "serializer.h"

/*
 * Special cases due to union & shared memory
 */

/*
 * Pack the destination buffer parameter to OEMCrypto_DecryptCENC and
 * OEMCrypto_CopyBuffer
 */
void ODK_Pack_OEMCrypto_DestBufferDesc(Message* message,
                                       const OEMCrypto_DestBufferDesc* obj) {
  if (obj == NULL) {
    SetStatus(message, MESSAGE_STATUS_NULL_POINTER_ERROR);
    return;
  }
  ODK_Pack_uint32_t(message, (uint32_t*)&obj->type);
  switch (obj->type) {
    case OEMCrypto_BufferType_Clear: {
      ODK_Pack_size_t(message, &obj->buffer.clear.max_length);
      ODK_PackSharedOutputBuffer(message, obj->buffer.clear.address,
                                 LengthFromSizeT(obj->buffer.clear.max_length));
      break;
    }
    case OEMCrypto_BufferType_Secure:
      /* secure memory - pass handle, length, offset */
      ODK_Pack_uint64_t(message, (uint64_t*)&obj->buffer.secure.handle);
      ODK_Pack_size_t(message, &obj->buffer.secure.max_length);
      ODK_Pack_size_t(message, &obj->buffer.secure.offset);
      break;
    case OEMCrypto_BufferType_Direct:
      ODK_Pack_bool(message, &obj->buffer.direct.is_video);
      break;
  }
}

/*
 * Unpack the destination buffer parameter to OEMCrypto_DecryptCENC and
 * OEMCrypto_CopyBuffer
 */
void ODK_Unpack_OEMCrypto_DestBufferDesc(Message* message,
                                         OEMCrypto_DestBufferDesc* obj) {
  if (obj == NULL) {
    SetStatus(message, MESSAGE_STATUS_NULL_POINTER_ERROR);
    return;
  }
  ODK_Unpack_uint32_t(message, (uint32_t*)&obj->type);
  switch (obj->type) {
    case OEMCrypto_BufferType_Clear: {
      ODK_Unpack_size_t(message, &obj->buffer.clear.max_length);
      uint8_t* shared_address =
          ODK_UnpackSharedBuffer(message, DEST_BUFFER_INDEX,
                                 LengthFromSizeT(obj->buffer.clear.max_length));
      if (!obj->buffer.clear.address) {
        // unpacking request in TEE - set the address to shared memory
        obj->buffer.clear.address = shared_address;
      }
      break;
    }
    case OEMCrypto_BufferType_Secure:
      /* secure memory handle - deal with later */
      ODK_Unpack_uint64_t(message, (uint64_t*)&obj->buffer.secure.handle);
      ODK_Unpack_size_t(message, &obj->buffer.secure.max_length);
      ODK_Unpack_size_t(message, &obj->buffer.secure.offset);
      break;
    case OEMCrypto_BufferType_Direct:
      ODK_Unpack_bool(message, &obj->buffer.direct.is_video);
      break;
  }
}

/*
 * Special serialization cases due to some parameters being defined as
 * pointers into other parameters.
 */

void ODK_Pack_RewrapDeviceRSAKey_Request(
    Message* msg, OEMCrypto_SESSION session,
    const uint8_t* message,
    size_t message_length,
    const uint8_t* signature,
    size_t signature_length,
    const uint32_t* unaligned_nonce,
    const uint8_t* enc_rsa_key,
    size_t enc_rsa_key_length,
    const uint8_t* enc_rsa_key_iv,
    const uint8_t* wrapped_rsa_key,
    const size_t* wrapped_rsa_key_length) {
  uint32_t api_value = 18; /* from _oecc18 */
  ODK_Pack_uint32_t(msg, &api_value);
  ODK_Pack_size_t(msg, &message_length);
  ODK_Pack_size_t(msg, &signature_length);
  ODK_Pack_size_t(msg, &enc_rsa_key_length);
  ODK_PackNullable_size_t(msg, wrapped_rsa_key_length);
  ODK_Pack_uint32_t(msg, &session);
  ODK_PackSharedInputBuffer(msg, 0, message, LengthFromSizeT(message_length));
  ODK_PackSharedInputBuffer(msg, 1, signature, LengthFromSizeT(signature_length));

  size_t unaligned_nonce_offset = 0;
  if (__builtin_sub_overflow((uintptr_t)unaligned_nonce, (uintptr_t)message,
                             &unaligned_nonce_offset)) {
    SetStatus(msg, MESSAGE_STATUS_PARSE_ERROR);
    return;
  }
  ODK_Pack_size_t(msg, &unaligned_nonce_offset);

  size_t enc_rsa_key_offset = 0;
  if (__builtin_sub_overflow((uintptr_t)enc_rsa_key, (uintptr_t)message,
                             &enc_rsa_key_offset)) {
    SetStatus(msg, MESSAGE_STATUS_PARSE_ERROR);
    return;
  }
  ODK_Pack_size_t(msg, &enc_rsa_key_offset);

  size_t enc_rsa_key_iv_offset = 0;
  if (__builtin_sub_overflow((uintptr_t)enc_rsa_key_iv, (uintptr_t)message,
                             &enc_rsa_key_iv_offset)) {
    SetStatus(msg, MESSAGE_STATUS_PARSE_ERROR);
    return;
  }
  ODK_Pack_size_t(msg, &enc_rsa_key_iv_offset);
  ODK_PackAlloc(msg, wrapped_rsa_key);
  ODK_PackEOM(msg);
}

void ODK_Unpack_RewrapDeviceRSAKey_Request(
    Message* msg, OEMCrypto_SESSION* session,
    SharedMemory** message,
    size_t* message_length,
    SharedMemory** signature,
    size_t* signature_length,
    uint32_t** unaligned_nonce,
    uint8_t** enc_rsa_key,
    size_t* enc_rsa_key_length,
    uint8_t** enc_rsa_key_iv,
    uint8_t** wrapped_rsa_key,
    size_t** wrapped_rsa_key_length) {
  uint32_t api_value = ~(uint32_t)0;
  ODK_Unpack_uint32_t(msg, &api_value);
  if (api_value != 18) SetStatus(msg, MESSAGE_STATUS_API_VALUE_ERROR);
  ODK_Unpack_size_t(msg, message_length);
  ODK_Unpack_size_t(msg, signature_length);
  ODK_Unpack_size_t(msg, enc_rsa_key_length);
  ODK_UnpackNullable_size_t(msg, wrapped_rsa_key_length);
  ODK_Unpack_uint32_t(msg, session);
  ODK_UnpackSharedBuffer(msg, 0, LengthFromSizeTPointer(message_length));
  ODK_UnpackSharedBuffer(msg, 1, LengthFromSizeTPointer(signature_length));
  *message = SharedMemory_GetAddress(0);
  *signature = SharedMemory_GetAddress(1);

  size_t unaligned_nonce_offset = 0;
  uintptr_t unaligned_nonce_ptr = 0;
  ODK_Unpack_size_t(msg, &unaligned_nonce_offset);
  if (__builtin_add_overflow((uintptr_t)*message, unaligned_nonce_offset,
                             &unaligned_nonce_ptr)) {
    SetStatus(msg, MESSAGE_STATUS_PARSE_ERROR);
    return;
  }
  *unaligned_nonce = (uint32_t*)unaligned_nonce_ptr;

  size_t enc_rsa_key_offset = 0;
  uintptr_t enc_rsa_key_ptr = 0;
  ODK_Unpack_size_t(msg, &enc_rsa_key_offset);
  if (__builtin_add_overflow((uintptr_t)*message, enc_rsa_key_offset,
                             &enc_rsa_key_ptr)) {
    SetStatus(msg, MESSAGE_STATUS_PARSE_ERROR);
    return;
  }
  *enc_rsa_key = (uint8_t*)enc_rsa_key_ptr;

  size_t enc_rsa_key_iv_offset = 0;
  uintptr_t enc_rsa_key_iv_ptr = 0;
  ODK_Unpack_size_t(msg, &enc_rsa_key_iv_offset);
  if (__builtin_add_overflow((uintptr_t)*message, enc_rsa_key_iv_offset,
                             &enc_rsa_key_iv_ptr)) {
    SetStatus(msg, MESSAGE_STATUS_PARSE_ERROR);
    return;
  }
  *enc_rsa_key_iv = (uint8_t*)enc_rsa_key_iv_ptr;
  *wrapped_rsa_key = (uint8_t*)ODK_UnpackAllocBuffer(
      msg, LengthFromSizeTDoublePointer(wrapped_rsa_key_length),
      sizeof(uint8_t));
  ODK_UnpackEOM(msg);
}