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

// These tests perform various end-to-end actions similar to what an application
// would do. They verify that policies specified on UAT are honored on the
// device.

#include "reboot_test.h"

#include <sstream>

#include "log.h"

namespace wvcdm {
FileSystem* RebootTest::file_system_;

namespace {
// We will encode a value string by wrapping it in braces, or as hex.
// If the string is not printable, or if it has unmatched braces, then we use
// hex. Otherwise, we surround the whole string with braces.
std::string EncodeString(const std::string& data) {
  int braces_count = 0;
  for (size_t i = 0; i < data.length(); i++) {
    if (data[i] == '{') braces_count++;
    if (data[i] == '}') braces_count--;
    // If printable or whitespace (because '\n' is not printable?!?).
    bool printable = isprint(data[i]) || isspace(data[i]);
    // If there are any unprintable characters, except whitespace, or if we
    // close a brace before we open it, then just use hex.
    if (!printable || braces_count < 0) {
      return "0x" + wvcdm::b2a_hex(data) + ",";
    }
  }
  // If we left any braces open, then use hex.
  if (braces_count != 0) return "0x" + wvcdm::b2a_hex(data) + ",";
  return "{" + data + "},";
}

// Encode a map key for dumping. When we encode a map, we expect the keys to be
// like filenames, so we can separate them with colons and whitespace. If the
// key has these special characters, we will encode as hex.
std::string EncodeKey(const std::string& data) {
  if (data.length() == 0) {
    LOGE("Encoding empty string as key!");
    return "EMPTY:";
  }
  // When decoding, we assume that a key starting with "0x" is in hex. So we
  // can't have any keys that start with "0x".
  if (data.substr(0, 2) == "0x") return "0x" + wvcdm::b2a_hex(data) + ":";
  // If the key is just  is not printable, or if it has unmatched braces, then
  // we use hex. Otherwise, we surround the whole string with braces.
  for (size_t i = 0; i < data.length(); i++) {
    if (!isprint(data[i]) || (data[i] == ':')) {
      return "0x" + wvcdm::b2a_hex(data) + ":";
    }
  }
  return data + ":";
}

// In between keys and values, we will ignore whitespace. This allows a human to
// edit the persistent data a little bit without breaking anything.
void SkipSpace(const std::string& encoded, size_t* index) {
  if (!index) return;
  while (*index < encoded.length() && isspace(encoded[*index])) (*index)++;
}

// Decode a string that was encoded using EncodeString.
std::string DecodeString(const std::string& encoded, size_t* index) {
  if (!index) return "";
  SkipSpace(encoded, index);
  if (*index + 2 >=
      encoded.length()) {  // Encoded string has at least 3 characters.
    LOGE("Error decoding short string from %s at %zd", encoded.c_str(), *index);
    *index = encoded.length();
    return "";
  }
  if (encoded[*index] == '{') {
    (*index)++;
    size_t start = *index;
    int braces_count = 1;
    while (*index < encoded.length()) {
      if (encoded[*index] == '{') braces_count++;
      if (encoded[*index] == '}') braces_count--;
      if (braces_count == 0) {
        size_t end = *index;
        (*index) += 2;  // absorb the comma and the '}', too.
        return encoded.substr(start, end - start);
      }
      (*index)++;
    }
    std::string tail = encoded.substr(start);
    LOGE("Non-terminated brace %s at %zd: %s", encoded.c_str(), start,
         tail.c_str());
    *index = encoded.length();
    return "";
  }
  if (encoded[*index] != '0' || encoded[*index + 1] != 'x') {
    std::string tail = encoded.substr(*index);
    LOGE("Hex should start with 0x in %s at %zd: %s", encoded.c_str(), *index,
         tail.c_str());
    *index = encoded.length();
    return "";
  }
  *index += 2;
  size_t start = *index;
  while (*index < encoded.length()) {
    if (encoded[*index] == ',') {
      size_t end = *index;
      std::vector<uint8_t> result =
          wvcdm::a2b_hex(encoded.substr(start, end - start));
      (*index)++;  // absorb the comma.
      return std::string(result.begin(), result.end());
    }
    (*index)++;
  }
  std::string tail = encoded.substr(start);
  LOGE("Bad encoding in %s at %zd: %s", encoded.c_str(), start, tail.c_str());
  *index = encoded.length();
  return "";
}

// Decode a string that was encoded with EncodeKey.
std::string DecodeKey(const std::string& encoded, size_t* index) {
  if (!index) return "";
  SkipSpace(encoded, index);
  if (*index + 1 >= encoded.length()) {
    LOGE("Error decoding key from %s at %zd", encoded.c_str(), *index);
    *index = encoded.length();
    return "";
  }
  // If it starts with 0x, then it is in hex.
  if (encoded[*index] == '0' && encoded[*index + 1] == 'x') {
    size_t start = *index + 2;
    while (*index < encoded.length() && encoded[*index] != ':') (*index)++;
    size_t end = *index;
    std::vector<uint8_t> result =
        wvcdm::a2b_hex(encoded.substr(start, end - start));
    (*index)++;  // skip the colon.
    return std::string(result.begin(), result.end());
  }
  size_t start = *index;
  while (*index < encoded.length() && encoded[*index] != ':') (*index)++;
  size_t end = *index;
  (*index)++;  // skip the colon.
  return encoded.substr(start, end - start);
}
}  // namespace

std::string RebootTest::DumpData(
    const std::map<std::string, std::string>& data) {
  std::ostringstream output;
  output << "{\n";
  for (const auto& entry : data) {
    output << "  " << EncodeKey(entry.first) << " "
           << EncodeString(entry.second) + "\n";
  }
  output << "}\n";
  return output.str();
}

bool RebootTest::ParseDump(const std::string& dump,
                           std::map<std::string, std::string>* data) {
  size_t index = 0;
  SkipSpace(dump, &index);
  if (index >= dump.length()) return false;
  if (dump[index] != '{') {
    LOGE("Dump does not start with '{'");
    return false;
  }
  index++;  // absorb '{'
  while (true) {
    SkipSpace(dump, &index);
    if (index >= dump.length()) return false;
    if (dump[index] == '}') {
      index++;  // absorb '}'
      SkipSpace(dump, &index);
      if (index != dump.length()) {
        std::string tail = dump.substr(index);
        LOGE("Trailing data in dump. %s at %zd: %s", dump.c_str(), index,
             tail.c_str());
        return false;
      }
      return true;
    }
    std::string tail = dump.substr(index);
    std::string key = DecodeKey(dump, &index);
    std::string value = DecodeString(dump, &index);
    (*data)[key] = value;
  }
}

void RebootTest::SetUp() {
  WvCdmTestBase::SetUp();
  if (!file_system_) file_system_ = new FileSystem();

  const ::testing::TestInfo* const test_info =
      ::testing::UnitTest::GetInstance()->current_test_info();
  std::string test_name =
      std::string(test_info->test_case_name()) + "-" + test_info->name();
  persistent_data_filename_ =
      config_.test_data_path() + "/" + test_name + ".dat";
  LOGD("Running test pass %d for %s", test_pass(), test_name.c_str());
  // Don't read data on the first pass, but do read data for all other passes.
  if (test_pass() > 0) {
    EXPECT_TRUE(file_system_->Exists(persistent_data_filename_));
    ssize_t file_size = file_system_->FileSize(persistent_data_filename_);
    auto file =
        file_system_->Open(persistent_data_filename_, file_system_->kReadOnly);
    ASSERT_TRUE(file);
    std::string dump(file_size, ' ');
    ssize_t read = file->Read(&dump[0], dump.size());
    EXPECT_EQ(read, file_size) << "Error reading persistent data file.";
    EXPECT_TRUE(ParseDump(dump, &persistent_data_));
  }
}

void RebootTest::TearDown() {
  auto file = file_system_->Open(persistent_data_filename_,
                                 FileSystem::kCreate | FileSystem::kTruncate);
  ASSERT_TRUE(file) << "Failed to open file: " << persistent_data_filename_;
  std::string dump = DumpData(persistent_data_);
  const ssize_t bytes_written = file->Write(dump.data(), dump.length());
  EXPECT_EQ(bytes_written, static_cast<ssize_t>(dump.length()));
  WvCdmTestBase::TearDown();
}

/** Test the dump and restore functions above. This does not test CDM
    functionality. */
TEST_F(RebootTest, TestDumpUtil) {
  // Check that an empty map can be saved.
  std::map<std::string, std::string> map1;
  const std::string dump = DumpData(map1);
  std::map<std::string, std::string> map2;
  EXPECT_TRUE(ParseDump(dump, &map2));
  EXPECT_EQ(map1, map2);
  // Now fill it with some data and try again.
  map1["key1"] = "this is a string.    ";
  map1["key2"] = "mismatch } {";
  map1["key3"] = "mismatch } ";
  map1["key4"] = "mismatch  {";
  map1["key5"] = "this: { has { matched } } braces { /.,)(**&^$&^% }";
  map1["key6"] = "";
  map1["00 whitespace in key 00"] = "value is ok";
  // This key looks like it might be hex. It should show up as hex in the
  // save file.
  map1["0x_bad_key_00"] = "value is ok";
  const std::string dump2 = DumpData(map1);
  std::map<std::string, std::string> map3;
  EXPECT_TRUE(ParseDump(dump2, &map3));
  EXPECT_EQ(map1, map3);
  if (test_pass() == 0) {
    persistent_data_ = map1;
  } else {
    EXPECT_EQ(persistent_data_, map1);
  }
}

/** Verify that the file system stores files from one test pass to the next. */
TEST_F(RebootTest, FilesArePersistent) {
  const std::string key = "saved_value";
  const std::string value = "the string that is saved";
  if (test_pass() == 0) {
    // There should be no persistent data on the first pass.
    EXPECT_NE(persistent_data_[key], value);
    // We will save some data for the next pass.
    persistent_data_[key] = value;
  } else {
    // There should be persistent data set in first pass.
    EXPECT_EQ(persistent_data_[key], value);
  }
}

}  // namespace wvcdm