// Copyright 2021 Google LLC
//
// 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 "tink/aead/internal/zero_copy_aead_wrapper.h"

#include <cstdint>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tink/aead.h"
#include "tink/aead/internal/mock_zero_copy_aead.h"
#include "tink/aead/internal/zero_copy_aead.h"
#include "tink/crypto_format.h"
#include "tink/primitive_set.h"
#include "tink/subtle/subtle_util.h"
#include "tink/util/test_matchers.h"
#include "proto/tink.pb.h"

namespace crypto {
namespace tink {
namespace internal {
namespace {

using ::crypto::tink::PrimitiveSet;
using ::crypto::tink::test::IsOk;
using ::crypto::tink::test::StatusIs;
using ::google::crypto::tink::KeysetInfo;
using ::google::crypto::tink::KeyStatusType;
using ::google::crypto::tink::OutputPrefixType;
using ::testing::_;
using ::testing::HasSubstr;
using ::testing::Return;
using ::testing::Unused;

constexpr absl::string_view kPlaintext = "Some data to encrypt.";
constexpr absl::string_view kAad = "Some data to authenticate.";
constexpr absl::string_view kCiphertext = "iv:Some data to encrypt.:tag";

using ZeroCopyAeadEntry =
    crypto::tink::PrimitiveSet<ZeroCopyAead>::Entry<ZeroCopyAead>;

TEST(ZeroCopyAeadWrapperEmptyTest, Nullptr) {
  ZeroCopyAeadWrapper wrapper;
  absl::StatusOr<std::unique_ptr<Aead>> aead_set = wrapper.Wrap(nullptr);
  EXPECT_THAT(aead_set.status(),
              StatusIs(absl::StatusCode::kInternal, HasSubstr("non-NULL")));
}

TEST(ZeroCopyAeadWrapperEmptyTest, Empty) {
  ZeroCopyAeadWrapper wrapper;
  absl::StatusOr<std::unique_ptr<Aead>> aead_set =
      wrapper.Wrap(absl::make_unique<PrimitiveSet<ZeroCopyAead>>());
  EXPECT_THAT(aead_set.status(), StatusIs(absl::StatusCode::kInvalidArgument,
                                          HasSubstr("no primary")));
}

// Returns an AEAD with expected return values for all its functions set via
// EXPECT_CALL. All values are derived from constants kPlaintext, kAad, and
// kCiphertext.
std::unique_ptr<MockZeroCopyAead> SetUpMockZeroCopyAead() {
  auto aead = absl::make_unique<MockZeroCopyAead>();

  EXPECT_CALL(*aead, MaxEncryptionSize(kPlaintext.size()))
      .WillRepeatedly(Return(kCiphertext.size()));
  EXPECT_CALL(*aead, Encrypt(kPlaintext, kAad, _))
      .WillRepeatedly([&](Unused, Unused, absl::Span<char> buffer) {
        memcpy(buffer.data(), kCiphertext.data(), kCiphertext.size());
        return kCiphertext.size();
      });
  EXPECT_CALL(*aead, MaxDecryptionSize(kCiphertext.size()))
      .WillRepeatedly(Return(kPlaintext.size()));
  EXPECT_CALL(*aead, Decrypt(kCiphertext, kAad, _))
      .WillRepeatedly([&](Unused, Unused, absl::Span<char> buffer) {
        std::memcpy(buffer.data(), kPlaintext.data(), kPlaintext.size());
        return kPlaintext.size();
      });

  return aead;
}

TEST(ZeroCopyAeadWrapperTest, EncryptDecrypt) {
  // Defines a Tink-type key.
  KeysetInfo::KeyInfo key_info;
  key_info.set_output_prefix_type(OutputPrefixType::TINK);
  key_info.set_key_id(1234543);
  key_info.set_status(KeyStatusType::ENABLED);

  // Creates a new AEAD set, adds a mock AEAD corresponding to the above key.
  PrimitiveSet<ZeroCopyAead>::Builder aead_set_builder;
  aead_set_builder.AddPrimaryPrimitive(SetUpMockZeroCopyAead(), key_info);
  absl::StatusOr<PrimitiveSet<ZeroCopyAead>> primitive_set =
      std::move(aead_set_builder).Build();
  ASSERT_THAT(primitive_set, IsOk());
  auto aead_set_ptr =
      std::make_unique<PrimitiveSet<ZeroCopyAead>>(std::move(*primitive_set));

  ZeroCopyAeadWrapper wrapper;
  absl::StatusOr<std::unique_ptr<Aead>> aead_set =
      wrapper.Wrap(std::move(aead_set_ptr));
  ASSERT_THAT(aead_set, IsOk());

  absl::StatusOr<std::string> ciphertext =
      (*aead_set)->Encrypt(kPlaintext, kAad);
  ASSERT_THAT(ciphertext, IsOk());
  absl::StatusOr<std::string> plaintext =
      (*aead_set)->Decrypt(*ciphertext, kAad);
  ASSERT_THAT(plaintext, IsOk());
  EXPECT_EQ(*plaintext, kPlaintext);
}

TEST(ZeroCopyAeadWrapperTest, EncryptMultipleKeys) {
  // Defines a Tink-type key.
  KeysetInfo::KeyInfo key_info;
  key_info.set_output_prefix_type(OutputPrefixType::TINK);
  key_info.set_key_id(1234543);
  key_info.set_status(KeyStatusType::ENABLED);

  // Creates a new AEAD set, adds a mock AEAD corresponding to the above key.
  PrimitiveSet<ZeroCopyAead>::Builder aead_set_builder;
  aead_set_builder.AddPrimaryPrimitive(SetUpMockZeroCopyAead(), key_info);
  absl::StatusOr<PrimitiveSet<ZeroCopyAead>> primitive_set =
      std::move(aead_set_builder).Build();
  ASSERT_THAT(primitive_set, IsOk());
  auto aead_set_ptr =
      std::make_unique<PrimitiveSet<ZeroCopyAead>>(std::move(*primitive_set));

  // Manually encrypt with the primary key.
  ZeroCopyAead& aead = aead_set_ptr->get_primary()->get_primitive();
  std::string ciphertext;
  subtle::ResizeStringUninitialized(
      &ciphertext, CryptoFormat::kNonRawPrefixSize +
                       aead.MaxEncryptionSize(kPlaintext.size()));
  absl::StatusOr<int64_t> ciphertext_size = aead.Encrypt(
      kPlaintext, kAad,
      absl::MakeSpan(ciphertext)
          .subspan(CryptoFormat::kNonRawPrefixSize, ciphertext.size()));
  ASSERT_THAT(ciphertext_size, IsOk());
  const std::string& key_id = aead_set_ptr->get_primary()->get_identifier();
  std::memcpy(&ciphertext[0], key_id.data(), key_id.size());
  ciphertext.resize(key_id.size() + *ciphertext_size);

  // Rebuild the PrimitiveSet with a second key.
  KeysetInfo::KeyInfo key_info2;
  key_info2.set_output_prefix_type(OutputPrefixType::TINK);
  key_info2.set_key_id(42);
  key_info2.set_status(KeyStatusType::ENABLED);
  std::unique_ptr<ZeroCopyAead> aead1 = absl::make_unique<MockZeroCopyAead>();

  PrimitiveSet<ZeroCopyAead>::Builder new_aead_set_builder;
  new_aead_set_builder.AddPrimaryPrimitive(SetUpMockZeroCopyAead(), key_info);
  new_aead_set_builder.AddPrimitive(std::move(aead1), key_info2);
  absl::StatusOr<PrimitiveSet<ZeroCopyAead>> new_aead_set =
      std::move(new_aead_set_builder).Build();
  ASSERT_THAT(new_aead_set, IsOk());
  aead_set_ptr =
      std::make_unique<PrimitiveSet<ZeroCopyAead>>(std::move(*new_aead_set));

  ZeroCopyAeadWrapper wrapper;
  absl::StatusOr<std::unique_ptr<Aead>> aead_set =
      wrapper.Wrap(std::move(aead_set_ptr));
  ASSERT_THAT(aead_set, IsOk());

  // Encrypt with the wrapped AEAD and check that the result is equal to
  // encrypting directly with the primary key.
  absl::StatusOr<std::string> wrap_ciphertext =
      (*aead_set)->Encrypt(kPlaintext, kAad);
  ASSERT_THAT(wrap_ciphertext, IsOk());
  EXPECT_EQ(*wrap_ciphertext, ciphertext);
}

TEST(ZeroCopyAeadWrapperTest, EncryptDecryptRawKey) {
  // Defines a Tink-type key.
  KeysetInfo::KeyInfo key_info;
  key_info.set_output_prefix_type(OutputPrefixType::TINK);
  key_info.set_key_id(1234543);
  key_info.set_status(KeyStatusType::ENABLED);

  // Add raw key to AEAD set.
  KeysetInfo::KeyInfo key_info2;
  key_info2.set_output_prefix_type(OutputPrefixType::RAW);
  key_info2.set_key_id(1234);
  key_info2.set_status(KeyStatusType::ENABLED);

  // Rebuild the PrimitiveSet with the raw key as primary.
  PrimitiveSet<ZeroCopyAead>::Builder aead_set_builder;
  aead_set_builder.AddPrimitive(SetUpMockZeroCopyAead(), key_info);
  aead_set_builder.AddPrimaryPrimitive(SetUpMockZeroCopyAead(), key_info2);
  absl::StatusOr<PrimitiveSet<ZeroCopyAead>> primitive_set =
      std::move(aead_set_builder).Build();
  ASSERT_THAT(primitive_set, IsOk());
  auto aead_set_ptr =
      std::make_unique<PrimitiveSet<ZeroCopyAead>>(std::move(*primitive_set));

  // Manually encrypt with the raw key.
  absl::StatusOr<const std::vector<std::unique_ptr<ZeroCopyAeadEntry>>*>
      raw_primitives = aead_set_ptr->get_raw_primitives();
  ASSERT_THAT(raw_primitives, IsOk());
  EXPECT_EQ((*raw_primitives)->size(), 1);
  ZeroCopyAead& aead = (*raw_primitives)->front()->get_primitive();
  std::string ciphertext;
  subtle::ResizeStringUninitialized(&ciphertext,
                                    aead.MaxEncryptionSize(kPlaintext.size()));
  absl::StatusOr<int64_t> ciphertext_size =
      aead.Encrypt(kPlaintext, kAad, absl::MakeSpan(ciphertext));
  ASSERT_THAT(ciphertext_size, IsOk());
  ciphertext.resize(*ciphertext_size);

  // Encrypt with the wrapped AEAD and check that the result is equal to
  // encrypting directly with the raw key.
  ZeroCopyAeadWrapper wrapper;
  absl::StatusOr<std::unique_ptr<Aead>> aead_set =
      wrapper.Wrap(std::move(aead_set_ptr));
  ASSERT_THAT(aead_set, IsOk());
  absl::StatusOr<std::string> wrap_ciphertext =
      (*aead_set)->Encrypt(kPlaintext, kAad);
  ASSERT_THAT(wrap_ciphertext, IsOk());
  EXPECT_EQ(*wrap_ciphertext, ciphertext);

  // Manually decrypt with the raw key.
  std::string plaintext;
  subtle::ResizeStringUninitialized(&plaintext,
                                    aead.MaxDecryptionSize(ciphertext.size()));
  absl::StatusOr<int64_t> plaintext_size =
      aead.Decrypt(ciphertext, kAad, absl::MakeSpan(plaintext));
  ASSERT_THAT(plaintext_size, IsOk());
  plaintext.resize(*plaintext_size);
  EXPECT_EQ(plaintext, kPlaintext);

  // Decrypt with the wrapped AEAD.
  absl::StatusOr<std::string> wrap_plaintext =
      (*aead_set)->Decrypt(ciphertext, kAad);
  ASSERT_THAT(wrap_plaintext, IsOk());
  EXPECT_EQ(*wrap_plaintext, kPlaintext);
}

TEST(ZeroCopyAeadWrapperTest, EncryptBadDecrypt) {
  // Defines a Tink-type key.
  KeysetInfo::KeyInfo key_info;
  key_info.set_output_prefix_type(OutputPrefixType::TINK);
  key_info.set_key_id(1234543);
  key_info.set_status(KeyStatusType::ENABLED);

  // Creates a new AEAD set, adds a mock AEAD corresponding to the above key.
  PrimitiveSet<ZeroCopyAead>::Builder aead_set_builder;
  aead_set_builder.AddPrimaryPrimitive(SetUpMockZeroCopyAead(), key_info);
  absl::StatusOr<PrimitiveSet<ZeroCopyAead>> primitive_set =
      std::move(aead_set_builder).Build();
  ASSERT_THAT(primitive_set, IsOk());
  auto aead_set_ptr =
      std::make_unique<PrimitiveSet<ZeroCopyAead>>(std::move(*primitive_set));

  ZeroCopyAeadWrapper wrapper;
  absl::StatusOr<std::unique_ptr<Aead>> aead_set =
      wrapper.Wrap(std::move(aead_set_ptr));
  ASSERT_THAT(aead_set, IsOk());

  absl::StatusOr<std::string> plaintext =
      (*aead_set)->Decrypt("some bad ciphertext", kAad);
  EXPECT_EQ(plaintext.status().code(), absl::StatusCode::kInvalidArgument);
}

}  // namespace
}  // namespace internal
}  // namespace tink
}  // namespace crypto
