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

#include <optional>
#include <vector>
#include "absl/log/log.h"

#include "cuttlefish/host/commands/secure_env/tpm_auth.h"
#include "cuttlefish/host/commands/secure_env/tpm_hmac.h"

namespace cuttlefish {

HmacSerializable::HmacSerializable(
    TpmResourceManager& resource_manager,
    std::function<TpmObjectSlot(TpmResourceManager&)> signing_key_fn,
    uint32_t digest_size, Serializable* wrapped, const Serializable* aad)
    : resource_manager_(resource_manager),
      signing_key_fn_(signing_key_fn),
      digest_size_(digest_size),
      wrapped_(wrapped),
      aad_(aad) {}

size_t HmacSerializable::SerializedSize() const {
  auto digest_size = sizeof(uint32_t) + digest_size_;
  auto data_size = sizeof(uint32_t) + wrapped_->SerializedSize();
  return digest_size + data_size;
}

uint8_t* HmacSerializable::Serialize(uint8_t* buf, const uint8_t* end) const {
  auto wrapped_size = wrapped_->SerializedSize();
  buf = keymaster::append_uint32_to_buf(buf, end, wrapped_size);
  auto signed_data = buf;
  buf = wrapped_->Serialize(buf, end);
  if (buf - signed_data != wrapped_size) {
    LOG(ERROR) << "Serialized wrapped data did not match expected size.";
    return buf;
  }
  auto key = signing_key_fn_(resource_manager_);
  if (!key) {
    LOG(ERROR) << "Could not retrieve key";
    return buf;
  }
  auto maced_data = AppendAad(signed_data, wrapped_size);
  if (!maced_data) {
    return buf;
  }
  auto hmac_data =
      TpmHmac(resource_manager_, key->get(), TpmAuth(ESYS_TR_PASSWORD),
              maced_data->data(), maced_data->size());
  if (!hmac_data) {
    LOG(ERROR) << "Failed to produce hmac";
    return buf;
  }
  if (hmac_data->size != digest_size_) {
    LOG(ERROR) << "Unexpected digest size. Wanted " << digest_size_
               << ", TPM produced " << hmac_data->size;
    return buf;
  }
  return keymaster::append_size_and_data_to_buf(
      buf, end, hmac_data->buffer, digest_size_);
}

bool HmacSerializable::Deserialize(const uint8_t** buf_ptr, const uint8_t* end) {
  size_t signed_data_size;
  keymaster::UniquePtr<uint8_t[]> signed_data;
  bool success =
      keymaster::copy_size_and_data_from_buf(
          buf_ptr, end, &signed_data_size, &signed_data);
  if (!success) {
    LOG(ERROR) << "Failed to retrieve signed data";
    return false;
  }
  size_t signature_size;
  keymaster::UniquePtr<uint8_t[]> signature;
  success =
      keymaster::copy_size_and_data_from_buf(
          buf_ptr, end, &signature_size, &signature);
  if (!success) {
    LOG(ERROR) << "Failed to retrieve signature";
    return false;
  }
  if (signature_size != digest_size_) {
    LOG(ERROR) << "Digest size did not match expected size.";
    return false;
  }
  auto key = signing_key_fn_(resource_manager_);
  if (!key) {
    LOG(ERROR) << "Could not retrieve key";
    return false;
  }
  auto maced_data = AppendAad(signed_data.get(), signed_data_size);
  if (!maced_data) {
    return false;
  }
  auto hmac_check =
      TpmHmac(resource_manager_, key->get(), TpmAuth(ESYS_TR_PASSWORD),
              maced_data->data(), maced_data->size());
  if (!hmac_check) {
    LOG(ERROR) << "Unable to calculate signature check";
    return false;
  }
  if (hmac_check->size != digest_size_) {
    LOG(ERROR) << "Unexpected signature check size. Wanted " << digest_size_
               << ", TPM produced " << hmac_check->size;
    return false;
  }
  if (memcmp(signature.get(), hmac_check->buffer, digest_size_) != 0) {
    LOG(ERROR) << "Signature check did not match original signature.";
    return false;
  }
  // Now that we've validated integrity on the data, do the inner deserialization
  auto inner_buf = signed_data.get();
  auto inner_buf_end = inner_buf + signed_data_size;
  return wrapped_->Deserialize(
      const_cast<const uint8_t**>(&inner_buf), inner_buf_end);
}

std::optional<std::vector<uint8_t>> HmacSerializable::AppendAad(
    const uint8_t* sensitive, size_t sensitive_size) const {
  if (!aad_) {
    return std::vector<uint8_t>(sensitive, sensitive + sensitive_size);
  }
  std::vector<uint8_t> output(sensitive_size + aad_->SerializedSize());
  std::copy(sensitive, sensitive + sensitive_size, output.begin());

  const uint8_t* actual_output_end =
      aad_->Serialize(&output[sensitive_size], output.data() + output.size());
  const ptrdiff_t actual_aad_size = actual_output_end - output.data();
  if (actual_aad_size != output.size()) {
    LOG(ERROR) << "Serialized aad did not match expected size. Expected: "
               << output.size() << ", actual: " << actual_aad_size;
    return std::nullopt;
  }
  return output;
}

}  // namespace cuttlefish
