chrome/browser/policy/messaging_layer/storage/storage_unittest.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 "chrome/browser/policy/messaging_layer/storage/storage.h"

#include <cstdint>
#include <tuple>
#include <utility>

#include "base/files/scoped_temp_dir.h"
#include "base/optional.h"
#include "base/sequenced_task_runner.h"
#include "base/strings/strcat.h"
#include "base/strings/string_number_conversions.h"
#include "base/synchronization/waitable_event.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/task_environment.h"
#include "chrome/browser/policy/messaging_layer/encryption/decryption.h"
#include "chrome/browser/policy/messaging_layer/encryption/encryption.h"
#include "chrome/browser/policy/messaging_layer/encryption/test_encryption_module.h"
#include "chrome/browser/policy/messaging_layer/storage/storage_configuration.h"
#include "chrome/browser/policy/messaging_layer/util/status.h"
#include "chrome/browser/policy/messaging_layer/util/status_macros.h"
#include "chrome/browser/policy/messaging_layer/util/statusor.h"
#include "components/policy/proto/record.pb.h"
#include "components/policy/proto/record_constants.pb.h"
#include "crypto/sha2.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/boringssl/src/include/openssl/curve25519.h"

using ::testing::_;
using ::testing::Between;
using ::testing::Eq;
using ::testing::Invoke;
using ::testing::NotNull;
using ::testing::Property;
using ::testing::Return;
using ::testing::Sequence;
using ::testing::StrEq;
using ::testing::WithArg;

namespace reporting {
namespace {

// Usage (in tests only):
//
//   TestEvent<ResType> e;
//   ... Do some async work passing e.cb() as a completion callback of
//       base::OnceCallback<void(ResType* res)> type which also may perform some
//       other action specified by |done| callback provided by the caller.
//   ... = e.result();  // Will wait for e.cb() to be called and return the
//       collected result.
//
template <typename ResType>
class TestEvent {
 public:
  TestEvent() : run_loop_(std::make_unique<base::RunLoop>()) {}
  ~TestEvent() { EXPECT_FALSE(run_loop_->running()) << "Not responded"; }
  TestEvent(const TestEvent& other) = delete;
  TestEvent& operator=(const TestEvent& other) = delete;
  ResType result() {
    run_loop_->Run();
    return std::forward<ResType>(result_);
  }

  // Completion callback to hand over to the processing method.
  base::OnceCallback<void(ResType res)> cb() {
    return base::BindOnce(
        [](base::RunLoop* run_loop, ResType* result, ResType res) {
          *result = std::forward<ResType>(res);
          run_loop->Quit();
        },
        base::Unretained(run_loop_.get()), base::Unretained(&result_));
  }

 private:
  std::unique_ptr<base::RunLoop> run_loop_;
  ResType result_;
};

// Context of single decryption. Self-destructs upon completion or failure.
class SingleDecryptionContext {
 public:
  SingleDecryptionContext(
      const EncryptedRecord& encrypted_record,
      scoped_refptr<Decryptor> decryptor,
      base::OnceCallback<void(StatusOr<base::StringPiece>)> response)
      : encrypted_record_(encrypted_record),
        decryptor_(decryptor),
        response_(std::move(response)) {}

  SingleDecryptionContext(const SingleDecryptionContext& other) = delete;
  SingleDecryptionContext& operator=(const SingleDecryptionContext& other) =
      delete;

  ~SingleDecryptionContext() {
    DCHECK(!response_) << "Self-destruct without prior response";
  }

  void Start() {
    base::ThreadPool::PostTask(
        FROM_HERE,
        base::BindOnce(&SingleDecryptionContext::RetrieveMatchingPrivateKey,
                       base::Unretained(this)));
  }

 private:
  void Respond(StatusOr<base::StringPiece> result) {
    std::move(response_).Run(result);
    delete this;
  }

  void RetrieveMatchingPrivateKey() {
    // Retrieve private key that matches public key hash.
    decryptor_->RetrieveMatchingPrivateKey(
        encrypted_record_.encryption_info().public_key_id(),
        base::BindOnce(
            [](SingleDecryptionContext* self,
               StatusOr<std::string> private_key_result) {
              if (!private_key_result.ok()) {
                self->Respond(private_key_result.status());
                return;
              }
              base::ThreadPool::PostTask(
                  FROM_HERE,
                  base::BindOnce(&SingleDecryptionContext::DecryptSharedSecret,
                                 base::Unretained(self),
                                 private_key_result.ValueOrDie()));
            },
            base::Unretained(this)));
  }

  void DecryptSharedSecret(base::StringPiece private_key) {
    // Decrypt shared secret from private key and peer public key.
    auto shared_secret_result = decryptor_->DecryptSecret(
        private_key, encrypted_record_.encryption_info().encryption_key());
    if (!shared_secret_result.ok()) {
      Respond(shared_secret_result.status());
      return;
    }
    base::ThreadPool::PostTask(
        FROM_HERE, base::BindOnce(&SingleDecryptionContext::OpenRecord,
                                  base::Unretained(this),
                                  shared_secret_result.ValueOrDie()));
  }

  void OpenRecord(base::StringPiece shared_secret) {
    decryptor_->OpenRecord(
        shared_secret,
        base::BindOnce(
            [](SingleDecryptionContext* self,
               StatusOr<Decryptor::Handle*> handle_result) {
              if (!handle_result.ok()) {
                self->Respond(handle_result.status());
                return;
              }
              base::ThreadPool::PostTask(
                  FROM_HERE,
                  base::BindOnce(&SingleDecryptionContext::AddToRecord,
                                 base::Unretained(self),
                                 base::Unretained(handle_result.ValueOrDie())));
            },
            base::Unretained(this)));
  }

  void AddToRecord(Decryptor::Handle* handle) {
    handle->AddToRecord(
        encrypted_record_.encrypted_wrapped_record(),
        base::BindOnce(
            [](SingleDecryptionContext* self, Decryptor::Handle* handle,
               Status status) {
              if (!status.ok()) {
                self->Respond(status);
                return;
              }
              base::ThreadPool::PostTask(
                  FROM_HERE,
                  base::BindOnce(&SingleDecryptionContext::CloseRecord,
                                 base::Unretained(self),
                                 base::Unretained(handle)));
            },
            base::Unretained(this), base::Unretained(handle)));
  }

  void CloseRecord(Decryptor::Handle* handle) {
    handle->CloseRecord(base::BindOnce(
        [](SingleDecryptionContext* self,
           StatusOr<base::StringPiece> decryption_result) {
          self->Respond(decryption_result);
        },
        base::Unretained(this)));
  }

 private:
  const EncryptedRecord encrypted_record_;
  const scoped_refptr<Decryptor> decryptor_;
  base::OnceCallback<void(StatusOr<base::StringPiece>)> response_;
};

class MockUploadClient : public Storage::UploaderInterface {
 public:
  // Mapping of <generation id, sequencing id> to matching record digest.
  // Whenever a record is uploaded and includes last record digest, this map
  // should have that digest already recorded. Only the first record in a
  // generation is uploaded without last record digest.
  using LastRecordDigestMap = std::map<std::tuple<Priority,
                                                  uint64_t /*generation id*/,
                                                  uint64_t /*sequencing id*/>,
                                       std::string /*digest*/>;

  explicit MockUploadClient(
      LastRecordDigestMap* last_record_digest_map,
      scoped_refptr<base::SequencedTaskRunner> sequenced_task_runner,
      scoped_refptr<Decryptor> decryptor)
      : last_record_digest_map_(last_record_digest_map),
        sequenced_task_runner_(sequenced_task_runner),
        decryptor_(decryptor) {}

  void ProcessRecord(EncryptedRecord encrypted_record,
                     base::OnceCallback<void(bool)> processed_cb) override {
    const auto& sequencing_information =
        encrypted_record.sequencing_information();
    if (!encrypted_record.has_encryption_info()) {
      // Wrapped record is not encrypted.
      WrappedRecord wrapped_record;
      ASSERT_TRUE(wrapped_record.ParseFromString(
          encrypted_record.encrypted_wrapped_record()));
      ScheduleVerifyRecord(sequencing_information, std::move(wrapped_record),
                           std::move(processed_cb));
      return;
    }
    // Decrypt encrypted_record.
    (new SingleDecryptionContext(
         encrypted_record, decryptor_,
         base::BindOnce(
             [](SequencingInformation sequencing_information,
                base::OnceCallback<void(bool)> processed_cb,
                MockUploadClient* client, StatusOr<base::StringPiece> result) {
               ASSERT_OK(result.status());
               WrappedRecord wrapped_record;
               ASSERT_TRUE(wrapped_record.ParseFromArray(
                   result.ValueOrDie().data(), result.ValueOrDie().size()));
               // Verify wrapped record once decrypted.
               client->ScheduleVerifyRecord(sequencing_information,
                                            std::move(wrapped_record),
                                            std::move(processed_cb));
             },
             sequencing_information, std::move(processed_cb),
             base::Unretained(this))))
        ->Start();
  }

  void ProcessGap(SequencingInformation start,
                  uint64_t count,
                  base::OnceCallback<void(bool)> processed_cb) override {
    LOG(FATAL) << "Gap not implemented yet";
  }

  void Completed(bool need_encryption_key, Status status) override {
    UploadComplete(need_encryption_key, status);
  }

  MOCK_METHOD(bool,
              UploadRecord,
              (Priority, uint64_t, base::StringPiece),
              (const));
  MOCK_METHOD(bool, UploadRecordFailure, (Status), (const));
  MOCK_METHOD(void, UploadComplete, (bool, Status), (const));

  // Helper class for setting up mock client expectations of a successful
  // completion.
  class SetUp {
   public:
    SetUp(Priority priority, MockUploadClient* client)
        : priority_(priority), client_(client) {}

    ~SetUp() {
      EXPECT_CALL(*client_, UploadRecordFailure(_))
          .Times(0)
          .InSequence(client_->test_upload_sequence_);
      EXPECT_CALL(*client_, UploadComplete(/*need_encryption_key=*/false,
                                           Eq(Status::StatusOK())))
          .Times(1)
          .InSequence(client_->test_upload_sequence_);
    }

    SetUp& Required(uint64_t sequencing_id, base::StringPiece value) {
      EXPECT_CALL(*client_, UploadRecord(Eq(priority_), Eq(sequencing_id),
                                         StrEq(std::string(value))))
          .InSequence(client_->test_upload_sequence_)
          .WillOnce(Return(true));
      return *this;
    }

    SetUp& Possible(uint64_t sequencing_id, base::StringPiece value) {
      EXPECT_CALL(*client_, UploadRecord(Eq(priority_), Eq(sequencing_id),
                                         StrEq(std::string(value))))
          .Times(Between(0, 1))
          .InSequence(client_->test_upload_sequence_)
          .WillRepeatedly(Return(true));
      return *this;
    }

   private:
    Priority priority_;
    MockUploadClient* const client_;
  };

  // Helper class for setting up mock client expectations on empty queue.
  class SetEmpty {
   public:
    SetEmpty(Priority priority, MockUploadClient* client)
        : priority_(priority), client_(client) {}

    ~SetEmpty() {
      EXPECT_CALL(*client_, UploadRecord(Eq(priority_), _, _)).Times(0);
      EXPECT_CALL(*client_, UploadRecordFailure(_)).Times(0);
      EXPECT_CALL(*client_, UploadComplete(/*need_encryption_key=*/false,
                                           Property(&Status::error_code,
                                                    Eq(error::OUT_OF_RANGE))))
          .Times(1);
    }

   private:
    Priority priority_;
    MockUploadClient* const client_;
  };

  // Helper class for setting up mock client expectations for key delivery.
  class SetKeyDelivery {
   public:
    SetKeyDelivery(Priority priority, MockUploadClient* client)
        : priority_(priority), client_(client) {}

    ~SetKeyDelivery() {
      EXPECT_CALL(*client_, UploadRecord(Eq(priority_), _, _)).Times(0);
      EXPECT_CALL(*client_, UploadRecordFailure(_)).Times(0);
      EXPECT_CALL(*client_, UploadComplete(/*need_encryption_key=*/true,
                                           Eq(Status::StatusOK())))
          .Times(1);
    }

   private:
    Priority priority_;
    MockUploadClient* const client_;
  };

 private:
  void ScheduleVerifyRecord(SequencingInformation sequencing_information,
                            WrappedRecord wrapped_record,
                            base::OnceCallback<void(bool)> processed_cb) {
    sequenced_task_runner_->PostTask(
        FROM_HERE,
        base::BindOnce(&MockUploadClient::VerifyRecord, base::Unretained(this),
                       sequencing_information, std::move(wrapped_record),
                       std::move(processed_cb)));
  }

  void VerifyRecord(SequencingInformation sequencing_information,
                    WrappedRecord wrapped_record,
                    base::OnceCallback<void(bool)> processed_cb) {
    // Verify generation match.
    if (generation_id_.has_value() &&
        generation_id_.value() != sequencing_information.generation_id()) {
      std::move(processed_cb)
          .Run(UploadRecordFailure(Status(
              error::DATA_LOSS,
              base::StrCat({"Generation id mismatch, expected=",
                            base::NumberToString(generation_id_.value()),
                            " actual=",
                            base::NumberToString(
                                sequencing_information.generation_id())}))));
      return;
    }
    if (!generation_id_.has_value()) {
      generation_id_ = sequencing_information.generation_id();
    }

    // Verify digest and its match.
    // Last record digest is not verified yet, since duplicate records are
    // accepted in this test.
    {
      std::string serialized_record;
      wrapped_record.record().SerializeToString(&serialized_record);
      const auto record_digest = crypto::SHA256HashString(serialized_record);
      DCHECK_EQ(record_digest.size(), crypto::kSHA256Length);
      if (record_digest != wrapped_record.record_digest()) {
        std::move(processed_cb)
            .Run(UploadRecordFailure(
                Status(error::DATA_LOSS, "Record digest mismatch")));
        return;
      }
      if (wrapped_record.has_last_record_digest()) {
        auto it = last_record_digest_map_->find(
            std::make_tuple(sequencing_information.priority(),
                            sequencing_information.sequencing_id() - 1,
                            sequencing_information.generation_id()));
        if (it == last_record_digest_map_->end() ||
            it->second != wrapped_record.last_record_digest()) {
          std::move(processed_cb)
              .Run(UploadRecordFailure(
                  Status(error::DATA_LOSS, "Last record digest mismatch")));
          return;
        }
      }
      last_record_digest_map_->emplace(
          std::make_tuple(sequencing_information.priority(),
                          sequencing_information.sequencing_id(),
                          sequencing_information.generation_id()),
          record_digest);
    }

    std::move(processed_cb)
        .Run(UploadRecord(sequencing_information.priority(),
                          sequencing_information.sequencing_id(),
                          wrapped_record.record().data()));
  }

  base::Optional<uint64_t> generation_id_;
  LastRecordDigestMap* const last_record_digest_map_;
  scoped_refptr<base::SequencedTaskRunner> sequenced_task_runner_;

  const scoped_refptr<Decryptor> decryptor_;

  Sequence test_upload_sequence_;
};

class StorageTest
    : public ::testing::TestWithParam<::testing::tuple<bool, size_t>> {
 protected:
  void SetUp() override {
    ASSERT_TRUE(location_.CreateUniqueTempDir());
    // Encryption is disabled by default.
    ASSERT_FALSE(EncryptionModule::is_enabled());
    if (::testing::get<0>(GetParam())) {
      // Enable encryption.
      scoped_feature_list_.InitFromCommandLine(
          {EncryptionModule::kEncryptedReporting}, {});
      // Create decryption module.
      auto decryptor_result = Decryptor::Create();
      ASSERT_OK(decryptor_result.status()) << decryptor_result.status();
      decryptor_ = std::move(decryptor_result.ValueOrDie());
      // First creation of Storage would need key delivered.
      expect_to_need_key_ = true;
    }
  }

  void TearDown() override { ResetTestStorage(); }

  StatusOr<scoped_refptr<Storage>> CreateTestStorage(
      const StorageOptions& options,
      scoped_refptr<EncryptionModule> encryption_module) {
    if (expect_to_need_key_) {
      // Set uploader expectations for any queue; expect no records and need
      // key. Make sure no uploads happen, and key is requested.
      EXPECT_CALL(set_mock_uploader_expectations_, Call(_, NotNull()))
          .WillOnce(Invoke(
              [](Priority priority, MockUploadClient* mock_upload_client) {
                MockUploadClient::SetKeyDelivery(priority, mock_upload_client);
              }));
    }
    // Initialize Storage with no key.
    TestEvent<StatusOr<scoped_refptr<Storage>>> e;
    Storage::Create(options,
                    base::BindRepeating(&StorageTest::BuildMockUploader,
                                        base::Unretained(this)),
                    encryption_module, e.cb());
    ASSIGN_OR_RETURN(auto storage, e.result());
    if (expect_to_need_key_) {
      // Provision the storage with a key.
      // Key delivery must have been requested above.
      GenerateAndDeliverKey(storage.get());
    }
    return storage;
  }

  void CreateTestStorageOrDie(
      const StorageOptions& options,
      scoped_refptr<EncryptionModule> encryption_module =
          base::MakeRefCounted<EncryptionModule>()) {
    ASSERT_FALSE(storage_) << "StorageTest already assigned";
    StatusOr<scoped_refptr<Storage>> storage_result =
        CreateTestStorage(options, encryption_module);
    ASSERT_OK(storage_result)
        << "Failed to create StorageTest, error=" << storage_result.status();
    storage_ = std::move(storage_result.ValueOrDie());
  }

  void ResetTestStorage() {
    task_environment_.RunUntilIdle();
    storage_.reset();
    expect_to_need_key_ = false;
  }

  StorageOptions BuildTestStorageOptions() const {
    return StorageOptions()
        .set_directory(base::FilePath(location_.GetPath()))
        .set_single_file_size(::testing::get<1>(GetParam()));
  }

  StatusOr<std::unique_ptr<Storage::UploaderInterface>> BuildMockUploader(
      Priority priority) {
    auto uploader = std::make_unique<MockUploadClient>(
        &last_record_digest_map_, sequenced_task_runner_, decryptor_);
    set_mock_uploader_expectations_.Call(priority, uploader.get());
    return uploader;
  }

  Status WriteString(Priority priority, base::StringPiece data) {
    EXPECT_TRUE(storage_) << "Storage not created yet";
    TestEvent<Status> w;
    Record record;
    record.set_data(std::string(data));
    record.set_destination(UPLOAD_EVENTS);
    record.set_dm_token("DM TOKEN");
    storage_->Write(priority, std::move(record), w.cb());
    return w.result();
  }

  void WriteStringOrDie(Priority priority, base::StringPiece data) {
    const Status write_result = WriteString(priority, data);
    ASSERT_OK(write_result) << write_result;
  }

  void ConfirmOrDie(Priority priority, std::uint64_t sequencing_id) {
    TestEvent<Status> c;
    storage_->Confirm(priority, sequencing_id, c.cb());
    const Status c_result = c.result();
    ASSERT_OK(c_result) << c_result;
  }

  void GenerateAndDeliverKey(Storage* storage) {
    ASSERT_TRUE(decryptor_) << "Decryptor not created";
    // Generate new pair of private key and public value.
    uint8_t public_value[X25519_PUBLIC_VALUE_LEN];
    uint8_t private_key[X25519_PRIVATE_KEY_LEN];
    X25519_keypair(public_value, private_key);
    TestEvent<StatusOr<Encryptor::PublicKeyId>> prepare_key_pair;
    decryptor_->RecordKeyPair(
        std::string(reinterpret_cast<const char*>(private_key),
                    X25519_PRIVATE_KEY_LEN),
        std::string(reinterpret_cast<const char*>(public_value),
                    X25519_PUBLIC_VALUE_LEN),
        prepare_key_pair.cb());
    auto prepare_key_result = prepare_key_pair.result();
    ASSERT_OK(prepare_key_result.status());
    Encryptor::PublicKeyId new_public_key_id = prepare_key_result.ValueOrDie();
    // Deliver public key to storage.
    SignedEncryptionInfo signed_encryption_key;
    signed_encryption_key.set_public_asymmetric_key(std::string(
        reinterpret_cast<const char*>(public_value), X25519_PUBLIC_VALUE_LEN));
    signed_encryption_key.set_public_key_id(new_public_key_id);
    // TODO(b/170054326): Add key signature.
    storage->UpdateEncryptionKey(signed_encryption_key);
  }

  base::test::TaskEnvironment task_environment_{
      base::test::TaskEnvironment::TimeSource::MOCK_TIME};

  base::test::ScopedFeatureList scoped_feature_list_;

  base::ScopedTempDir location_;
  scoped_refptr<Decryptor> decryptor_;
  scoped_refptr<Storage> storage_;
  bool expect_to_need_key_{false};

  // Test-wide global mapping of <generation id, sequencing id> to record
  // digest. Serves all MockUploadClients created by test fixture.
  MockUploadClient::LastRecordDigestMap last_record_digest_map_;
  // Guard Access to last_record_digest_map_
  scoped_refptr<base::SequencedTaskRunner> sequenced_task_runner_{
      base::ThreadPool::CreateSequencedTaskRunner(base::TaskTraits())};

  ::testing::MockFunction<void(Priority, MockUploadClient*)>
      set_mock_uploader_expectations_;
};

constexpr std::array<const char*, 3> data = {"Rec1111", "Rec222", "Rec33"};
constexpr std::array<const char*, 3> more_data = {"More1111", "More222",
                                                  "More33"};

TEST_P(StorageTest, WriteIntoNewStorageAndReopen) {
  CreateTestStorageOrDie(BuildTestStorageOptions());
  EXPECT_CALL(set_mock_uploader_expectations_, Call(_, NotNull())).Times(0);
  WriteStringOrDie(FAST_BATCH, data[0]);
  WriteStringOrDie(FAST_BATCH, data[1]);
  WriteStringOrDie(FAST_BATCH, data[2]);

  ResetTestStorage();

  CreateTestStorageOrDie(BuildTestStorageOptions());
}

TEST_P(StorageTest, WriteIntoNewStorageReopenAndWriteMore) {
  CreateTestStorageOrDie(BuildTestStorageOptions());
  EXPECT_CALL(set_mock_uploader_expectations_, Call(_, NotNull())).Times(0);
  WriteStringOrDie(FAST_BATCH, data[0]);
  WriteStringOrDie(FAST_BATCH, data[1]);
  WriteStringOrDie(FAST_BATCH, data[2]);

  ResetTestStorage();

  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(FAST_BATCH, more_data[0]);
  WriteStringOrDie(FAST_BATCH, more_data[1]);
  WriteStringOrDie(FAST_BATCH, more_data[2]);
}

TEST_P(StorageTest, WriteIntoNewStorageAndUpload) {
  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(FAST_BATCH, data[0]);
  WriteStringOrDie(FAST_BATCH, data[1]);
  WriteStringOrDie(FAST_BATCH, data[2]);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Required(2, data[2]);
          }));

  // Trigger upload.
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));
}

TEST_P(StorageTest, WriteIntoNewStorageReopenWriteMoreAndUpload) {
  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(FAST_BATCH, data[0]);
  WriteStringOrDie(FAST_BATCH, data[1]);
  WriteStringOrDie(FAST_BATCH, data[2]);

  ResetTestStorage();

  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(FAST_BATCH, more_data[0]);
  WriteStringOrDie(FAST_BATCH, more_data[1]);
  WriteStringOrDie(FAST_BATCH, more_data[2]);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Required(2, data[2])
                .Required(3, more_data[0])
                .Required(4, more_data[1])
                .Required(5, more_data[2]);
          }));

  // Trigger upload.
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));
}

TEST_P(StorageTest, WriteIntoNewStorageAndFlush) {
  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(MANUAL_BATCH, data[0]);
  WriteStringOrDie(MANUAL_BATCH, data[1]);
  WriteStringOrDie(MANUAL_BATCH, data[2]);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_,
              Call(Eq(MANUAL_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Required(2, data[2]);
          }));

  // Trigger upload.
  EXPECT_OK(storage_->Flush(MANUAL_BATCH));
}

TEST_P(StorageTest, WriteIntoNewStorageReopenWriteMoreAndFlush) {
  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(MANUAL_BATCH, data[0]);
  WriteStringOrDie(MANUAL_BATCH, data[1]);
  WriteStringOrDie(MANUAL_BATCH, data[2]);

  ResetTestStorage();

  CreateTestStorageOrDie(BuildTestStorageOptions());
  WriteStringOrDie(MANUAL_BATCH, more_data[0]);
  WriteStringOrDie(MANUAL_BATCH, more_data[1]);
  WriteStringOrDie(MANUAL_BATCH, more_data[2]);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_,
              Call(Eq(MANUAL_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Required(2, data[2])
                .Required(3, more_data[0])
                .Required(4, more_data[1])
                .Required(5, more_data[2]);
          }));

  // Trigger upload.
  EXPECT_OK(storage_->Flush(MANUAL_BATCH));
}

TEST_P(StorageTest, WriteAndRepeatedlyUploadWithConfirmations) {
  CreateTestStorageOrDie(BuildTestStorageOptions());

  WriteStringOrDie(FAST_BATCH, data[0]);
  WriteStringOrDie(FAST_BATCH, data[1]);
  WriteStringOrDie(FAST_BATCH, data[2]);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Required(2, data[2]);
          }));

  // Forward time to trigger upload
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));

  // Confirm #0 and forward time again, removing data #0
  ConfirmOrDie(FAST_BATCH, /*sequencing_id=*/0);
  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(1, data[1])
                .Required(2, data[2]);
          }));
  // Forward time to trigger upload
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));

  // Confirm #1 and forward time again, removing data #1
  ConfirmOrDie(FAST_BATCH, /*sequencing_id=*/1);
  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(2, data[2]);
          }));
  // Forward time to trigger upload
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));

  // Add more records and verify that #2 and new records are returned.
  WriteStringOrDie(FAST_BATCH, more_data[0]);
  WriteStringOrDie(FAST_BATCH, more_data[1]);
  WriteStringOrDie(FAST_BATCH, more_data[2]);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(2, data[2])
                .Required(3, more_data[0])
                .Required(4, more_data[1])
                .Required(5, more_data[2]);
          }));
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));

  // Confirm #2 and forward time again, removing data #2
  ConfirmOrDie(FAST_BATCH, /*sequencing_id=*/2);

  // Set uploader expectations.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(3, more_data[0])
                .Required(4, more_data[1])
                .Required(5, more_data[2]);
          }));
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(1));
}

TEST_P(StorageTest, WriteAndRepeatedlyImmediateUpload) {
  CreateTestStorageOrDie(BuildTestStorageOptions());

  // Upload is initiated asynchronously, so it may happen after the next
  // record is also written. Because of that we set expectations for the
  // records after the current one as |Possible|.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Possible(1, data[1])
                .Possible(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE,
                   data[0]);  // Immediately uploads and verifies.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Possible(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE,
                   data[1]);  // Immediately uploads and verifies.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, data[0])
                .Required(1, data[1])
                .Required(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE,
                   data[2]);  // Immediately uploads and verifies.
}

TEST_P(StorageTest, WriteAndRepeatedlyImmediateUploadWithConfirmations) {
  CreateTestStorageOrDie(BuildTestStorageOptions());

  // Upload is initiated asynchronously, so it may happen after the next
  // record is also written. Because of the Confirmation below, we set
  // expectations for the records that may be eliminated by Confirmation as
  // |Possible|.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Possible(0, data[0])
                .Possible(1, data[1])
                .Possible(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, data[0]);
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Possible(0, data[0])
                .Possible(1, data[1])
                .Possible(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, data[1]);
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Possible(0, data[0])
                .Possible(1, data[1])
                .Required(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, data[2]);

  // Confirm #1, removing data #0 and #1
  ConfirmOrDie(IMMEDIATE, /*sequencing_id=*/1);

  // Add more records and verify that #2 and new records are returned.
  // Upload is initiated asynchronously, so it may happen after the next
  // record is also written. Because of that we set expectations for the
  // records after the current one as |Possible|.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(2, data[2])
                .Required(3, more_data[0])
                .Possible(4, more_data[1])
                .Possible(5, more_data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, more_data[0]);
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(2, data[2])
                .Required(3, more_data[0])
                .Required(4, more_data[1])
                .Possible(5, more_data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, more_data[1]);
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(2, data[2])
                .Required(3, more_data[0])
                .Required(4, more_data[1])
                .Required(5, more_data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, more_data[2]);
}

TEST_P(StorageTest, WriteAndRepeatedlyUploadMultipleQueues) {
  CreateTestStorageOrDie(BuildTestStorageOptions());

  // Upload is initiated asynchronously, so it may happen after the next
  // record is also written. Because of the Confirmation below, we set
  // expectations for the records that may be eliminated by Confirmation as
  // |Possible|.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Possible(0, data[0])
                .Possible(1, data[1])
                .Possible(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, data[0]);
  WriteStringOrDie(SLOW_BATCH, more_data[0]);
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Possible(0, data[0])
                .Possible(1, data[1])
                .Possible(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, data[1]);
  WriteStringOrDie(SLOW_BATCH, more_data[1]);

  // Set uploader expectations for SLOW_BATCH.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillRepeatedly(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetEmpty(priority, mock_upload_client);
          }));
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(SLOW_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Required(0, more_data[0])
                .Required(1, more_data[1]);
          }));
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(20));

  // Confirm #0 SLOW_BATCH, removing data #0
  ConfirmOrDie(SLOW_BATCH, /*sequencing_id=*/0);

  // Confirm #1 IMMEDIATE, removing data #0 and #1
  ConfirmOrDie(IMMEDIATE, /*sequencing_id=*/1);

  // Add more data
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(IMMEDIATE), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(priority, mock_upload_client)
                .Possible(1, data[1])
                .Required(2, data[2]);
          }));
  WriteStringOrDie(IMMEDIATE, data[2]);
  WriteStringOrDie(SLOW_BATCH, more_data[2]);

  // Set uploader expectations for SLOW_BATCH.
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(FAST_BATCH), NotNull()))
      .WillRepeatedly(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetEmpty(priority, mock_upload_client);
          }));
  EXPECT_CALL(set_mock_uploader_expectations_, Call(Eq(SLOW_BATCH), NotNull()))
      .WillOnce(
          Invoke([](Priority priority, MockUploadClient* mock_upload_client) {
            MockUploadClient::SetUp(SLOW_BATCH, mock_upload_client)
                .Required(1, more_data[1])
                .Required(2, more_data[2]);
          }));
  task_environment_.FastForwardBy(base::TimeDelta::FromSeconds(20));
}

TEST_P(StorageTest, WriteEncryptFailure) {
  auto test_encryption_module =
      base::MakeRefCounted<test::TestEncryptionModule>();
  CreateTestStorageOrDie(BuildTestStorageOptions(), test_encryption_module);
  EXPECT_CALL(*test_encryption_module, EncryptRecord(_, _))
      .WillOnce(WithArg<1>(
          Invoke([](base::OnceCallback<void(StatusOr<EncryptedRecord>)> cb) {
            std::move(cb).Run(Status(error::UNKNOWN, "Failing for tests"));
          })));
  const Status result = WriteString(FAST_BATCH, "TEST_MESSAGE");
  EXPECT_FALSE(result.ok());
  EXPECT_EQ(result.error_code(), error::UNKNOWN);
}

INSTANTIATE_TEST_SUITE_P(
    VaryingFileSize,
    StorageTest,
    ::testing::Combine(::testing::Bool() /* true - encryption enabled */,
                       ::testing::Values(128 * 1024LL * 1024LL,
                                         256 /* two records in file */,
                                         1 /* single record in file */)));

}  // namespace
}  // namespace reporting