// Copyright 2020 Google LLC. All Rights Reserved.

#include <vector>

#include "api/aead_whitebox.h"
#include "api/test_data.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace {
// These tests focus on functionality that should and should not exist between
// different Aead White-box instances. They assume that all other functionality
// of the white-box has already been verified.
class AeadWhiteboxCrossInstanceTest : public ::testing::Test {
 protected:
  void SetUp() override { init_data_ = GetValidAeadInitData(); }

  void TearDown() override {
    // Even if the pointer is null, as per the API, it should be safe to call
    // WB_Aead_Delete().
    WB_Aead_Delete(whitebox_a_);
    WB_Aead_Delete(whitebox_b_);
  }

  WB_Aead_Whitebox* whitebox_a_ = nullptr;
  WB_Aead_Whitebox* whitebox_b_ = nullptr;

  // Save a copy of the init data so that it is easier to create instances. It
  // would be nice if this could be const, but we need to set it in SetUp().
  std::vector<uint8_t> init_data_;

  // We need two different contexts so that we can have two unique instances.
  // Each context was generated using a random number generator. There is no
  // meaning to the bytes, the only condition is that they be different from
  // each other.
  const std::vector<uint8_t> context_a_ = {
      0xa0, 0xaf, 0x75, 0xbf, 0x6b, 0xff, 0x2f, 0x76,
      0x90, 0x47, 0xc4, 0x21, 0xb5, 0xf2, 0xb9, 0x62,
  };

  const std::vector<uint8_t> context_b_ = {
      0xfe, 0x6d, 0x01, 0x45, 0xbd, 0xa9, 0xcc, 0xd4,
      0x74, 0xaf, 0xca, 0xf8, 0xf9, 0x79, 0x2d, 0xf4,
  };

  // Both instances will always use the same plaintext. Since we are going to
  // verify decryption, we need the plaintext to be recognizable after it is
  // encrypted and then decrypted. To make it easier to read, we also avoided
  // random data.
  const std::vector<uint8_t> common_plaintext_ = {
      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
      0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
  };

  // We don't know how big the ciphertext will be, so we don't know how much
  // memory to allocate for our buffers. Since our plaintext is small, 256
  // should be large enough.
  const size_t safe_buffer_size_ = 256;
};

// Create two instances, each with unique contexts. Since they have different
// contexts, they should not be able to decrypt each other's data.
TEST_F(AeadWhiteboxCrossInstanceTest,
       DataVerificationErrorForDifferentContext) {
  ASSERT_EQ(WB_Aead_Create(init_data_.data(), init_data_.size(),
                           context_a_.data(), context_a_.size(), &whitebox_a_),
            WB_RESULT_OK);
  ASSERT_EQ(WB_Aead_Create(init_data_.data(), init_data_.size(),
                           context_b_.data(), context_b_.size(), &whitebox_b_),
            WB_RESULT_OK);

  size_t ciphertext_size = safe_buffer_size_;
  std::vector<uint8_t> ciphertext(ciphertext_size);

  size_t plaintext_size = safe_buffer_size_;
  std::vector<uint8_t> plaintext(plaintext_size);

  // The flow looks like:
  //   |common plaintext| --> [WB A] --> ciphertext --> [WB B] --> plaintext
  //
  // But it should fail after white-box tries to decrypt it, so |plaintext|
  // should never get written to.
  ASSERT_EQ(WB_Aead_Encrypt(whitebox_a_, common_plaintext_.data(),
                            common_plaintext_.size(), ciphertext.data(),
                            &ciphertext_size),
            WB_RESULT_OK);
  ciphertext.resize(ciphertext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_b_, ciphertext.data(), ciphertext.size(),
                            plaintext.data(), &plaintext_size),
            WB_RESULT_DATA_VERIFICATION_ERROR);
  plaintext.resize(plaintext_size);
}

// Create two instances, each using the same contexts. Since they have the same
// context, they should be able to decrypt each other's data.
TEST_F(AeadWhiteboxCrossInstanceTest, DataVerificationErrorForCommonContext) {
  ASSERT_EQ(WB_Aead_Create(init_data_.data(), init_data_.size(),
                           context_a_.data(), context_a_.size(), &whitebox_a_),
            WB_RESULT_OK);
  ASSERT_EQ(WB_Aead_Create(init_data_.data(), init_data_.size(),
                           context_a_.data(), context_a_.size(), &whitebox_b_),
            WB_RESULT_OK);

  size_t ciphertext_size = safe_buffer_size_;
  std::vector<uint8_t> ciphertext(ciphertext_size);

  size_t plaintext_size = safe_buffer_size_;
  std::vector<uint8_t> plaintext(plaintext_size);

  // The flow looks like:
  //   |common plaintext| --> [WB A] --> ciphertext --> [WB B] --> plaintext
  ASSERT_EQ(WB_Aead_Encrypt(whitebox_a_, common_plaintext_.data(),
                            common_plaintext_.size(), ciphertext.data(),
                            &ciphertext_size),
            WB_RESULT_OK);
  ciphertext.resize(ciphertext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_b_, ciphertext.data(), ciphertext.size(),
                            plaintext.data(), &plaintext_size),
            WB_RESULT_OK);
  plaintext.resize(plaintext_size);

  ASSERT_EQ(plaintext, common_plaintext_);
}
}  // namespace