// Copyright 2020 Google LLC. All Rights Reserved.

#include <vector>

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

namespace {
// Regardless of implementation, we need to have enough room in our output
// buffer to hold the additional information added by WB_Aead_Encrypt(). So
// this number just needs to be reasonably big compared to the input.
constexpr size_t kDefaultOutputSize = 256;

// For our encrypt tests, we assume that WB_Aead_Create() works and each test
// will have a functional instance. In these tests, we will not be checking the
// encrypt and decrypt flow, just the return codes of encrypt. For checking the
// encrypt-decrypt flow, see the decrypt tests as they will assume encryption
// works.
class AeadWhiteboxEncryptTest : 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);
  }

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

  WB_Aead_Whitebox* whitebox_ = nullptr;

  // Because we are not going to verify the encryption in these tests, we don't
  // need the input to be recognizable.
  const std::vector<uint8_t> input_ = {
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  };
};

TEST_F(AeadWhiteboxEncryptTest, Success) {
  size_t output_size = kDefaultOutputSize;
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(WB_Aead_Encrypt(whitebox_, input_.data(), input_.size(),
                            output.data(), &output_size),
            WB_RESULT_OK);

  // Additional information should have been added to the output (compared to
  // the input).
  ASSERT_GT(output_size, input_.size());
}

TEST_F(AeadWhiteboxEncryptTest, InvalidParameterForNullWhitebox) {
  size_t output_size = kDefaultOutputSize;
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(WB_Aead_Encrypt(nullptr, input_.data(), input_.size(),
                            output.data(), &output_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxEncryptTest, InvalidParameterForNullInputData) {
  size_t output_size = kDefaultOutputSize;
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(WB_Aead_Encrypt(whitebox_, nullptr, input_.size(), output.data(),
                            &output_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxEncryptTest, InvalidParameterForZeroInputSize) {
  size_t output_size = kDefaultOutputSize;
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(
      WB_Aead_Encrypt(whitebox_, input_.data(), 0, output.data(), &output_size),
      WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxEncryptTest, InvalidParameterForNullOutputData) {
  size_t output_size = kDefaultOutputSize;
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(WB_Aead_Encrypt(whitebox_, input_.data(), input_.size(), nullptr,
                            &output_size),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxEncryptTest, InvalidParameterForNullOutputDataSize) {
  size_t output_size = kDefaultOutputSize;
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(WB_Aead_Encrypt(whitebox_, input_.data(), input_.size(),
                            output.data(), nullptr),
            WB_RESULT_INVALID_PARAMETER);
}

TEST_F(AeadWhiteboxEncryptTest, BufferTooSmallForSmallOutputBuffer) {
  // Because WB_Aead_Encrypt() will add more data, having the output be the
  // same size as the input will make it by definition too small.
  size_t output_size = input_.size();
  std::vector<uint8_t> output(output_size);

  ASSERT_EQ(WB_Aead_Encrypt(whitebox_, input_.data(), input_.size(),
                            output.data(), &output_size),
            WB_RESULT_BUFFER_TOO_SMALL);

  // Additional room should be needed. We don't know how much more, so we can
  // only check that |output_size| is larger than |input_.size()|.
  ASSERT_GT(output_size, input_.size());
}
}  // namespace