// 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 "test_host.h"

#include <gtest/gtest.h>
#include <chrono>
#include <unordered_set>

#include "cdm_version.h"
#include "device_files.pb.h"
#include "file_store.h"
#include "log.h"

using namespace CDM_NAMESPACE;
using video_widevine_client::sdk::SavedStorage;

namespace {

// Some files are expected to go in global storage. All other files are expected
// to go in per-origin storage. To help us enforce this in tests, this set
// tracks the filenames that belong in global storage. TestHost::Storage will
// reject attempts to access these files via per-origin storage or to access
// files not in this list via global storage.
const std::unordered_set<std::string> kGlobalFilenames = {
    "usgtable.bin",          // CDM usage table data
    "StoredUsageTime.dat",   // Reference OEMCrypto usage table data
    "GenerationNumber.dat",  // Reference OEMCrypto master generation number
    "persistent.dat",        // Persistent data storage for certain TEE
                             // implementations
    "keybox.dat",  // Legacy file for storing keybox in non-secure storage.
    // CDM data for OTA keybox renewal.
    "okp.bin",
    "debug_ignore_keybox_count.txt",
    "debug_allow_test_keybox.txt",
    // Widevine L3 data files.
    "ay64.dat",
    "ay64.dat2",
    "ay64.dat3",
    "ay64.dat4",
    "ay64.dat5",
    "ay64.dat6",
    "l3_failure_file",
    wvutil::kOemCertificateFileName,
};

}  // namespace

TestHost::TestHost() : global_storage_(true), per_origin_storage_(false) {
  Reset();
}
TestHost::~TestHost() { wvutil::TestSleep::RemoveCallback(this); }

void TestHost::Reset() {
  auto now = std::chrono::system_clock().now();
  now_ = now.time_since_epoch() / std::chrono::milliseconds(1);
  wvutil::TestSleep::AddCallback(this);

  // Surprisingly, std::priority_queue has no clear().
  while (!timers_.empty()) {
    timers_.pop();
  }

  global_storage_.Reset();
  per_origin_storage_.Reset();
}

void TestHost::ElapseTime(int64_t milliseconds) {
  // Note that, during the time rollback tests, milliseconds will be negative,
  // so we cannot assume goal_time > now_.
  int64_t goal_time = now_ + milliseconds;

  // Walk forward from now_ to goal_time, stepping at each timer along the way
  // to fire its callback.
  while (!timers_.empty() && now_ < goal_time) {
    Timer t = timers_.top();
    ASSERT_GE(t.expiry_time(), now_);
    if (t.expiry_time() <= goal_time) {
      timers_.pop();
      now_ = t.expiry_time();
      t.client()->onTimerExpired(t.context());
    } else {
      // The next timer is further in the future than goal_time, so we are done
      // processing the timers.
      break;
    }
  }

  // No matter what happened with the timers, update now_ to the goal_time.
  now_ = goal_time;
}

size_t TestHost::NumTimers() const { return timers_.size(); }

int64_t TestHost::now() { return now_; }

void TestHost::setTimeout(int64_t delay_ms, IClient* client, void* context) {
  int64_t expiry_time = now_ + delay_ms;
  timers_.push(Timer(expiry_time, client, context));
}

void TestHost::cancel(IClient* client) {
  // Filter out the timers for this client and put the rest into |others|.
  std::priority_queue<Timer> others;

  while (timers_.size()) {
    Timer t = timers_.top();
    timers_.pop();

    if (t.client() != client) {
      others.push(t);
    }
  }

  // Now swap the queues.
  std::swap(timers_, others);
}

TestHost::Storage::Storage(bool is_global) : is_global_(is_global) { Reset(); }

void TestHost::Storage::Reset() {
  files_.clear();
}

bool TestHost::Storage::LoadFromString(const std::string& data) {
  SavedStorage proto;
  if (!proto.ParseFromString(data)) return false;
  Reset();
  files_.insert(proto.files().begin(), proto.files().end());
  return true;
}

bool TestHost::Storage::SaveToString(std::string* data) const {
  SavedStorage proto;
  proto.mutable_files()->insert(files_.begin(), files_.end());
  return proto.SerializeToString(data);
}

bool TestHost::Storage::read(const std::string& name, std::string* data) {
  if (wvutil::kLegacyCertificateFileName == name &&
      !g_host->baked_in_cert().empty()) {
    *data = g_host->baked_in_cert();
    return true;
  }
  StorageMap::iterator it = files_.find(name);
  bool ok = it != files_.end();
  LOGV("read file: %s: %s", name.c_str(), ok ? "ok" : "fail");
  if (!CheckFilename(name) || !ok) return false;
  *data = it->second;
  return true;
}

bool TestHost::Storage::write(const std::string& name,
                              const std::string& data) {
  LOGV("write file: %s", name.c_str());
  if (wvutil::kLegacyCertificateFileName == name &&
      !g_host->baked_in_cert().empty()) {
    return false;
  }
  if (!CheckFilename(name)) return false;
  files_[name] = data;
  return true;
}

bool TestHost::Storage::exists(const std::string& name) {
  if (wvutil::kLegacyCertificateFileName == name &&
      !g_host->baked_in_cert().empty()) {
    LOGV("exists? %s: always", name.c_str());
    return true;
  }
  StorageMap::iterator it = files_.find(name);
  bool ok = it != files_.end();
  LOGV("exists? %s: %s", name.c_str(), ok ? "true" : "false");
  if (!CheckFilename(name)) return false;
  return ok;
}

bool TestHost::Storage::remove(const std::string& name) {
  if (name.empty()) {
    // If no name, delete all files (see DeviceFiles::DeleteAllFiles())
    LOGV("remove all files");
    files_.clear();
    return true;
  }
  LOGV("remove: %s", name.c_str());
  if (!CheckFilename(name)) return false;
  return files_.erase(name) > 0;
}

int32_t TestHost::Storage::size(const std::string& name) {
  if (wvutil::kLegacyCertificateFileName == name &&
      !g_host->baked_in_cert().empty()) {
    LOGV("size? %s: always", name.c_str());
    return static_cast<int32_t>(g_host->baked_in_cert().size());
  }
  StorageMap::iterator it = files_.find(name);
  bool ok = (it != files_.end());
  LOGV("size? %s: %s", name.c_str(), ok ? "ok" : "fail");
  if (!CheckFilename(name) || !ok) return -1;
  return static_cast<int32_t>(it->second.size());
}

bool TestHost::Storage::list(std::vector<std::string>* names) {
  names->clear();
  for (StorageMap::iterator it = files_.begin(); it != files_.end(); it++) {
    names->push_back(it->first);
  }
  return true;
}

bool TestHost::Storage::CheckFilename(const std::string& name) {
  const bool is_global_filename =
      (kGlobalFilenames.find(name) != kGlobalFilenames.end());
  if (is_global_ != is_global_filename) {
    LOGE("Attempt to access %s in %s storage rejected.", name.c_str(),
         is_global_ ? "global" : "per-origin");
    return false;
  }
  return true;
}