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

#include <stddef.h>
#include <stdint.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"
#include "serialization_base.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 (sizeof(ODK_CoreMessage) > 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;
    }
    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;
}

static OEMCryptoResult ODK_ParseResponse(
    const uint8_t* message, size_t message_length, size_t core_message_length,
    ODK_MessageType message_type, const ODK_NonceValues* nonce_values,
    void* response_buffer, uint32_t response_buffer_length) {
  if (message == NULL || response_buffer == NULL ||
      core_message_length > message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }

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

  /* The core message should be at the beginning of the buffer, and with a
   * shorter length. The core message is the part we are parsing. */
  ODK_Message_SetSize(&msg, core_message_length);

  /* Parse message and unpack it into response buffer. */
  switch (message_type) {
    case ODK_License_Response_Type: {
      if (sizeof(ODK_LicenseResponse) > response_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Unpack_ODK_LicenseResponse(&msg, (ODK_LicenseResponse*)response_buffer);
      break;
    }
    case ODK_Renewal_Response_Type: {
      if (sizeof(ODK_RenewalResponse) > response_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Unpack_ODK_RenewalResponse(&msg, (ODK_RenewalResponse*)response_buffer);
      break;
    }
    case ODK_Provisioning_Response_Type: {
      if (sizeof(ODK_ProvisioningResponse) > response_buffer_length) {
        return ODK_ERROR_CORE_MESSAGE;
      }
      Unpack_ODK_ProvisioningResponse(
          &msg, (ODK_ProvisioningResponse*)response_buffer);
      break;
    }
    default: {
      return ODK_ERROR_CORE_MESSAGE;
    }
  }

  ODK_CoreMessage* core_message = (ODK_CoreMessage*)response_buffer;
  if (ODK_Message_GetStatus(&msg) != MESSAGE_STATUS_OK ||
      message_type != core_message->message_type ||
      ODK_Message_GetOffset(&msg) != core_message->message_length) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  if (nonce_values) {
    /* always verify nonce_values for Renewal and Provisioning responses */
    if (!ODK_NonceValuesEqual(nonce_values, &(core_message->nonce_values))) {
      return OEMCrypto_ERROR_INVALID_NONCE;
    }
  }

  return OEMCrypto_SUCCESS;
}

/* @ public odk functions */

/* @@ prepare request functions */

OEMCryptoResult ODK_PrepareCoreLicenseRequest(
    uint8_t* message, size_t message_length, size_t* core_message_length,
    const ODK_NonceValues* nonce_values) {
  if (core_message_length == NULL || nonce_values == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_PreparedLicenseRequest license_request = {
      {0, 0, {}},
  };
  return ODK_PrepareRequest(
      message, message_length, core_message_length, 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, 0, {}}, 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 uint8_t* device_id,
    size_t device_id_length) {
  if (core_message_length == NULL || nonce_values == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  ODK_PreparedProvisioningRequest provisioning_request = {
      {0, 0, {}},
      0,
      {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);
  }
  return ODK_PrepareRequest(message, message_length, core_message_length,
                            ODK_Provisioning_Request_Type, nonce_values,
                            &provisioning_request,
                            sizeof(ODK_PreparedProvisioningRequest));
}

/* @@ 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,
    const uint8_t* request_hash, ODK_TimerLimits* timer_limits,
    ODK_ClockValues* clock_values, ODK_NonceValues* nonce_values,
    ODK_ParsedLicense* parsed_license) {
  if (message == NULL || request_hash == NULL || timer_limits == NULL ||
      clock_values == NULL || nonce_values == NULL || parsed_license == NULL) {
    return ODK_ERROR_CORE_MESSAGE;
  }

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

  const OEMCryptoResult err = ODK_ParseResponse(
      message, message_length, core_message_length, ODK_License_Response_Type,
      NULL, &license_response, sizeof(ODK_LicenseResponse));

  if (err != OEMCrypto_SUCCESS) {
    return err;
  }

  /* We do not support future API version. Also, this function should not be
   * used for legacy licenses. */
  if (license_response.request.core_message.nonce_values.api_major_version >
          ODK_MAJOR_VERSION ||
      license_response.request.core_message.nonce_values.api_major_version <
          ODK_FIRST_VERSION) {
    return ODK_UNSUPPORTED_API;
  }

  /* If the server sent us an older format, record the license's API version. */
  if (nonce_values->api_major_version >
      license_response.request.core_message.nonce_values.api_major_version) {
    nonce_values->api_major_version =
        license_response.request.core_message.nonce_values.api_major_version;
    nonce_values->api_minor_version =
        license_response.request.core_message.nonce_values.api_minor_version;
  } else if (nonce_values->api_minor_version >
             license_response.request.core_message.nonce_values
                 .api_minor_version) {
    nonce_values->api_minor_version =
        license_response.request.core_message.nonce_values.api_minor_version;
  }
  /* 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.request.core_message.nonce_values.nonce ||
        nonce_values->session_id !=
            license_response.request.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.request.core_message.nonce_values.nonce;
    nonce_values->session_id =
        license_response.request.core_message.nonce_values.session_id;
  }
  /* For v16, in order to be backwards compatible with a v15 license server,
   * OEMCrypto stores a hash of the core license request and only signs the
   * message body. Here, when we process the license response, we verify that
   * the server has the same hash of the core request.  */
  if (initial_license_load && parsed_license->nonce_required &&
      crypto_memcmp(request_hash, license_response.request_hash,
                    ODK_SHA256_HASH_SIZE)) {
    return ODK_ERROR_CORE_MESSAGE;
  }
  *timer_limits = parsed_license->timer_limits;
  /* And update the clock values state. */
  clock_values->timer_status = ODK_CLOCK_TIMER_STATUS_LICENSE_LOADED;
  return OEMCrypto_SUCCESS;
}

OEMCryptoResult ODK_ParseRenewal(const uint8_t* message, size_t message_length,
                                 size_t core_message_length,
                                 const 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;
  }

  ODK_RenewalResponse renewal_response = {
      {{0, 0, {}}, 0},
      0,
  };
  const OEMCryptoResult err = ODK_ParseResponse(
      message, message_length, core_message_length, ODK_Renewal_Response_Type,
      nonce_values, &renewal_response, sizeof(ODK_RenewalResponse));

  if (err != OEMCrypto_SUCCESS) {
    return err;
  }

  /* 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,
    const 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;
  }

  ODK_ProvisioningResponse provisioning_response = {{{0, 0, {}}, 0, {0}}, NULL};
  provisioning_response.parsed_provisioning = parsed_response;

  if (device_id_length > ODK_DEVICE_ID_LEN_MAX) {
    return ODK_ERROR_CORE_MESSAGE;
  }

  const OEMCryptoResult err = ODK_ParseResponse(
      message, message_length, core_message_length,
      ODK_Provisioning_Response_Type, nonce_values, &provisioning_response,
      sizeof(ODK_ProvisioningResponse));

  if (err != OEMCrypto_SUCCESS) {
    return err;
  }

  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;
  }

  return OEMCrypto_SUCCESS;
}