// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <memory>
#include <utility>

#include "chaps/tpm_thread_utility_impl.h"

#include <base/bind.h>
#include <base/synchronization/waitable_event.h>
#include <base/task/bind_post_task.h>
#include <base/threading/thread.h>
#include <base/threading/thread_task_runner_handle.h>

namespace {
constexpr char kTPMThreadName[] = "tpm_thread";
}

namespace chaps {

TPMThreadUtilityImpl::TPMThread::TPMThread(const std::string& name,
                                           TPMThreadUtilityImpl* utility)
    : base::Thread(name), utility_(utility) {
  CHECK(utility_);
}

TPMThreadUtilityImpl::TPMThread::~TPMThread() {
  Stop();
}

void TPMThreadUtilityImpl::TPMThread::CleanUp() {
  utility_->inner_tpm_.reset();
}

TPMThreadUtilityImpl::TPMThreadUtilityImpl(
    std::unique_ptr<TPMUtility> inner_tpm)
    : inner_tpm_(std::move(inner_tpm)),
      tpm_thread_(std::make_unique<TPMThread>(kTPMThreadName, this)) {
  base::Thread::Options options;
  options.message_pump_type = base::MessagePumpType::IO;
  tpm_thread_->StartWithOptions(std::move(options));
  task_runner_ = tpm_thread_->task_runner();
}

TPMThreadUtilityImpl::~TPMThreadUtilityImpl() {
  // The |inner_tpm_| would be destructed on |tpm_thread_|.
  tpm_thread_.reset();
}

template <typename MethodType, typename ResultType, typename... Args>
void TPMThreadUtilityImpl::SendRequestAndWaitResult(const MethodType& method,
                                                    ResultType* result,
                                                    Args&&... args) {
  base::WaitableEvent event(base::WaitableEvent::ResetPolicy::MANUAL,
                            base::WaitableEvent::InitialState::NOT_SIGNALED);
  base::OnceClosure task =
      base::BindOnce(method, base::Unretained(inner_tpm_.get()),
                     std::forward<Args>(args)...)
          .Then(base::BindOnce(
              [](ResultType* result, ResultType value) {
                *result = std::move(value);
              },
              result))
          .Then(base::BindOnce(&base::WaitableEvent::Signal,
                               base::Unretained(&event)));
  task_runner_->PostTask(FROM_HERE, std::move(task));
  event.Wait();
}

template <typename MethodType, typename... Args>
void TPMThreadUtilityImpl::SendRequestAndWait(const MethodType& method,
                                              Args&&... args) {
  base::WaitableEvent event(base::WaitableEvent::ResetPolicy::MANUAL,
                            base::WaitableEvent::InitialState::NOT_SIGNALED);
  base::OnceClosure task =
      base::BindOnce(method, base::Unretained(inner_tpm_.get()),
                     std::forward<Args>(args)...)
          .Then(base::BindOnce(&base::WaitableEvent::Signal,
                               base::Unretained(&event)));
  task_runner_->PostTask(FROM_HERE, std::move(task));
  event.Wait();
}

TPMVersion TPMThreadUtilityImpl::GetTPMVersion() {
  // This simple function doesn't need to be called on the TPM thread.
  return inner_tpm_->GetTPMVersion();
}

size_t TPMThreadUtilityImpl::MinRSAKeyBits() {
  // This simple function doesn't need to be called on the TPM thread.
  return inner_tpm_->MinRSAKeyBits();
}

size_t TPMThreadUtilityImpl::MaxRSAKeyBits() {
  // This simple function doesn't need to be called on the TPM thread.
  return inner_tpm_->MaxRSAKeyBits();
}

bool TPMThreadUtilityImpl::Init() {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::Init, &result);
  return result;
}

bool TPMThreadUtilityImpl::IsTPMAvailable() {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::IsTPMAvailable, &result);
  return result;
}

bool TPMThreadUtilityImpl::Authenticate(const brillo::SecureBlob& auth_data,
                                        const std::string& auth_key_blob,
                                        const std::string& encrypted_root_key,
                                        brillo::SecureBlob* root_key) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::Authenticate, &result, auth_data,
                           auth_key_blob, encrypted_root_key, root_key);
  return result;
}

bool TPMThreadUtilityImpl::ChangeAuthData(
    const brillo::SecureBlob& old_auth_data,
    const brillo::SecureBlob& new_auth_data,
    const std::string& old_auth_key_blob,
    std::string* new_auth_key_blob) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::ChangeAuthData, &result, old_auth_data,
                           new_auth_data, old_auth_key_blob, new_auth_key_blob);
  return result;
}

bool TPMThreadUtilityImpl::GenerateRandom(int num_bytes,
                                          std::string* random_data) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::GenerateRandom, &result, num_bytes,
                           random_data);
  return result;
}

bool TPMThreadUtilityImpl::StirRandom(const std::string& entropy_data) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::StirRandom, &result, entropy_data);
  return result;
}

bool TPMThreadUtilityImpl::GenerateRSAKey(int slot,
                                          int modulus_bits,
                                          const std::string& public_exponent,
                                          const brillo::SecureBlob& auth_data,
                                          std::string* key_blob,
                                          int* key_handle) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::GenerateRSAKey, &result, slot,
                           modulus_bits, public_exponent, auth_data, key_blob,
                           key_handle);
  return result;
}

bool TPMThreadUtilityImpl::GetRSAPublicKey(int key_handle,
                                           std::string* public_exponent,
                                           std::string* modulus) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::GetRSAPublicKey, &result, key_handle,
                           public_exponent, modulus);
  return result;
}

bool TPMThreadUtilityImpl::IsECCurveSupported(int curve_nid) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::IsECCurveSupported, &result, curve_nid);
  return result;
}

bool TPMThreadUtilityImpl::GenerateECCKey(int slot,
                                          int nid,
                                          const brillo::SecureBlob& auth_data,
                                          std::string* key_blob,
                                          int* key_handle) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::GenerateECCKey, &result, slot, nid,
                           auth_data, key_blob, key_handle);
  return result;
}

bool TPMThreadUtilityImpl::GetECCPublicKey(int key_handle,
                                           std::string* public_point) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::GetECCPublicKey, &result, key_handle,
                           public_point);
  return result;
}

bool TPMThreadUtilityImpl::WrapRSAKey(int slot,
                                      const std::string& public_exponent,
                                      const std::string& modulus,
                                      const std::string& prime_factor,
                                      const brillo::SecureBlob& auth_data,
                                      std::string* key_blob,
                                      int* key_handle) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::WrapRSAKey, &result, slot,
                           public_exponent, modulus, prime_factor, auth_data,
                           key_blob, key_handle);
  return result;
}

bool TPMThreadUtilityImpl::WrapECCKey(int slot,
                                      int curve_nid,
                                      const std::string& public_point_x,
                                      const std::string& public_point_y,
                                      const std::string& private_value,
                                      const brillo::SecureBlob& auth_data,
                                      std::string* key_blob,
                                      int* key_handle) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::WrapECCKey, &result, slot, curve_nid,
                           public_point_x, public_point_y, private_value,
                           auth_data, key_blob, key_handle);
  return result;
}

bool TPMThreadUtilityImpl::LoadKey(int slot,
                                   const std::string& key_blob,
                                   const brillo::SecureBlob& auth_data,
                                   int* key_handle) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::LoadKey, &result, slot, key_blob,
                           auth_data, key_handle);
  return result;
}

bool TPMThreadUtilityImpl::LoadKeyWithParent(
    int slot,
    const std::string& key_blob,
    const brillo::SecureBlob& auth_data,
    int parent_key_handle,
    int* key_handle) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::LoadKeyWithParent, &result, slot,
                           key_blob, auth_data, parent_key_handle, key_handle);
  return result;
}

void TPMThreadUtilityImpl::UnloadKeysForSlot(int slot) {
  SendRequestAndWait(&TPMUtility::UnloadKeysForSlot, slot);
}

bool TPMThreadUtilityImpl::Bind(int key_handle,
                                const std::string& input,
                                std::string* output) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::Bind, &result, key_handle, input,
                           output);
  return result;
}

bool TPMThreadUtilityImpl::Unbind(int key_handle,
                                  const std::string& input,
                                  std::string* output) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::Unbind, &result, key_handle, input,
                           output);
  return result;
}

bool TPMThreadUtilityImpl::Sign(int key_handle,
                                CK_MECHANISM_TYPE signing_mechanism,
                                const std::string& mechanism_parameter,
                                const std::string& input,
                                std::string* signature) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::Sign, &result, key_handle,
                           signing_mechanism, mechanism_parameter, input,
                           signature);
  return result;
}

bool TPMThreadUtilityImpl::IsSRKReady() {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::IsSRKReady, &result);
  return result;
}

bool TPMThreadUtilityImpl::SealData(const std::string& unsealed_data,
                                    const brillo::SecureBlob& auth_value,
                                    std::string* key_blob,
                                    std::string* encrypted_data) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::SealData, &result, unsealed_data,
                           auth_value, key_blob, encrypted_data);
  return result;
}

bool TPMThreadUtilityImpl::UnsealData(const std::string& key_blob,
                                      const std::string& encrypted_data,
                                      const brillo::SecureBlob& auth_value,
                                      brillo::SecureBlob* unsealed_data) {
  bool result;
  SendRequestAndWaitResult(&TPMUtility::UnsealData, &result, key_blob,
                           encrypted_data, auth_value, unsealed_data);
  return result;
}

void TPMThreadUtilityImpl::GenerateRandomAsync(
    int num_bytes, GenerateRandomCallback callback) {
  CHECK(base::SequencedTaskRunnerHandle::IsSet())
      << "Caller doesn't have task runner.";

  auto random_data = std::make_unique<std::string>();

  base::OnceCallback<bool(void)> task = base::BindOnce(
      &TPMUtility::GenerateRandom, base::Unretained(inner_tpm_.get()),
      num_bytes, random_data.get());

  base::OnceCallback<void(bool)> run_callback = base::BindOnce(
      [](GenerateRandomCallback callback,
         std::unique_ptr<std::string> random_data,
         bool result) { std::move(callback).Run(result, *random_data); },
      std::move(callback), std::move(random_data));

  task_runner_->PostTaskAndReplyWithResult(FROM_HERE, std::move(task),
                                           std::move(run_callback));
}

void TPMThreadUtilityImpl::UnloadKeysForSlotAsync(
    int slot, UnloadKeysForSlotCallback callback) {
  CHECK(base::SequencedTaskRunnerHandle::IsSet())
      << "Caller doesn't have task runner.";

  base::OnceClosure task =
      base::BindOnce(&TPMUtility::UnloadKeysForSlot,
                     base::Unretained(inner_tpm_.get()), slot)
          .Then(base::BindPostTask(base::SequencedTaskRunnerHandle::Get(),
                                   std::move(callback)));

  task_runner_->PostTask(FROM_HERE, std::move(task));
}

void TPMThreadUtilityImpl::SealDataAsync(const std::string& unsealed_data,
                                         const brillo::SecureBlob& auth_value,
                                         SealDataCallback callback) {
  CHECK(base::SequencedTaskRunnerHandle::IsSet())
      << "Caller doesn't have task runner.";

  auto key_blob = std::make_unique<std::string>();
  auto encrypted_data = std::make_unique<std::string>();

  base::OnceCallback<bool(void)> task = base::BindOnce(
      &TPMUtility::SealData, base::Unretained(inner_tpm_.get()), unsealed_data,
      auth_value, key_blob.get(), encrypted_data.get());

  base::OnceCallback<void(bool)> run_callback = base::BindOnce(
      [](SealDataCallback callback, std::unique_ptr<std::string> key_blob,
         std::unique_ptr<std::string> encrypted_data, bool result) {
        std::move(callback).Run(result, *key_blob, *encrypted_data);
      },
      std::move(callback), std::move(key_blob), std::move(encrypted_data));

  task_runner_->PostTaskAndReplyWithResult(FROM_HERE, std::move(task),
                                           std::move(run_callback));
}

void TPMThreadUtilityImpl::UnsealDataAsync(const std::string& key_blob,
                                           const std::string& encrypted_data,
                                           const brillo::SecureBlob& auth_value,
                                           UnsealDataCallback callback) {
  CHECK(base::SequencedTaskRunnerHandle::IsSet())
      << "Caller doesn't have task runner.";

  auto unsealed_data = std::make_unique<brillo::SecureBlob>();

  base::OnceCallback<bool(void)> task = base::BindOnce(
      &TPMUtility::UnsealData, base::Unretained(inner_tpm_.get()), key_blob,
      encrypted_data, auth_value, unsealed_data.get());

  base::OnceCallback<void(bool)> run_callback = base::BindOnce(
      [](UnsealDataCallback callback,
         std::unique_ptr<brillo::SecureBlob> unsealed_data,
         bool result) { std::move(callback).Run(result, *unsealed_data); },
      std::move(callback), std::move(unsealed_data));

  task_runner_->PostTaskAndReplyWithResult(FROM_HERE, std::move(task),
                                           std::move(run_callback));
}

}  // namespace chaps