// Copyright 2019 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/win/conflicts/inspection_results_cache.h"

#include <algorithm>
#include <array>
#include <cstdint>
#include <string>
#include <string_view>
#include <utility>

#include "base/containers/span.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/files/important_file_writer.h"
#include "base/pickle.h"
#include "base/strings/string_util.h"
#include "crypto/obsolete/md5.h"

std::array<uint8_t, crypto::obsolete::Md5::kSize>
base::Md5ForWinInspectionResultsCache(base::span<const uint8_t> payload) {
  return crypto::obsolete::Md5::Hash(payload);
}

namespace {

// The current version of the cache format. Must be incremented every time the
// format changes.
constexpr int kVersion = 1;

// Serializes a [ModuleInfoKey, ModuleInspectionResult, time stamp] tuple into
// |pickle|. The |time_stamp| represents the last time this specific inspection
// result was used, and was calculated using CalculateTimeStamp().
void SerializeInspectionResult(const ModuleInfoKey& module_key,
                               const ModuleInspectionResult& inspection_result,
                               uint32_t time_stamp,
                               base::Pickle* pickle) {
  // ModuleInfoKey:
  module_key.module_path.WriteToPickle(pickle);
  pickle->WriteUInt32(module_key.module_size);
  pickle->WriteUInt32(module_key.module_time_date_stamp);

  // ModuleInspectionResult:
  pickle->WriteString16(inspection_result.location);
  pickle->WriteString16(inspection_result.basename);
  pickle->WriteString16(inspection_result.product_name);
  pickle->WriteString16(inspection_result.description);
  pickle->WriteString16(inspection_result.version);
  pickle->WriteInt(static_cast<int>(inspection_result.certificate_info.type));
  inspection_result.certificate_info.path.WriteToPickle(pickle);
  pickle->WriteString16(inspection_result.certificate_info.subject);

  // Time stamp:
  pickle->WriteUInt32(time_stamp);
}

// Deserializes a [ModuleInfoKey, ModuleInspectionResult] pair into |result| by
// reading from |pickle_iterator|. Skips pairs whose time stamp is older than
// |min_time_stamp|. Returns true if |result| contains a valid inspection
// result.
bool DeserializeInspectionResult(uint32_t min_time_stamp,
                                 base::PickleIterator* pickle_iterator,
                                 InspectionResultsCache* result) {
  DCHECK(pickle_iterator);
  DCHECK(result);

  // ModuleInfoKey:
  base::FilePath module_path;
  uint32_t module_size = 0;
  uint32_t module_time_date_stamp = 0;
  if (!module_path.ReadFromPickle(pickle_iterator) ||
      !pickle_iterator->ReadUInt32(&module_size) ||
      !pickle_iterator->ReadUInt32(&module_time_date_stamp)) {
    return false;
  }

  std::pair<ModuleInfoKey, std::pair<ModuleInspectionResult, uint32_t>> value(
      std::piecewise_construct,
      std::forward_as_tuple(std::move(module_path), module_size,
                            module_time_date_stamp),
      std::forward_as_tuple());

  // ModuleInspectionResult and time stamp:
  ModuleInspectionResult& inspection_result = value.second.first;
  uint32_t& time_stamp = value.second.second;

  std::u16string location, basename;
  if (!pickle_iterator->ReadString16(&inspection_result.location) ||
      !pickle_iterator->ReadString16(&inspection_result.basename) ||
      !pickle_iterator->ReadString16(&inspection_result.product_name) ||
      !pickle_iterator->ReadString16(&inspection_result.description) ||
      !pickle_iterator->ReadString16(&inspection_result.version) ||
      !pickle_iterator->ReadInt(
          reinterpret_cast<int*>(&inspection_result.certificate_info.type)) ||
      !inspection_result.certificate_info.path.ReadFromPickle(
          pickle_iterator) ||
      !pickle_iterator->ReadString16(
          &inspection_result.certificate_info.subject) ||
      !pickle_iterator->ReadUInt32(&time_stamp)) {
    return false;
  }

  // Only keep this element if it is not expired. An expired entry is not an
  // error.
  if (time_stamp >= min_time_stamp)
    result->insert(std::move(value));

  return true;
}

// Serializes an InspectionResultsCache into a base::Pickle. The serialized data
// contains a version number and a MD5 checksum at the end.
base::Pickle SerializeInspectionResultsCache(
    const InspectionResultsCache& inspection_results_cache) {
  base::Pickle pickle;

  pickle.WriteInt(kVersion);
  pickle.WriteUInt64(inspection_results_cache.size());

  for (const auto& inspection_result : inspection_results_cache) {
    SerializeInspectionResult(inspection_result.first,
                              inspection_result.second.first,
                              inspection_result.second.second, &pickle);
  }

  // Append the md5 digest of the data to detect serializations errors.
  std::array<uint8_t, crypto::obsolete::Md5::kSize> md5_digest =
      Md5ForWinInspectionResultsCache(pickle.payload_bytes());
  pickle.WriteBytes(md5_digest);
  return pickle;
}

// Deserializes an InspectionResultsCache from |pickle_iterator|. This function
// ensures that both the version and the checksum of the data are valid. Returns
// a ReadCacheResult value indicating what failed if unsuccessful.
ReadCacheResult DeserializeInspectionResultsCache(
    uint32_t min_time_stamp,
    base::PickleIterator pickle_iterator,
    InspectionResultsCache* result) {
  DCHECK(result);

  // Make a copy of the iterator before iterating to be able to consume the full
  // payload later when computing the MD5 hash.
  base::PickleIterator pickle_iterator_copy = pickle_iterator;

  // Check the version number.
  int version = 0;
  if (!pickle_iterator.ReadInt(&version))
    return ReadCacheResult::kFailDeserializeVersion;
  if (version != kVersion)
    return ReadCacheResult::kFailInvalidVersion;

  // Retrieve the count of inspection results.
  uint64_t inspection_result_count = 0;
  if (!pickle_iterator.ReadUInt64(&inspection_result_count))
    return ReadCacheResult::kFailDeserializeCount;
  if (inspection_result_count < 0)
    return ReadCacheResult::kFailInvalidCount;

  // Now deserialize all the ModuleInspectionResults. Failure to deserialize one
  // inspection result counts as an invalid cache.
  for (uint64_t i = 0; i < inspection_result_count; i++) {
    if (!DeserializeInspectionResult(min_time_stamp, &pickle_iterator, result))
      return ReadCacheResult::kFailDeserializeInspectionResult;
  }

  // Now check the md5 checksum.
  const std::array<uint8_t, crypto::obsolete::Md5::kSize>* read_md5_digest =
      nullptr;
  if (!pickle_iterator.ReadBytes(
          reinterpret_cast<const char**>(&read_md5_digest),
          crypto::obsolete::Md5::kSize)) {
    return ReadCacheResult::kFailDeserializeMD5;
  }

  // Check if the md5 checksum matches.
  std::optional<base::span<const uint8_t>> payload =
      pickle_iterator_copy.ReadBytes(pickle_iterator_copy.RemainingBytes());
  CHECK(payload.has_value());
  if (!std::ranges::equal(
          *read_md5_digest,
          Md5ForWinInspectionResultsCache(payload->first(
              payload->size() - crypto::obsolete::Md5::kSize)))) {
    return ReadCacheResult::kFailInvalidMD5;
  }

  return ReadCacheResult::kSuccess;
}

}  // namespace

void AddInspectionResultToCache(
    const ModuleInfoKey& module_key,
    const ModuleInspectionResult& inspection_result,
    InspectionResultsCache* inspection_results_cache) {
  auto insert_result = inspection_results_cache->insert(std::make_pair(
      module_key, std::make_pair(inspection_result,
                                 CalculateTimeStamp(base::Time::Now()))));
  // An insertion should always succeed because the user of this code never
  // tries to insert an existing value.
  DCHECK(insert_result.second);
}

std::optional<ModuleInspectionResult> GetInspectionResultFromCache(
    const ModuleInfoKey& module_key,
    InspectionResultsCache* inspection_results_cache) {
  std::optional<ModuleInspectionResult> inspection_result;

  auto it = inspection_results_cache->find(module_key);
  if (it != inspection_results_cache->end()) {
    // Update the time stamp.
    it->second.second = CalculateTimeStamp(base::Time::Now());
    inspection_result = it->second.first;
  }

  return inspection_result;
}

ReadCacheResult ReadInspectionResultsCache(
    const base::FilePath& file_path,
    uint32_t min_time_stamp,
    InspectionResultsCache* inspection_results_cache) {
  std::string contents;
  if (!ReadFileToString(file_path, &contents))
    return ReadCacheResult::kFailReadFile;

  InspectionResultsCache temporary_result;
  ReadCacheResult read_result = DeserializeInspectionResultsCache(
      min_time_stamp,
      base::PickleIterator::WithData(base::as_byte_span(contents)),
      &temporary_result);

  // Only update the output cache when successful.
  if (read_result == ReadCacheResult::kSuccess)
    *inspection_results_cache = std::move(temporary_result);

  return read_result;
}

bool WriteInspectionResultsCache(
    const base::FilePath& file_path,
    const InspectionResultsCache& inspection_results_cache) {
  base::Pickle pickle =
      SerializeInspectionResultsCache(inspection_results_cache);

  // TODO(crbug.com/40106434): Investigate if using WriteFileAtomically() in a
  // CONTINUE_ON_SHUTDOWN sequence can cause too many corrupted caches.
  return base::ImportantFileWriter::WriteFileAtomically(file_path,
                                                        pickle.AsStringView());
}