//
// Copyright (C) 2020 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tpm_hmac.h"

#include <android-base/logging.h>
#include <tss2/tss2_esys.h>
#include <tss2/tss2_rc.h>

#include "cuttlefish/host/commands/secure_env/primary_key_builder.h"
#include "cuttlefish/host/commands/secure_env/tpm_resource_manager.h"

namespace cuttlefish {

/* For data large enough to fit in a single TPM2_HMAC call. */
static UniqueEsysPtr<TPM2B_DIGEST> OneshotHmac(
    TpmResourceManager& resource_manager,
    ESYS_TR key_handle,
    TpmAuth auth,
    const uint8_t* data,
    size_t data_size) {
  if (data_size  > TPM2_MAX_DIGEST_BUFFER) {
    LOG(ERROR) << "Logic error: OneshotSign called with data_size "
               << data_size << " (> " << TPM2_MAX_DIGEST_BUFFER << ")";
    return {};
  }
  TPM2B_MAX_BUFFER buffer;
  static_assert(sizeof(buffer.buffer) >= TPM2_MAX_DIGEST_BUFFER);
  buffer.size = data_size;
  memcpy(buffer.buffer, data, data_size);
  TPM2B_DIGEST* out_hmac = nullptr;
  auto rc =
      Esys_HMAC(*resource_manager.Esys(), key_handle, auth.auth1(),
                auth.auth2(), auth.auth3(), &buffer, TPM2_ALG_NULL, &out_hmac);
  if (rc != TPM2_RC_SUCCESS) {
    LOG(ERROR) << "TPM2_HMAC failed: " << Tss2_RC_Decode(rc) << "(" << rc << ")";
    return {};
  }
  if (out_hmac == nullptr) {
    LOG(ERROR) << "out_hmac unset";
    return {};
  }
  return UniqueEsysPtr<TPM2B_DIGEST>(out_hmac);
}

/* For data too large to fit in a single TPM2_HMAC call. */
static UniqueEsysPtr<TPM2B_DIGEST> SegmentedHmac(
    TpmResourceManager& resource_manager,
    ESYS_TR key_handle,
    TpmAuth key_auth,
    const uint8_t* data,
    size_t data_size) {
  // TODO(schuffelen): Pipeline commands where possible.
  TPM2B_AUTH sequence_auth;
  sequence_auth.size = sizeof(rand());
  *reinterpret_cast<decltype(rand())*>(sequence_auth.buffer) = rand();
  ESYS_TR sequence_handle;
  auto slot = resource_manager.ReserveSlot();
  if (!slot) {
    LOG(ERROR) << "No slots available";
    return {};
  }
  auto locked_esys = resource_manager.Esys();
  auto rc = Esys_HMAC_Start(*locked_esys, key_handle, key_auth.auth1(),
                            key_auth.auth2(), key_auth.auth3(), &sequence_auth,
                            TPM2_ALG_NULL, &sequence_handle);
  if (rc != TPM2_RC_SUCCESS) {
    LOG(ERROR) << "TPM2_HMAC_Start failed: " << Tss2_RC_Decode(rc)
               << "(" << rc << ")";
    return {};
  }
  slot->set(sequence_handle);
  rc = Esys_TR_SetAuth(*locked_esys, sequence_handle,
                       &sequence_auth);
  if (rc != TPM2_RC_SUCCESS) {
    LOG(ERROR) << "Esys_TR_SetAuth failed: " << Tss2_RC_Decode(rc)
               << "(" << rc << ")";
    return {};
  }
  auto hashed = 0;
  TPM2B_MAX_BUFFER buffer;
  while (data_size - hashed > TPM2_MAX_DIGEST_BUFFER) {
    buffer.size = TPM2_MAX_DIGEST_BUFFER;
    memcpy(buffer.buffer, &data[hashed], TPM2_MAX_DIGEST_BUFFER);
    hashed += TPM2_MAX_DIGEST_BUFFER;
    rc = Esys_SequenceUpdate(*locked_esys, sequence_handle, ESYS_TR_PASSWORD,
                             ESYS_TR_NONE, ESYS_TR_NONE, &buffer);
    if (rc != TPM2_RC_SUCCESS) {
      LOG(ERROR) << "Esys_SequenceUpdate failed: " << Tss2_RC_Decode(rc)
                << "(" << rc << ")";
      return {};
    }
  }
  buffer.size = data_size - hashed;
  memcpy(buffer.buffer, &data[hashed], buffer.size);
  TPM2B_DIGEST* out_hmac = nullptr;
  TPMT_TK_HASHCHECK* validation = nullptr;
  rc = Esys_SequenceComplete(*locked_esys, sequence_handle, ESYS_TR_PASSWORD,
                             ESYS_TR_NONE, ESYS_TR_NONE, &buffer, TPM2_RH_OWNER,
                             &out_hmac, &validation);
  if (rc != TPM2_RC_SUCCESS) {
    LOG(ERROR) << "Esys_SequenceComplete failed: " << Tss2_RC_Decode(rc)
               << "(" << rc << ")";
    return {};
  }
  // TPM2_SequenceComplete already flushes the sequence context on success.
  slot->set(ESYS_TR_NONE);
  if (out_hmac == nullptr) {
    LOG(ERROR) << "out_hmac was null";
    return {};
  }
  Esys_Free(validation);
  return UniqueEsysPtr<TPM2B_DIGEST>(out_hmac);
}

UniqueEsysPtr<TPM2B_DIGEST> TpmHmac(
    TpmResourceManager& resource_manager,
    ESYS_TR key_handle,
    TpmAuth auth,
    const uint8_t* data,
    size_t data_size) {
  auto fn = data_size > TPM2_MAX_DIGEST_BUFFER ? SegmentedHmac : OneshotHmac;

  return fn(resource_manager, key_handle, auth, data, data_size);
}

UniqueEsysPtr<TPM2B_DIGEST> TpmHmacWithContext(
    TpmResourceManager& resource_manager, const std::string& context,
    const uint8_t* data, size_t data_size) {
  // Use the same context for all HMAC operations (ignoring the provided
  // `context` parameter) to better comply with the KeyMint spec.
  auto key_slot =
      PrimaryKeyBuilder::CreateSigningKey(resource_manager, "TpmHmac_context");
  if (!key_slot) {
    LOG(ERROR) << "Could not make signing key for " << context;
    return nullptr;
  }
  return TpmHmac(resource_manager, key_slot->get(), TpmAuth(ESYS_TR_PASSWORD),
                 data, data_size);
}

}  // namespace cuttlefish