// 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 {
// For our decrypt tests, we assume that WB_Aead_Create() and WB_Aead_Encrypt()
// work and each test will have valid encrypted data.
class AeadWhiteboxDecryptTest : public ::testing::Test {
 protected:
  void SetUp() override {
    const std::vector<uint8_t> init_data = GetValidAeadInitData();

    const std::vector<uint8_t> context = {
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    };

    ASSERT_EQ(WB_Aead_Create(init_data.data(), init_data.size(), context.data(),
                             context.size(), &whitebox_),
              WB_RESULT_OK);

    // Regardless of implementation, we need to have enough room in our cipher
    // text buffer to hold the additional information added by
    // WB_Aead_Encrypt(). So this number just needs to be reasonably big
    // compared to the input.
    size_t ciphertext_size = 256;
    ciphertext_.resize(ciphertext_size);
    ASSERT_EQ(WB_Aead_Encrypt(whitebox_, plaintext_.data(), plaintext_.size(),
                              ciphertext_.data(), &ciphertext_size),
              WB_RESULT_OK);
    ciphertext_.resize(ciphertext_size);
  }

  void TearDown() override { WB_Aead_Delete(whitebox_); }

  WB_Aead_Whitebox* whitebox_ = nullptr;

  // Since we are going to verify decryption, we need the plaintext to be
  // recognizable after it is encrypted and then decrypted.
  const std::vector<uint8_t> plaintext_ = {
      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
      0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F,
  };

  // Because the ciphertext will be implementation specific, we can't use
  // golden data for it. |ciphertext_| will be initialized in SetUp().
  std::vector<uint8_t> ciphertext_;
};

TEST_F(AeadWhiteboxDecryptTest, Success) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_, ciphertext_.data(), ciphertext_.size(),
                            plaintext.data(), &plaintext_size),
            WB_RESULT_OK);
  plaintext.resize(plaintext_size);

  ASSERT_EQ(plaintext, plaintext_);
}

TEST_F(AeadWhiteboxDecryptTest, InvalidParameterForNullWhitebox) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  ASSERT_EQ(WB_Aead_Decrypt(nullptr, ciphertext_.data(), ciphertext_.size(),
                            plaintext.data(), &plaintext_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxDecryptTest, InvalidParameterForNullInputData) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_, nullptr, ciphertext_.size(),
                            plaintext.data(), &plaintext_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxDecryptTest, InvalidParameterForInvalidInputDataSize) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  // Because we don't know how big the authentication tag will be, we need to
  // keep the number small. However, we don't use zero as "no data" is different
  // than "invalid data".
  ASSERT_EQ(WB_Aead_Decrypt(whitebox_, ciphertext_.data(), 4, plaintext.data(),
                            &plaintext_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxDecryptTest, InvalidParameterForNullOutputData) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_, ciphertext_.data(), ciphertext_.size(),
                            nullptr, &plaintext_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxDecryptTest, InvalidParameterForNullOutputDataSize) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_, ciphertext_.data(), ciphertext_.size(),
                            plaintext.data(), nullptr),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxDecryptTest, BufferTooSmall) {
  // We know what the plaintext should be, so if we make the plaintext buffer
  // just a little too small, it should fail without the test needing to know
  // anything out the implementation.
  size_t plaintext_size = plaintext_.size() - 1;
  std::vector<uint8_t> plaintext(plaintext_size);

  ASSERT_EQ(WB_Aead_Decrypt(whitebox_, ciphertext_.data(), ciphertext_.size(),
                            plaintext.data(), &plaintext_size),
            WB_RESULT_BUFFER_TOO_SMALL);
  ASSERT_EQ(plaintext_size, plaintext_.size());
}

TEST_F(AeadWhiteboxDecryptTest, DataVerificationError) {
  // The plaintext should be smaller than the input, but we use the input so
  // that we will always have enough.
  size_t plaintext_size = ciphertext_.size();
  std::vector<uint8_t> plaintext(plaintext_size);

  // Create invalid input by flipping one byte. Use bitwise-not so that, no
  // matter the original value, it will be different (we don't know what the
  // cipher text will actually be).
  std::vector<uint8_t> invalid_input = ciphertext_;
  invalid_input[0] = ~invalid_input[0];

  ASSERT_EQ(
      WB_Aead_Decrypt(whitebox_, invalid_input.data(), invalid_input.size(),
                      plaintext.data(), &plaintext_size),
      WB_RESULT_DATA_VERIFICATION_ERROR);
}
}  // namespace