// Copyright 2019 Google LLC. All rights reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.

#include "odk.h"

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

#include "odk_overflow.h"
#include "odk_serialize.h"
#include "odk_structs.h"
#include "odk_structs_priv.h"
#include "odk_util.h"

/* @ private odk functions */

static OEMCryptoResult ODK_PrepareRequest(
    uint8_t* message, size_t message_length, size_t* core_message_length,
    ODK_MessageType message_type, const ODK_NonceValues* nonce_values,
    void* prepared_request_buffer, size_t prepared_request_buffer_length) {
  if (nonce_values == NULL || core_message_length == NULL ||
      prepared_request_buffer == NULL ||
      *core_message_length > message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  ODK_Message msg = ODK_Message_Create(message, *core_message_length);

  /* The core message should be at the beginning of the buffer, and with a
   * shorter length. */
  if (ODK_CORE_MESSAGE_SIZE > prepared_request_buffer_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_CoreMessage* core_message = (ODK_CoreMessage*)prepared_request_buffer;
  *core_message = (ODK_CoreMessage){
      message_type,
      0,
      *nonce_values,
  };

  /* Set core message length, and pack prepared request into message if the
   * message buffer has been correctly initialized by the caller. */
  switch (message_type) {
    case ODK_License_Request_Type: {
      core_message->message_length = ODK_LICENSE_REQUEST_SIZE;
      if (sizeof(ODK_PreparedLicenseRequest) > prepared_request_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Pack_ODK_PreparedLicenseRequest(
          &msg, (ODK_PreparedLicenseRequest*)prepared_request_buffer);
      break;
    }
    case ODK_Renewal_Request_Type: {
      core_message->message_length = ODK_RENEWAL_REQUEST_SIZE;
      if (sizeof(ODK_PreparedRenewalRequest) > prepared_request_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Pack_ODK_PreparedRenewalRequest(
          &msg, (ODK_PreparedRenewalRequest*)prepared_request_buffer);
      break;
    }
    case ODK_Provisioning_Request_Type: {
      core_message->message_length = ODK_PROVISIONING_REQUEST_SIZE;
      if (sizeof(ODK_PreparedProvisioningRequest) >
          prepared_request_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Pack_ODK_PreparedProvisioningRequest(
          &msg, (ODK_PreparedProvisioningRequest*)prepared_request_buffer);
      break;
    }
    case ODK_Provisioning40_Request_Type: {
      core_message->message_length = ODK_PROVISIONING40_REQUEST_SIZE;
      if (sizeof(ODK_PreparedProvisioning40Request) >
          prepared_request_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Pack_ODK_PreparedProvisioning40Request(
          &msg, (ODK_PreparedProvisioning40Request*)prepared_request_buffer);
      break;
    }
    case ODK_Renewed_Provisioning_Request_Type: {
      core_message->message_length = ODK_RENEWED_PROVISIONING_REQUEST_SIZE;
      if (sizeof(ODK_PreparedRenewedProvisioningRequest) >
          prepared_request_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Pack_ODK_PreparedRenewedProvisioningRequest(
          &msg,
          (ODK_PreparedRenewedProvisioningRequest*)prepared_request_buffer);
      break;
    }
    default: {
      return ODK_ERROR_CORE_MESSAGE;
    }
  }

  *core_message_length = core_message->message_length;
  if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK) {
    /* This is to indicate the caller that the core_message_length has been
     * appropriately set, but the message buffer is either empty or too small,
     * which needs to be initialized and filled in the subsequent call. */
    return OEMCrypto_ERROR_SHORT_BUFFER;
  }
  if (ODK_Message_GetSize(&msg) != *core_message_length) {
    /* This should not happen. Something is wrong. */
    return ODK_ERROR_CORE_MESSAGE;
  }
  return OEMCrypto_SUCCESS;
}

/* Parse the core message and verify that it has the right type. The nonce
 *  values are updated to hold the response's API version.
 */
static OEMCryptoResult ODK_ParseCoreHeader(const uint8_t* message,
                                           size_t message_length,
                                           size_t core_message_length,
                                           ODK_MessageType message_type,
                                           ODK_NonceValues* nonce_values) {
  // The core_message_length is the length of the core message, which is a
  // substring of the complete message.
  if (message == NULL || core_message_length > message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_CoreMessage core_message;
  ODK_Message msg = ODK_Message_Create((uint8_t*)message, message_length);

  /* The core message should be at the beginning of the buffer. The core message
   * is the part we are parsing. */
  ODK_Message_SetSize(&msg, core_message_length);
  Unpack_ODK_CoreMessage(&msg, &core_message);

  if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
      message_type != core_message.message_type) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  // The current offset should be the end of the header, which is the message
  // type, message length, api version, and nonce fields. The header can't be
  // larger than the whole core message. Also, the core message specifies its
  // length, which should be exactly the length of the core message buffer.
  if (ODK_Message_GetOffset(&msg) > core_message.message_length ||
      core_message.message_length != core_message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  /* We do not support future API version. Also, this function should not be
   * used for legacy licenses without a core message. */
  if (core_message.nonce_values.api_major_version > ODK_MAJOR_VERSION ||
      core_message.nonce_values.api_major_version < ODK_FIRST_VERSION) {
    return ODK_UNSUPPORTED_API;
  }
  if (nonce_values) {
    /* If the server sent us an older format, record the message's API version.
     */
    if (nonce_values->api_major_version >
        core_message.nonce_values.api_major_version) {
      // If the major version is smaller, use both values from the server.
      nonce_values->api_major_version =
          core_message.nonce_values.api_major_version;
      nonce_values->api_minor_version =
          core_message.nonce_values.api_minor_version;
    } else if (nonce_values->api_major_version ==
                   core_message.nonce_values.api_major_version &&
               nonce_values->api_minor_version >
                   core_message.nonce_values.api_minor_version) {
      // Otherwise, if the major versions are equal, but the minor is smaller,
      // then we should lower the minor version.
      nonce_values->api_minor_version =
          core_message.nonce_values.api_minor_version;
    }
  }
  return OEMCrypto_SUCCESS;
}

/* @ public odk functions */

/* @@ prepare request functions */

OEMCryptoResult ODK_PrepareCoreLicenseRequest(
    uint8_t* message, size_t message_length, size_t* core_message_size,
    const ODK_NonceValues* nonce_values,
    const ODK_MessageCounterInfo* counter_info) {
  if (core_message_size == NULL || nonce_values == NULL ||
      counter_info == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_PreparedLicenseRequest license_request = {0};
  memcpy(&license_request.counter_info, counter_info,
         sizeof(license_request.counter_info));
  return ODK_PrepareRequest(
      message, message_length, core_message_size, ODK_License_Request_Type,
      nonce_values, &license_request, sizeof(ODK_PreparedLicenseRequest));
}

OEMCryptoResult ODK_PrepareCoreRenewalRequest(uint8_t* message,
                                              size_t message_length,
                                              size_t* core_message_size,
                                              ODK_NonceValues* nonce_values,
                                              ODK_ClockValues* clock_values,
                                              uint64_t system_time_seconds) {
  if (core_message_size == NULL || nonce_values == NULL ||
      clock_values == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  /* If the license has not been loaded, then this is release instead of a
   * renewal. All releases use v15. */
  if (clock_values->timer_status == ODK_CLOCK_TIMER_STATUS_LICENSE_NOT_LOADED ||
      clock_values->timer_status == ODK_CLOCK_TIMER_STATUS_LICENSE_INACTIVE) {
    nonce_values->api_major_version = ODK_FIRST_VERSION - 1;
  }
  if (nonce_values->api_major_version < ODK_FIRST_VERSION) {
    *core_message_size = 0;
    return OEMCrypto_SUCCESS;
  }

  ODK_PreparedRenewalRequest renewal_request = {0};
  /* First, we compute the time this request was made relative to the playback
   * clock. */
  if (clock_values->time_of_first_decrypt == 0) {
    /* It is OK to preemptively request a renewal before playback starts.
     * We'll treat this as asking for a renewal at playback time 0. */
    renewal_request.playback_time = 0;
  } else {
    /* Otherwise, playback_time is relative to the first decrypt. */
    if (odk_sub_overflow_u64(system_time_seconds,
                             clock_values->time_of_first_decrypt,
                             &renewal_request.playback_time)) {
      return ODK_ERROR_CORE_MESSAGE;
    }
  }

  /* Save time for this request so that we can verify the response. This makes
   * all earlier requests invalid. If preparing this request fails, then all
   * requests will be bad. */
  clock_values->time_of_renewal_request = renewal_request.playback_time;

  return ODK_PrepareRequest(
      message, message_length, core_message_size, ODK_Renewal_Request_Type,
      nonce_values, &renewal_request, sizeof(ODK_PreparedRenewalRequest));
}

OEMCryptoResult ODK_PrepareCoreProvisioningRequest(
    uint8_t* message, size_t message_length, size_t* core_message_length,
    const ODK_NonceValues* nonce_values,
    const ODK_MessageCounterInfo* counter_info) {
  if (core_message_length == NULL || nonce_values == NULL ||
      counter_info == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_PreparedProvisioningRequest provisioning_request = {0};
  memcpy(&provisioning_request.counter_info, counter_info,
         sizeof(ODK_MessageCounterInfo));

  return ODK_PrepareRequest(message, message_length, core_message_length,
                            ODK_Provisioning_Request_Type, nonce_values,
                            &provisioning_request,
                            sizeof(ODK_PreparedProvisioningRequest));
}

OEMCryptoResult ODK_PrepareCoreProvisioning40Request(
    uint8_t* message, size_t message_length, size_t* core_message_length,
    const ODK_NonceValues* nonce_values, const uint8_t* device_info,
    size_t device_info_length, const ODK_MessageCounterInfo* counter_info) {
  if (core_message_length == NULL || nonce_values == NULL ||
      counter_info == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_PreparedProvisioning40Request provisioning_request = {0};
  if (device_info_length > sizeof(provisioning_request.device_info)) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  provisioning_request.device_info_length = (uint32_t)device_info_length;
  if (device_info) {
    memcpy(provisioning_request.device_info, device_info, device_info_length);
  }
  memcpy(&provisioning_request.counter_info, counter_info,
         sizeof(provisioning_request.counter_info));

  return ODK_PrepareRequest(message, message_length, core_message_length,
                            ODK_Provisioning40_Request_Type, nonce_values,
                            &provisioning_request,
                            sizeof(provisioning_request));
}

OEMCryptoResult ODK_PrepareCoreRenewedProvisioningRequest(
    uint8_t* message, size_t message_length, size_t* core_message_length,
    const ODK_NonceValues* nonce_values, const uint8_t* device_id,
    size_t device_id_length, uint16_t renewal_type, const uint8_t* renewal_data,
    size_t renewal_data_length) {
  if (core_message_length == NULL || nonce_values == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_PreparedRenewedProvisioningRequest provisioning_request = {0};
  if (device_id_length > sizeof(provisioning_request.device_id)) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  provisioning_request.device_id_length = (uint32_t)device_id_length;
  if (device_id) {
    memcpy(provisioning_request.device_id, device_id, device_id_length);
  }
  if (renewal_data_length > sizeof(provisioning_request.renewal_data)) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  provisioning_request.renewal_type = renewal_type;
  provisioning_request.renewal_data_length = (uint32_t)renewal_data_length;
  if (renewal_data) {
    memcpy(provisioning_request.renewal_data, renewal_data,
           renewal_data_length);
  }
  return ODK_PrepareRequest(message, message_length, core_message_length,
                            ODK_Renewed_Provisioning_Request_Type, nonce_values,
                            &provisioning_request,
                            sizeof(provisioning_request));
}

/* @@ parse response functions */

OEMCryptoResult ODK_ParseLicense(
    const uint8_t* message, size_t message_length, size_t core_message_length,
    bool initial_license_load, bool usage_entry_present,
    uint64_t system_time_seconds, ODK_TimerLimits* timer_limits,
    ODK_ClockValues* clock_values, ODK_NonceValues* nonce_values,
    ODK_ParsedLicense* parsed_license, uint64_t* timer_value) {
  if (message == NULL || timer_limits == NULL || clock_values == NULL ||
      nonce_values == NULL || parsed_license == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  OEMCryptoResult err =
      ODK_ParseCoreHeader(message, message_length, core_message_length,
                          ODK_License_Response_Type, nonce_values);
  if (err != OEMCrypto_SUCCESS) {
    return err;
  }

  ODK_LicenseResponse license_response = {0};
  license_response.parsed_license = parsed_license;

  ODK_Message msg = ODK_Message_Create((uint8_t*)message, message_length);

  ODK_Message_SetSize(&msg, core_message_length);

  Unpack_ODK_LicenseResponse(&msg, &license_response);

  if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
      ODK_Message_GetOffset(&msg) != core_message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  /* If the license has a provider session token (pst), then OEMCrypto should
   * have a usage entry loaded. The opposite is also an error.  */
  if ((usage_entry_present && parsed_license->pst.length == 0) ||
      (!usage_entry_present && parsed_license->pst.length > 0)) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  /* If this is the first time we load this license, then we verify that the
   * nonce values are the correct, otherwise we copy the nonce values. If the
   * nonce values are not required to be correct, then we don't know if this is
   * an initial load or not. In that case, we also copy the values so that we
   * can use the nonce values later for a renewal.
   */
  if (parsed_license->nonce_required && initial_license_load) {
    if (nonce_values->nonce !=
            license_response.core_message.nonce_values.nonce ||
        nonce_values->session_id !=
            license_response.core_message.nonce_values.session_id) {
      return OEMCrypto_ERROR_INVALID_NONCE;
    }
  } else { /* !initial_license_load, or can't tell if initial. */
    nonce_values->nonce = license_response.core_message.nonce_values.nonce;
    nonce_values->session_id =
        license_response.core_message.nonce_values.session_id;
    /* Start the rental clock if not already started for reloading an offline
     * license without a nonce.  */
    if (!parsed_license->nonce_required &&
        clock_values->time_of_license_request_signed == 0) {
      clock_values->time_of_license_request_signed = system_time_seconds;
    }
  }
  bool license_load =
      (parsed_license->renewal_delay_base == OEMCrypto_License_Load);
  *timer_limits = parsed_license->timer_limits;
  /* And update the clock values state. */
  clock_values->timer_status = ODK_CLOCK_TIMER_STATUS_LICENSE_LOADED;
  if (nonce_values->api_major_version == 18 && license_load) {
    err = ODK_AttemptFirstPlayback(system_time_seconds, timer_limits,
                                   clock_values, timer_value);
    return err;
  }
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult ODK_ParseRenewal(const uint8_t* message, size_t message_length,
                                 size_t core_message_length,
                                 ODK_NonceValues* nonce_values,
                                 uint64_t system_time,
                                 const ODK_TimerLimits* timer_limits,
                                 ODK_ClockValues* clock_values,
                                 uint64_t* timer_value) {
  if (message == NULL || nonce_values == NULL || timer_limits == NULL ||
      clock_values == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  const OEMCryptoResult err =
      ODK_ParseCoreHeader(message, message_length, core_message_length,
                          ODK_Renewal_Response_Type, nonce_values);
  if (err != OEMCrypto_SUCCESS) {
    return err;
  }
  ODK_RenewalResponse renewal_response = {0};
  ODK_Message msg = ODK_Message_Create((uint8_t*)message, message_length);
  ODK_Message_SetSize(&msg, core_message_length);
  Unpack_ODK_RenewalResponse(&msg, &renewal_response);

  if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
      ODK_Message_GetOffset(&msg) != core_message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  /* always verify nonce_values for Renewal and Provisioning responses */
  if (!ODK_NonceValuesEqualExcludingVersion(
          nonce_values,
          &(renewal_response.request.core_message.nonce_values))) {
    return OEMCrypto_ERROR_INVALID_NONCE;
  }

  /* Reference:
   *   Doc: License Duration and Renewal (Changes for OEMCrypto v16)
   *     Section: Renewal Message
   */
  /* If a renewal request is lost in transit, we should throw it out and create
   * a new one. We use the timestamp to make sure we have the latest request.
   * We only do this if playback has already started. This allows us to reload
   * an offline license and also reload a renewal before starting playback.
   */
  if (clock_values->timer_status != ODK_CLOCK_TIMER_STATUS_LICENSE_LOADED &&
      clock_values->time_of_renewal_request <
          renewal_response.request.playback_time) {
    return ODK_STALE_RENEWAL;
  }
  return ODK_ComputeRenewalDuration(timer_limits, clock_values, system_time,
                                    renewal_response.renewal_duration_seconds,
                                    timer_value);
}

OEMCryptoResult ODK_ParseProvisioning(
    const uint8_t* message, size_t message_length, size_t core_message_length,
    ODK_NonceValues* nonce_values, const uint8_t* device_id,
    size_t device_id_length, ODK_ParsedProvisioning* parsed_response) {
  if (message == NULL || nonce_values == NULL || device_id == NULL ||
      parsed_response == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  const OEMCryptoResult err =
      ODK_ParseCoreHeader(message, message_length, core_message_length,
                          ODK_Provisioning_Response_Type, nonce_values);
  if (err != OEMCrypto_SUCCESS) {
    return err;
  }

  if (nonce_values->api_major_version <= 17) {
    // Do v16/v17
    ODK_ProvisioningResponseV16 provisioning_response = {0};
    provisioning_response.parsed_provisioning = parsed_response;

    if (device_id_length > ODK_DEVICE_ID_LEN_MAX) {
      return ODK_ERROR_CORE_MESSAGE;
    }

    ODK_Message msg = ODK_Message_Create((uint8_t*)message, message_length);
    ODK_Message_SetSize(&msg, core_message_length);
    Unpack_ODK_ProvisioningResponseV16(&msg, &provisioning_response);
    if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
        ODK_Message_GetOffset(&msg) != core_message_length) {
      return ODK_ERROR_CORE_MESSAGE;
    }
    /* always verify nonce_values for Renewal and Provisioning responses */
    if (!ODK_NonceValuesEqualExcludingVersion(
            nonce_values,
            &(provisioning_response.request.core_message.nonce_values))) {
      return OEMCrypto_ERROR_INVALID_NONCE;
    }

    if (crypto_memcmp(device_id, provisioning_response.request.device_id,
                      device_id_length) != 0) {
      return ODK_ERROR_CORE_MESSAGE;
    }

    const uint8_t zero[ODK_DEVICE_ID_LEN_MAX] = {0};
    /* check bytes beyond device_id_length are 0 */
    if (crypto_memcmp(
            zero, provisioning_response.request.device_id + device_id_length,
            ODK_DEVICE_ID_LEN_MAX - device_id_length) != 0) {
      return ODK_ERROR_CORE_MESSAGE;
    }
  } else {
    // v18
    ODK_ProvisioningResponse provisioning_response = {0};
    provisioning_response.parsed_provisioning = parsed_response;

    ODK_Message msg = ODK_Message_Create((uint8_t*)message, message_length);
    ODK_Message_SetSize(&msg, core_message_length);
    Unpack_ODK_ProvisioningResponse(&msg, &provisioning_response);
    if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
        ODK_Message_GetOffset(&msg) != core_message_length) {
      return ODK_ERROR_CORE_MESSAGE;
    }
    /* always verify nonce_values for Renewal and Provisioning responses */
    if (!ODK_NonceValuesEqualExcludingVersion(
            nonce_values, &(provisioning_response.core_message.nonce_values))) {
      return OEMCrypto_ERROR_INVALID_NONCE;
    }
  }

  return OEMCrypto_SUCCESS;
}

OEMCryptoResult ODK_ParseProvisioning40(const uint8_t* message,
                                        size_t message_length,
                                        size_t core_message_length,
                                        ODK_NonceValues* nonce_values) {
  if (message == NULL || nonce_values == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  const OEMCryptoResult err =
      ODK_ParseCoreHeader(message, message_length, core_message_length,
                          ODK_Provisioning_Response_Type, nonce_values);
  if (err != OEMCrypto_SUCCESS) {
    return err;
  }
  ODK_Provisioning40Response provisioning_response = {0};

  ODK_Message msg = ODK_Message_Create((uint8_t*)message, message_length);
  ODK_Message_SetSize(&msg, core_message_length);
  Unpack_ODK_Provisioning40Response(&msg, &provisioning_response);
  if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
      ODK_Message_GetOffset(&msg) != core_message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  /* always verify nonce_values for Renewal and Provisioning responses */
  if (!ODK_NonceValuesEqualExcludingVersion(
          nonce_values, &(provisioning_response.core_message.nonce_values))) {
    return OEMCrypto_ERROR_INVALID_NONCE;
  }

  return OEMCrypto_SUCCESS;
}

bool CheckApiVersionAtMost(const ODK_NonceValues* nonce_values,
                           uint16_t major_version, uint16_t minor_version) {
  return nonce_values->api_major_version < major_version ||
         (nonce_values->api_major_version == major_version &&
          nonce_values->api_minor_version <= minor_version);
}