components/federated_learning/floc_sorting_lsh_clusters_service.cc

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

#include "components/federated_learning/floc_sorting_lsh_clusters_service.h"

#include <utility>

#include "base/files/file.h"
#include "base/files/file_util.h"
#include "base/task/post_task.h"
#include "base/task/thread_pool.h"
#include "base/task_runner_util.h"
#include "components/federated_learning/floc_constants.h"
#include "third_party/protobuf/src/google/protobuf/io/coded_stream.h"
#include "third_party/protobuf/src/google/protobuf/io/zero_copy_stream_impl_lite.h"

namespace federated_learning {

namespace {

class CopyingFileInputStream : public google::protobuf::io::CopyingInputStream {
 public:
  explicit CopyingFileInputStream(base::File file) : file_(std::move(file)) {}

  CopyingFileInputStream(const CopyingFileInputStream&) = delete;
  CopyingFileInputStream& operator=(const CopyingFileInputStream&) = delete;

  ~CopyingFileInputStream() override = default;

  // google::protobuf::io::CopyingInputStream:
  int Read(void* buffer, int size) override {
    return file_.ReadAtCurrentPosNoBestEffort(static_cast<char*>(buffer), size);
  }

 private:
  base::File file_;
};

absl::optional<uint64_t> ApplySortingLshOnBackgroundThread(
    uint64_t sim_hash,
    const base::FilePath& file_path) {
  base::File sorting_lsh_clusters_file(
      file_path, base::File::FLAG_OPEN | base::File::FLAG_READ);
  if (!sorting_lsh_clusters_file.IsValid())
    return absl::nullopt;

  CopyingFileInputStream copying_stream(std::move(sorting_lsh_clusters_file));
  google::protobuf::io::CopyingInputStreamAdaptor zero_copy_stream_adaptor(
      &copying_stream);

  google::protobuf::io::CodedInputStream input_stream(
      &zero_copy_stream_adaptor);

  // The file should contain a list of integers. The 7th-order bit represents
  // whether the cohort should be blocked. The number represented by the 1st-6th
  // bits should be within the range [0, MaxNumberOfBitsInFloc]. Suppose the
  // list is l, then S(i) = 2^(l[i] & 0b111111) represents the the number of
  // hashes that can be associated with this floc id. The cumulative sum of S(i)
  // represents the boundary sim_hash values. We will use the higher index to
  // encode |sim_hash|, i.e. if sim_hash is within range
  // [CumSum(S(i-1)), CumSum(S(i))), |i| will be output floc.
  //
  // 0 is always an implicit CumSum boundary, i.e. if
  // 0 <= |sim_hash| < 2^(l[0]), then the index 0 will be the output floc.
  //
  // However, if the is_blocked bit (i.e. l[i] & 0b1000000) indicates that the
  // cohort should be blocked, we will output an invalid floc id.

  // Input sanitization: As we compute on the fly, we will check to make sure
  // each encountered entry, after dropping the is_blocked bit, is within
  // [0, MaxNumberOfBitsInFloc]. Besides, the cumulative sum should be no
  // greater than 2^MaxNumberOfBitsInFloc at any given time. If we cannot find
  // an index i, it means the the final cumulative sum is less than
  // 2^MaxNumberOfBitsInFloc, while we expect it to be exactly
  // 2^MaxNumberOfBitsInFloc, and we should also fail in this case. When some
  // check fails, we will also output an invalid floc id.
  //
  // A stricter sanitization would be to always stream all numbers and check
  // properties. We skip doing this to save some computation cost.
  const uint64_t kExpectedFinalCumulativeSum = (1ULL << kMaxNumberOfBitsInFloc);
  DCHECK_LT(sim_hash, kExpectedFinalCumulativeSum);

  uint64_t cumulative_sum = 0;
  uint32_t next_combined;

  // TODO(yaoxia): Add metrics for when the file has unexpected format.

  for (uint64_t index = 0; input_stream.ReadVarint32(&next_combined); ++index) {
    // Sanitizing error: the entry used more than |kSortingLshMaxBits| bits.
    if ((next_combined >> kSortingLshMaxBits) > 0)
      return absl::nullopt;

    bool is_blocked = next_combined & kSortingLshBlockedMask;
    uint32_t next = next_combined & kSortingLshSizeMask;

    // Sanitizing error
    if (next > kMaxNumberOfBitsInFloc)
      return absl::nullopt;

    cumulative_sum += (1ULL << next);

    // Sanitizing error
    if (cumulative_sum > kExpectedFinalCumulativeSum)
      return absl::nullopt;

    // Found the sim-hash upper bound. Use the index as the new floc.
    if (cumulative_sum > sim_hash) {
      if (is_blocked)
        return absl::nullopt;

      return index;
    }
  }

  // Sanitizing error: we didn't find a sim-hash upper bound, but we expect to
  // always find it after finish iterating through the list.
  return absl::nullopt;
}

}  // namespace

FlocSortingLshClustersService::FlocSortingLshClustersService()
    : background_task_runner_(base::ThreadPool::CreateSequencedTaskRunner(
          {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
           base::TaskShutdownBehavior::SKIP_ON_SHUTDOWN})),
      weak_ptr_factory_(this) {}

FlocSortingLshClustersService::~FlocSortingLshClustersService() = default;

void FlocSortingLshClustersService::AddObserver(Observer* observer) {
  observers_.AddObserver(observer);
}

void FlocSortingLshClustersService::RemoveObserver(Observer* observer) {
  observers_.RemoveObserver(observer);
}

bool FlocSortingLshClustersService::IsSortingLshClustersFileReady() const {
  return first_file_ready_seen_;
}

void FlocSortingLshClustersService::OnSortingLshClustersFileReady(
    const base::FilePath& file_path,
    const base::Version& version) {
  sorting_lsh_clusters_file_path_ = file_path;
  sorting_lsh_clusters_version_ = version;
  first_file_ready_seen_ = true;

  for (auto& observer : observers_)
    observer.OnSortingLshClustersFileReady();
}

void FlocSortingLshClustersService::ApplySortingLsh(
    uint64_t sim_hash,
    ApplySortingLshCallback callback) {
  DCHECK(first_file_ready_seen_);

  base::PostTaskAndReplyWithResult(
      background_task_runner_.get(), FROM_HERE,
      base::BindOnce(&ApplySortingLshOnBackgroundThread, sim_hash,
                     sorting_lsh_clusters_file_path_),
      base::BindOnce(&FlocSortingLshClustersService::DidApplySortingLsh,
                     weak_ptr_factory_.GetWeakPtr(), std::move(callback),
                     sorting_lsh_clusters_version_));
}

void FlocSortingLshClustersService::SetBackgroundTaskRunnerForTesting(
    scoped_refptr<base::SequencedTaskRunner> background_task_runner) {
  background_task_runner_ = background_task_runner;
}

void FlocSortingLshClustersService::DidApplySortingLsh(
    ApplySortingLshCallback callback,
    base::Version version,
    absl::optional<uint64_t> final_hash) {
  std::move(callback).Run(std::move(final_hash), std::move(version));
}

}  // namespace federated_learning