// Copyright 2022 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/mac/internal/chunked_mac_wrapper.h"

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

#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/str_cat.h"
#include "absl/strings/string_view.h"
#include "tink/chunked_mac.h"
#include "tink/mac/internal/chunked_mac_impl.h"
#include "tink/mac/internal/stateful_mac.h"
#include "tink/primitive_set.h"
#include "tink/secret_data.h"
#include "tink/util/secret_data.h"
#include "tink/util/test_matchers.h"
#include "proto/tink.pb.h"

namespace crypto {
namespace tink {
namespace internal {
namespace {

using ::crypto::tink::test::IsOk;
using ::crypto::tink::test::IsOkAndHolds;
using ::crypto::tink::test::StatusIs;
using ::crypto::tink::util::SecretDataFromStringView;
using ::google::crypto::tink::KeysetInfo;
using ::google::crypto::tink::KeyStatusType;
using ::google::crypto::tink::OutputPrefixType;
using ::testing::Values;

class FakeStatefulMac : public StatefulMac {
 public:
  explicit FakeStatefulMac(absl::string_view name) : name_(name) {}

  absl::Status Update(absl::string_view data) override {
    absl::StrAppend(&buffer_, data);
    return absl::OkStatus();
  }

  absl::StatusOr<SecretData> FinalizeAsSecretData() override {
    return SecretDataFromStringView(absl::StrCat(name_, buffer_));
  }

 private:
  const std::string name_;
  std::string buffer_ = "";
};

class FakeStatefulMacFactory : public StatefulMacFactory {
 public:
  explicit FakeStatefulMacFactory(absl::string_view name) : name_(name) {}

  absl::StatusOr<std::unique_ptr<StatefulMac>> Create() const override {
    return std::unique_ptr<StatefulMac>(
        absl::make_unique<FakeStatefulMac>(name_));
  }

 private:
  std::string name_;
};

TEST(ChunkedMacWrapperTest, WrapNullptr) {
  EXPECT_THAT(ChunkedMacWrapper().Wrap(nullptr).status(),
              StatusIs(absl::StatusCode::kInternal));
}

TEST(ChunkedMacWrapperTest, WrapEmpty) {
  auto mac_set = std::make_unique<PrimitiveSet<ChunkedMac>>();
  EXPECT_THAT(ChunkedMacWrapper().Wrap(std::move(mac_set)).status(),
              StatusIs(absl::StatusCode::kInvalidArgument));
}

std::unique_ptr<ChunkedMac> CreateFakeChunkedMac(absl::string_view name) {
  return absl::make_unique<ChunkedMacImpl>(
      absl::make_unique<FakeStatefulMacFactory>(name));
}

KeysetInfo::KeyInfo CreateKeyInfo(uint32_t key_id,
                                  OutputPrefixType output_prefix_type) {
  KeysetInfo::KeyInfo key_info;
  key_info.set_output_prefix_type(output_prefix_type);
  key_info.set_key_id(key_id);
  key_info.set_status(KeyStatusType::ENABLED);
  return key_info;
}

TEST(ChunkedMacWrapperTest, ComputeMac) {
  PrimitiveSet<ChunkedMac>::Builder mac_set_builder;

  // Add primitives to the primitive set.
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac0:"),
      CreateKeyInfo(/*key_id=*/0x12d66f, OutputPrefixType::TINK));
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac1:"),
      CreateKeyInfo(/*key_id=*/0xb1539, OutputPrefixType::LEGACY));
  mac_set_builder.AddPrimaryPrimitive(
      CreateFakeChunkedMac("chunkedmac2:"),
      CreateKeyInfo(/*key_id=*/0x6e12af, OutputPrefixType::TINK));

  absl::StatusOr<PrimitiveSet<ChunkedMac>> mac_set =
      std::move(mac_set_builder).Build();
  ASSERT_THAT(mac_set, IsOk());

  // Wrap primitive set into a ChunkedMac.
  absl::StatusOr<std::unique_ptr<crypto::tink::ChunkedMac>> chunked_mac =
      ChunkedMacWrapper().Wrap(
          absl::make_unique<PrimitiveSet<ChunkedMac>>(*std::move(mac_set)));
  ASSERT_THAT(chunked_mac.status(), IsOk());

  absl::StatusOr<std::unique_ptr<ChunkedMacComputation>> computation =
      (*chunked_mac)->CreateComputation();
  EXPECT_THAT(computation.status(), IsOk());
  EXPECT_THAT((*computation)->Update("inputdata"), IsOk());
  absl::StatusOr<std::string> tag = (*computation)->ComputeMac();
  const std::string output_prefix = std::string("\x01\x00\x6e\x12\xaf", 5);
  const std::string raw_tag = "chunkedmac2:inputdata";
  EXPECT_THAT(tag, IsOkAndHolds(absl::StrCat(output_prefix, raw_tag)));
}

TEST(ChunkedMacWrapperTest, VerifyMacWithUniquePrefix) {
  PrimitiveSet<ChunkedMac>::Builder mac_set_builder;

  // Add primitives to primitive set.
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac0:"),
      CreateKeyInfo(/*key_id=*/0x12d66f, OutputPrefixType::TINK));
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac1:"),
      CreateKeyInfo(/*key_id=*/0xb1539, OutputPrefixType::LEGACY));
  mac_set_builder.AddPrimaryPrimitive(
      CreateFakeChunkedMac("chunkedmac2:"),
      CreateKeyInfo(/*key_id=*/0x6e12af, OutputPrefixType::TINK));

  absl::StatusOr<PrimitiveSet<ChunkedMac>> mac_set =
      std::move(mac_set_builder).Build();
  ASSERT_THAT(mac_set, IsOk());

  // Wrap primitive set into a ChunkedMac.
  absl::StatusOr<std::unique_ptr<crypto::tink::ChunkedMac>> chunked_mac =
      ChunkedMacWrapper().Wrap(
          absl::make_unique<PrimitiveSet<ChunkedMac>>(*std::move(mac_set)));
  ASSERT_THAT(chunked_mac.status(), IsOk());

  const std::string output_prefix = std::string("\x01\x00\x6e\x12\xaf", 5);
  const std::string raw_tag = "chunkedmac2:inputdata";
  absl::StatusOr<std::unique_ptr<ChunkedMacVerification>> verification =
      (*chunked_mac)->CreateVerification(absl::StrCat(output_prefix, raw_tag));
  EXPECT_THAT(verification.status(), IsOk());
  EXPECT_THAT((*verification)->Update("inputdata"), IsOk());
  EXPECT_THAT((*verification)->VerifyMac(), IsOk());
}

TEST(ChunkedMacWrapperTest, VerifyMacWithDuplicatePrefix) {
  PrimitiveSet<ChunkedMac>::Builder mac_set_builder;

  // Add primitives to primitive set.
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac0:"),
      CreateKeyInfo(/*key_id=*/0x12d66f, OutputPrefixType::LEGACY));
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac1:"),
      CreateKeyInfo(/*key_id=*/0x6e12af, OutputPrefixType::TINK));
  mac_set_builder.AddPrimaryPrimitive(
      CreateFakeChunkedMac("chunkedmac2:"),
      CreateKeyInfo(/*key_id=*/0x6e12af, OutputPrefixType::TINK));

  absl::StatusOr<PrimitiveSet<ChunkedMac>> mac_set =
      std::move(mac_set_builder).Build();
  ASSERT_THAT(mac_set, IsOk());

  // Wrap primitive set into a ChunkedMac.
  absl::StatusOr<std::unique_ptr<crypto::tink::ChunkedMac>> chunked_mac =
      ChunkedMacWrapper().Wrap(
          absl::make_unique<PrimitiveSet<ChunkedMac>>(*std::move(mac_set)));
  ASSERT_THAT(chunked_mac.status(), IsOk());

  const std::string output_prefix = std::string("\x01\x00\x6e\x12\xaf", 5);
  const std::string raw_tag = "chunkedmac1:inputdata";
  absl::StatusOr<std::unique_ptr<ChunkedMacVerification>> verification =
      (*chunked_mac)->CreateVerification(absl::StrCat(output_prefix, raw_tag));
  EXPECT_THAT(verification.status(), IsOk());
  EXPECT_THAT((*verification)->Update("inputdata"), IsOk());
  EXPECT_THAT((*verification)->VerifyMac(), IsOk());
}

TEST(ChunkedMacWrapperTest, VerifyMacWithRawTagStartingWithKeyId) {
  PrimitiveSet<ChunkedMac>::Builder mac_set_builder;

  const std::string key_id0 = std::string("\x01\x00\x12\xd6\x6f", 5);

  // Add primitives to primitive set.
  mac_set_builder.AddPrimitive(
      CreateFakeChunkedMac("chunkedmac0:"),
      CreateKeyInfo(/*key_id=*/0x12d66f, OutputPrefixType::TINK));
  mac_set_builder.AddPrimaryPrimitive(
      CreateFakeChunkedMac(/*name=*/absl::StrCat(key_id0, ":chunkedmac1:")),
      CreateKeyInfo(/*key_id=*/0x6e12af, OutputPrefixType::RAW));

  absl::StatusOr<PrimitiveSet<ChunkedMac>> mac_set =
      std::move(mac_set_builder).Build();
  ASSERT_THAT(mac_set, IsOk());

  // Wrap primitive set into a ChunkedMac.
  absl::StatusOr<std::unique_ptr<crypto::tink::ChunkedMac>> chunked_mac =
      ChunkedMacWrapper().Wrap(
          absl::make_unique<PrimitiveSet<ChunkedMac>>(*std::move(mac_set)));
  ASSERT_THAT(chunked_mac.status(), IsOk());

  const std::string raw_tag = absl::StrCat(key_id0, ":chunkedmac1:inputdata");
  absl::StatusOr<std::unique_ptr<ChunkedMacVerification>> verification =
      (*chunked_mac)->CreateVerification(raw_tag);
  EXPECT_THAT(verification.status(), IsOk());
  EXPECT_THAT((*verification)->Update("inputdata"), IsOk());
  EXPECT_THAT((*verification)->VerifyMac(), IsOk());
}

class ChunkedMacWrapperOutputPrefixTest
    : public testing::TestWithParam<OutputPrefixType> {};

INSTANTIATE_TEST_SUITE_P(
    ChunkedMacWrapperOutputPrefixTestSuite, ChunkedMacWrapperOutputPrefixTest,
    Values(OutputPrefixType::LEGACY, OutputPrefixType::RAW,
           OutputPrefixType::CRUNCHY, OutputPrefixType::TINK));

TEST_P(ChunkedMacWrapperOutputPrefixTest, ComputeVerifyMac) {
  OutputPrefixType output_prefix_type = GetParam();

  PrimitiveSet<ChunkedMac>::Builder mac_set_builder;

  // Add primitives to primitive set.
  mac_set_builder.AddPrimaryPrimitive(
      CreateFakeChunkedMac("chunkedmac:"),
      CreateKeyInfo(/*key_id=*/0x12d66f, output_prefix_type));

  absl::StatusOr<PrimitiveSet<ChunkedMac>> mac_set =
      std::move(mac_set_builder).Build();
  ASSERT_THAT(mac_set, IsOk());

  // Wrap primitive set into a ChunkedMac.
  absl::StatusOr<std::unique_ptr<crypto::tink::ChunkedMac>> chunked_mac =
      ChunkedMacWrapper().Wrap(
          absl::make_unique<PrimitiveSet<ChunkedMac>>(*std::move(mac_set)));
  ASSERT_THAT(chunked_mac.status(), IsOk());

  // Compute MAC via wrapper.
  absl::StatusOr<std::unique_ptr<ChunkedMacComputation>> mac_computation =
      (*chunked_mac)->CreateComputation();
  ASSERT_THAT(mac_computation.status(), IsOk());
  ASSERT_THAT((*mac_computation)->Update("inputdata"), IsOk());
  absl::StatusOr<std::string> tag = (*mac_computation)->ComputeMac();
  ASSERT_THAT(tag.status(), IsOk());

  // Verify MAC via wrapper.
  absl::StatusOr<std::unique_ptr<ChunkedMacVerification>> mac_verification =
      (*chunked_mac)->CreateVerification(*tag);
  ASSERT_THAT(mac_verification.status(), IsOk());
  ASSERT_THAT((*mac_verification)->Update("inputdata"), IsOk());
  ASSERT_THAT((*mac_verification)->VerifyMac(), IsOk());
}

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