Add tests for last record digest match.

Bug: b:153364303
Bug: b:153659559
Change-Id: I0a7ecedbcf043b93473334e666c93532586eca3b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2399122
Auto-Submit: Leonid Baraz <lbaraz@chromium.org>
Reviewed-by: Zach Trudo <zatrudo@google.com>
Commit-Queue: Leonid Baraz <lbaraz@chromium.org>
Cr-Commit-Position: refs/heads/master@{#806744}

chrome/browser/policy/messaging_layer/storage/storage_queue.cc

@@ -920,9 +920,7 @@ class StorageQueue::WriteContext : public TaskRunnerContext<Status> {
     }
     scoped_refptr<SingleFile> last_file = assign_result.ValueOrDie();
 
-    // Prepare and initiate writing generation into to a new file.
-    // Pick up when both WriteMetadata and WriteHeaderAndBlock
-    // have finished.
+    // Writing metadata ahead of the data write.
     Status write_result = storage_queue_->WriteMetadata();
     if (!write_result.ok()) {
       Response(write_result);

chrome/browser/policy/messaging_layer/storage/storage_queue_unittest.cc

@@ -76,7 +76,16 @@ class TestEvent {
 
 class MockUploadClient : public StorageQueue::UploaderInterface {
  public:
-  MockUploadClient() = default;
+  // Mapping of (generation, seq number) 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::pair<uint64_t /*generation */, uint64_t /*sequencing number*/>,
+      std::string /*digest*/>;
+
+  explicit MockUploadClient(LastRecordDigestMap* last_record_digest_map)
+      : last_record_digest_map_(last_record_digest_map) {}
 
   void ProcessRecord(StatusOr<EncryptedRecord> encrypted_record,
                      base::OnceCallback<void(bool)> processed_cb) override {
@@ -89,25 +98,22 @@ class MockUploadClient : public StorageQueue::UploaderInterface {
     ASSERT_TRUE(wrapped_record.ParseFromString(
         encrypted_record.ValueOrDie().encrypted_wrapped_record()));
     // Verify generation match.
+    const auto& sequencing_information =
+        encrypted_record.ValueOrDie().sequencing_information();
     if (generation_id_.has_value() &&
-        generation_id_.value() != encrypted_record.ValueOrDie()
-                                      .sequencing_information()
-                                      .generation_id()) {
+        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(encrypted_record.ValueOrDie()
-                                                     .sequencing_information()
-                                                     .generation_id())}))));
+                            base::NumberToString(
+                                sequencing_information.generation_id())}))));
       return;
     }
     if (!generation_id_.has_value()) {
-      generation_id_ = encrypted_record.ValueOrDie()
-                           .sequencing_information()
-                           .generation_id();
+      generation_id_ = sequencing_information.generation_id();
     }
 
     // Verify digest and its match.
@@ -124,12 +130,26 @@ class MockUploadClient : public StorageQueue::UploaderInterface {
                 Status(error::DATA_LOSS, "Record digest mismatch")));
         return;
       }
+      if (wrapped_record.has_last_record_digest()) {
+        auto it = last_record_digest_map_->find(
+            std::make_pair(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_pair(sequencing_information.sequencing_id(),
+                         sequencing_information.generation_id()),
+          record_digest);
     }
 
     std::move(processed_cb)
-        .Run(UploadRecord(encrypted_record.ValueOrDie()
-                              .sequencing_information()
-                              .sequencing_id(),
+        .Run(UploadRecord(sequencing_information.sequencing_id(),
                           wrapped_record.record().data()));
   }
 
@@ -176,6 +196,8 @@ class MockUploadClient : public StorageQueue::UploaderInterface {
 
  private:
   base::Optional<uint64_t> generation_id_;
+  LastRecordDigestMap* const last_record_digest_map_;
+
   Sequence test_upload_sequence_;
 };
 
@@ -218,7 +240,8 @@ class StorageQueueTest : public ::testing::TestWithParam<size_t> {
 
   StatusOr<std::unique_ptr<StorageQueue::UploaderInterface>>
   BuildMockUploader() {
-    auto uploader = std::make_unique<MockUploadClient>();
+    auto uploader =
+        std::make_unique<MockUploadClient>(&last_record_digest_map_);
     set_mock_uploader_expectations_.Call(uploader.get());
     return uploader;
   }
@@ -250,6 +273,10 @@ class StorageQueueTest : public ::testing::TestWithParam<size_t> {
   scoped_refptr<test::TestEncryptionModule> test_encryption_module_;
   scoped_refptr<StorageQueue> storage_queue_;
 
+  // Test-wide global mapping of seq number to record digest.
+  // Serves all MockUploadClients created by test fixture.
+  MockUploadClient::LastRecordDigestMap last_record_digest_map_;
+
   ::testing::MockFunction<void(MockUploadClient*)>
       set_mock_uploader_expectations_;
 

chrome/browser/policy/messaging_layer/storage/storage_unittest.cc

@@ -5,8 +5,8 @@
 #include "chrome/browser/policy/messaging_layer/storage/storage.h"
 
 #include <cstdint>
+#include <tuple>
 #include <utility>
-#include <vector>
 
 #include "base/files/scoped_temp_dir.h"
 #include "base/optional.h"
@@ -78,7 +78,18 @@ class TestEvent {
 
 class MockUploadClient : public Storage::UploaderInterface {
  public:
-  MockUploadClient() = default;
+  // Mapping of (generation, seq number) 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 */,
+                          uint64_t /*sequencing number*/>,
+               std::string /*digest*/>;
+
+  explicit MockUploadClient(LastRecordDigestMap* last_record_digest_map)
+      : last_record_digest_map_(last_record_digest_map) {}
 
   void ProcessRecord(StatusOr<EncryptedRecord> encrypted_record,
                      base::OnceCallback<void(bool)> processed_cb) override {
@@ -91,25 +102,22 @@ class MockUploadClient : public Storage::UploaderInterface {
     ASSERT_TRUE(wrapped_record.ParseFromString(
         encrypted_record.ValueOrDie().encrypted_wrapped_record()));
     // Verify generation match.
+    const auto& sequencing_information =
+        encrypted_record.ValueOrDie().sequencing_information();
     if (generation_id_.has_value() &&
-        generation_id_.value() != encrypted_record.ValueOrDie()
-                                      .sequencing_information()
-                                      .generation_id()) {
+        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(encrypted_record.ValueOrDie()
-                                                     .sequencing_information()
-                                                     .generation_id())}))));
+                            base::NumberToString(
+                                sequencing_information.generation_id())}))));
       return;
     }
     if (!generation_id_.has_value()) {
-      generation_id_ = encrypted_record.ValueOrDie()
-                           .sequencing_information()
-                           .generation_id();
+      generation_id_ = sequencing_information.generation_id();
     }
 
     // Verify digest and its match.
@@ -126,15 +134,30 @@ class MockUploadClient : public Storage::UploaderInterface {
                 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(
-            encrypted_record.ValueOrDie().sequencing_information().priority(),
-            encrypted_record.ValueOrDie()
-                .sequencing_information()
-                .sequencing_id(),
-            wrapped_record.record().data()));
+        .Run(UploadRecord(sequencing_information.priority(),
+                          sequencing_information.sequencing_id(),
+                          wrapped_record.record().data()));
   }
 
   void Completed(Status status) override { UploadComplete(status); }
@@ -205,6 +228,8 @@ class MockUploadClient : public Storage::UploaderInterface {
 
  private:
   base::Optional<uint64_t> generation_id_;
+  LastRecordDigestMap* const last_record_digest_map_;
+
   Sequence test_upload_sequence_;
 };
 
@@ -234,14 +259,15 @@ class StorageTest : public ::testing::Test {
 
   StatusOr<std::unique_ptr<Storage::UploaderInterface>> BuildMockUploader(
       Priority priority) {
-    auto uploader = std::make_unique<MockUploadClient>();
+    auto uploader =
+        std::make_unique<MockUploadClient>(&last_record_digest_map_);
     set_mock_uploader_expectations_.Call(priority, uploader.get());
     return uploader;
   }
 
   Status WriteString(Priority priority, base::StringPiece data) {
-    TestEvent<Status> w;
     EXPECT_TRUE(storage_) << "Storage not created yet";
+    TestEvent<Status> w;
     Record record;
     record.set_data(std::string(data));
     record.set_destination(UPLOAD_EVENTS);
@@ -266,6 +292,10 @@ class StorageTest : public ::testing::Test {
   scoped_refptr<test::TestEncryptionModule> test_encryption_module_;
   scoped_refptr<Storage> storage_;
 
+  // Test-wide global mapping of seq number to record digest.
+  // Serves all MockUploadClients created by test fixture.
+  MockUploadClient::LastRecordDigestMap last_record_digest_map_;
+
   ::testing::MockFunction<void(Priority, MockUploadClient*)>
       set_mock_uploader_expectations_;