// Copyright 2018 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.

#include <gtest/gtest.h>

#include <utility>

#include "log.h"
#include "string_conversions.h"

namespace wvutil {

namespace {

// Test vectors as suggested by http://tools.ietf.org/html/rfc4648#section-10
const std::string kNullString("");
const std::string kf("f");
const std::string kfo("fo");
const std::string kfoo("foo");
const std::string kfoob("foob");
const std::string kfooba("fooba");
const std::string kfoobar("foobar");
const std::string kfB64("Zg==");
const std::string kfoB64("Zm8=");
const std::string kfooB64("Zm9v");
const std::string kfoobB64("Zm9vYg==");
const std::string kfoobaB64("Zm9vYmE=");
const std::string kfoobarB64("Zm9vYmFy");

// Arbitrary clear test vectors
const std::string kMultipleOf24BitsData("Good day!");
const std::string kOneByteOverData("Hello Friend!");
const std::string kTwoBytesOverData("Hello Friend!!");
const std::string kTestData =
    "\030\361\\\366\267> \331\210\360\\-\311:\324\256\376"
    "\261\234\241\326d\326\177\346\346\223\333Y\305\214\330";

// Arbitrary encoded test vectors
const std::string kMultipleOf24BitsB64Data("R29vZCBkYXkh");
const std::string kOneByteOverB64Data("SGVsbG8gRnJpZW5kIQ==");
const std::string kTwoBytesOverB64Data("SGVsbG8gRnJpZW5kISE=");
const std::string kB64TestData = "GPFc9rc+INmI8FwtyTrUrv6xnKHWZNZ/5uaT21nFjNg=";

const std::pair<const std::string*, const std::string*> kBase64TestVectors[] = {
    make_pair(&kNullString, &kNullString),
    make_pair(&kf, &kfB64),
    make_pair(&kfo, &kfoB64),
    make_pair(&kfoo, &kfooB64),
    make_pair(&kfoob, &kfoobB64),
    make_pair(&kfooba, &kfoobaB64),
    make_pair(&kfoobar, &kfoobarB64),
    make_pair(&kMultipleOf24BitsData, &kMultipleOf24BitsB64Data),
    make_pair(&kOneByteOverData, &kOneByteOverB64Data),
    make_pair(&kTwoBytesOverData, &kTwoBytesOverB64Data),
    make_pair(&kTestData, &kB64TestData)};

// Arbitrary invalid base64 test vectors
const std::string kBase64ErrorVectors[] = {"Foo$sa",
                                           "Foo\x99\x23\xfa\02",
                                           "Foo==Foo",
                                           "FooBa",
                                           "SGVsbG8sIFdvcmxkI===",
                                           "SGVsbG8sIFdvcmxkI======",
                                           "SGVsbG8sIFdvcmxkIQp=="};

std::string ConvertToBase64WebSafe(const std::string& std_base64_string) {
  std::string str(std_base64_string);
  for (size_t i = 0; i < str.size(); ++i) {
    if (str[i] == '+')
      str[i] = '-';
    else if (str[i] == '/')
      str[i] = '_';
  }
  return str;
}

}  // namespace

class Base64EncodeDecodeTest
    : public ::testing::TestWithParam<
          std::pair<const std::string*, const std::string*> > {};

TEST_P(Base64EncodeDecodeTest, EncodeDecodeTest) {
  std::pair<const std::string*, const std::string*> values = GetParam();
  const std::string& plain_text_string = *(values.first);
  const std::string& expected_encoded = *(values.second);

  // Encode from string.
  const std::string b64_string_encoded = Base64Encode(plain_text_string);
  EXPECT_EQ(b64_string_encoded, expected_encoded);

  // Encode from vector.
  const std::vector<uint8_t> plain_text_vector(plain_text_string.begin(),
                                               plain_text_string.end());
  const std::string b64_vector_encoded = Base64Encode(plain_text_vector);
  EXPECT_EQ(b64_vector_encoded, expected_encoded);

  // Decode from string.
  const std::vector<uint8_t> decoded_vector = Base64Decode(expected_encoded);
  EXPECT_EQ(decoded_vector, plain_text_vector);
}

TEST_P(Base64EncodeDecodeTest, WebSafeEncodeDecodeTest) {
  std::pair<const std::string*, const std::string*> values = GetParam();
  const std::string& plain_text_string = *(values.first);
  const std::string& expected_encoded =
      ConvertToBase64WebSafe(*(values.second));

  // Encode from string.
  const std::string b64_string_encoded = Base64SafeEncode(plain_text_string);
  EXPECT_EQ(b64_string_encoded, expected_encoded);

  // Encode from vector.
  const std::vector<uint8_t> plain_text_vector(plain_text_string.begin(),
                                               plain_text_string.end());
  const std::string b64_vector_encoded = Base64SafeEncode(plain_text_vector);
  EXPECT_EQ(b64_vector_encoded, expected_encoded);

  // Decode from string.
  const std::vector<uint8_t> decoded_vector =
      Base64SafeDecode(expected_encoded);
  EXPECT_EQ(decoded_vector, plain_text_vector);
}

TEST_P(Base64EncodeDecodeTest, WebSafeEncodeNoPad) {
  std::pair<const std::string*, const std::string*> values = GetParam();
  const std::string& plain_text_string = *(values.first);
  const std::string& padded_encoded = ConvertToBase64WebSafe(*(values.second));

  // Encode from string.
  const std::string b64_string_encoded =
      Base64SafeEncodeNoPad(plain_text_string);

  // If input is empty, output will be empty.
  if (plain_text_string.empty()) {
    EXPECT_TRUE(b64_string_encoded.empty());
    return;
  }

  if (padded_encoded.back() == '=') {
    // If padding is present in the regular encoding, then it should be
    // striped from the result.
    EXPECT_NE(b64_string_encoded.back(), '=');
    const std::string expected_encoded =
        padded_encoded.substr(0, b64_string_encoded.size());
    EXPECT_EQ(b64_string_encoded, expected_encoded);
  } else {
    // If no padding is present, then results should be equal.
    EXPECT_EQ(b64_string_encoded, padded_encoded);
  }

  // Encode from vector.
  const std::vector<uint8_t> plain_text_vector(plain_text_string.begin(),
                                               plain_text_string.end());
  const std::string b64_vector_encoded =
      Base64SafeEncodeNoPad(plain_text_vector);
  // Assuming the above has passed, the results should be the same as
  // a result encoded from a string.
  EXPECT_EQ(b64_vector_encoded, b64_string_encoded);
}

class Base64ErrorDecodeTest : public ::testing::TestWithParam<std::string> {};

TEST_P(Base64ErrorDecodeTest, EncoderErrors) {
  const std::vector<uint8_t> standard_result = Base64Decode(GetParam());
  EXPECT_TRUE(standard_result.empty());
  const std::vector<uint8_t> safe_result = Base64SafeDecode(GetParam());
  EXPECT_TRUE(safe_result.empty());
}

INSTANTIATE_TEST_SUITE_P(ExecutesBase64Test, Base64EncodeDecodeTest,
                         ::testing::ValuesIn(kBase64TestVectors));

INSTANTIATE_TEST_SUITE_P(ExecutesBase64Test, Base64ErrorDecodeTest,
                         ::testing::ValuesIn(kBase64ErrorVectors));

class HtoNLL64Test : public ::testing::Test {};

TEST_F(HtoNLL64Test, PositiveNumber) {
  uint8_t data[8] = {1, 2, 3, 4, 5, 6, 7, 8};
  int64_t* network_byte_order = reinterpret_cast<int64_t*>(data);
  int64_t host_byte_order = htonll64(*network_byte_order);
  EXPECT_EQ(0x0102030405060708, host_byte_order);
}
TEST_F(HtoNLL64Test, NegativeNumber) {
  uint8_t data[8] = {0xfe, 2, 3, 4, 5, 6, 7, 8};
  int64_t* network_byte_order = reinterpret_cast<int64_t*>(data);
  int64_t host_byte_order = htonll64(*network_byte_order);
  EXPECT_EQ(-0x01FdFcFbFaF9F8F8, host_byte_order);
}

TEST(SafeByteIdToString, EmptyId) {
  const std::string kEmptyString;
  EXPECT_EQ(SafeByteIdToString(kEmptyString), "<empty>");

  const std::vector<uint8_t> kEmptyVector;
  EXPECT_EQ(SafeByteIdToString(kEmptyVector), "<empty>");
}

TEST(SafeByteIdToString, AllAlphaNumeric_NoEscape) {
  const std::string sid("Hello, World!");
  const std::vector<uint8_t> vid(sid.begin(), sid.end());
  EXPECT_EQ(SafeByteIdToString(sid), "\"Hello, World!\"");
  EXPECT_EQ(SafeByteIdToString(vid), "\"Hello, World!\"");
}

TEST(SafeByteIdToString, AllAlphaNumeric_NoEscape_Exhaustive) {
  for (char ch = 'a'; ch <= 'z'; ch++) {
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }

  for (char ch = 'A'; ch <= 'Z'; ch++) {
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }

  for (char ch = '0'; ch <= '9'; ch++) {
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }
}

TEST(SafeByteIdToString, AllSymbols_NoEscape) {
  for (char ch = ' '; ch <= '/'; ch++) {
    if (ch == '"') continue;
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }

  for (char ch = ':'; ch <= '@'; ch++) {
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }

  for (char ch = '['; ch <= '`'; ch++) {
    if (ch == '\\') continue;
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }

  for (char ch = '{'; ch <= '~'; ch++) {
    const std::string sid(1, ch);
    const std::vector<uint8_t> vid(1, ch);
    std::string expected(1, '"');
    expected.push_back(ch);
    expected.push_back('"');
    EXPECT_EQ(SafeByteIdToString(sid), expected);
    EXPECT_EQ(SafeByteIdToString(vid), expected);
  }
}

TEST(SafeByteIdToString, Escapable_NoEscape) {
  const std::string quote_sid(1, '"');
  const std::vector<uint8_t> quote_vid(1, '"');
  EXPECT_EQ(SafeByteIdToString(quote_sid), "\"\\\"\"");
  EXPECT_EQ(SafeByteIdToString(quote_vid), "\"\\\"\"");

  const std::string backslash_sid(1, '\\');
  const std::vector<uint8_t> backslash_vid(1, '\\');
  EXPECT_EQ(SafeByteIdToString(backslash_sid), "\"\\\\\"");
  EXPECT_EQ(SafeByteIdToString(backslash_vid), "\"\\\\\"");
}

TEST(SafeByteIdToString, AllNonPrintable) {
  const std::vector<uint8_t> vid = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  const std::string sid(vid.begin(), vid.end());

  EXPECT_EQ(SafeByteIdToString(vid), "00010203040506070809");
  EXPECT_EQ(SafeByteIdToString(sid), "00010203040506070809");
}
}  // namespace wvutil