// Copyright 2020 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 "arc/keymaster/context/chaps_crypto_operation.h"

#include <chaps/pkcs11/cryptoki.h>

#include "arc/keymaster/context/chaps_client.h"
#include "arc/keymaster/context/context_adaptor.h"

namespace arc {
namespace keymaster {
namespace context {

constexpr MechanismDescription kCkmRsaPkcsSign = {
    OperationType::kSign, Algorithm::kRsa, Digest::kNone, Padding ::kPkcs1,
    BlockMode::kNone};

constexpr MechanismDescription kCkmSha256RsaPkcsSign = {
    OperationType::kSign, Algorithm::kRsa, Digest::kSha256, Padding ::kPkcs1,
    BlockMode::kNone};

ChapsCryptoOperation::ChapsCryptoOperation(
    base::WeakPtr<ContextAdaptor> context_adaptor,
    const std::string& label,
    const brillo::Blob& id)
    : context_adaptor_(context_adaptor), label_(label), id_(id) {}

ChapsCryptoOperation::~ChapsCryptoOperation() = default;

bool ChapsCryptoOperation::IsSupportedMechanism(
    MechanismDescription description) const {
  return kCkmRsaPkcsSign == description || kCkmSha256RsaPkcsSign == description;
}

base::Optional<uint64_t> ChapsCryptoOperation::Begin(
    MechanismDescription mechanism_description) {
  if (!IsSupportedMechanism(mechanism_description)) {
    LOG(ERROR) << "Mechanism not implemented for chaps keys: " << description();
    return base::nullopt;
  }
  set_description(mechanism_description);

  chaps_ = std::make_unique<ChapsClient>(context_adaptor_);
  base::Optional<CK_OBJECT_HANDLE> handle =
      chaps_->FindObject(CKO_PRIVATE_KEY, label_, id_);
  if (!handle.has_value())
    return base::nullopt;

  CK_OBJECT_HANDLE key_handle = handle.value();

  bool success = false;

  if (description() == kCkmSha256RsaPkcsSign) {
    success = chaps_->InitializeSignature(CKM_SHA256_RSA_PKCS, key_handle);
  } else if (description() == kCkmRsaPkcsSign) {
    success = chaps_->InitializeSignature(CKM_RSA_PKCS, key_handle);
  } else {
    LOG(ERROR) << "Unsupported operation " << description();
  }

  return success ? chaps_->session_handle() : base::nullopt;
}

base::Optional<brillo::Blob> ChapsCryptoOperation::Update(
    const brillo::Blob& input) {
  switch (description().type) {
    case OperationType::kSign:
      return chaps_->UpdateSignature(input)
                 ? base::Optional<brillo::Blob>(brillo::Blob())
                 : base::nullopt;
    case OperationType::kUnsupported:
      return base::nullopt;
  }
}

base::Optional<brillo::Blob> ChapsCryptoOperation::Finish() {
  base::Optional<brillo::Blob> result;

  switch (description().type) {
    case OperationType::kSign:
      result = chaps_->FinalizeSignature();
      break;
    case OperationType::kUnsupported:
      result = base::nullopt;
      break;
  }

  chaps_.reset();
  return result;
}

bool ChapsCryptoOperation::Abort() {
  chaps_.reset();
  return true;
}

}  // namespace context
}  // namespace keymaster
}  // namespace arc