Add need_encryption_keys to UploadClient.

Added has_encryption_key to relevant modules.

Bug: b:170054326
Change-Id: Ic1a0b8cc548e5c29015a4fd604bcc3138092d867
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2581006
Commit-Queue: Leonid Baraz <lbaraz@chromium.org>
Reviewed-by: Zach Trudo <zatrudo@google.com>
Cr-Commit-Position: refs/heads/master@{#835787}

chrome/browser/policy/messaging_layer/encryption/encryption_module.cc

@@ -4,6 +4,8 @@
 
 #include "chrome/browser/policy/messaging_layer/encryption/encryption_module.h"
 
+#include <atomic>
+
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/feature_list.h"
@@ -90,8 +92,21 @@ void EncryptionModule::UpdateAsymmetricKey(
     base::StringPiece new_public_key,
     Encryptor::PublicKeyId new_public_key_id,
     base::OnceCallback<void(Status)> response_cb) {
-  encryptor_->UpdateAsymmetricKey(new_public_key, new_public_key_id,
-                                  std::move(response_cb));
+  encryptor_->UpdateAsymmetricKey(
+      new_public_key, new_public_key_id,
+      base::BindOnce(
+          [](EncryptionModule* encryption_module,
+             base::OnceCallback<void(Status)> response_cb, Status status) {
+            if (status.ok()) {
+              encryption_module->has_encryption_key_.store(true);
+            }
+            std::move(response_cb).Run(status);
+          },
+          base::Unretained(this), std::move(response_cb)));
+}
+
+bool EncryptionModule::has_encryption_key() const {
+  return has_encryption_key_.load();
 }
 
 }  // namespace reporting

chrome/browser/policy/messaging_layer/encryption/encryption_module.h

@@ -5,6 +5,8 @@
 #ifndef CHROME_BROWSER_POLICY_MESSAGING_LAYER_ENCRYPTION_ENCRYPTION_MODULE_H_
 #define CHROME_BROWSER_POLICY_MESSAGING_LAYER_ENCRYPTION_ENCRYPTION_MODULE_H_
 
+#include <atomic>
+
 #include "base/callback.h"
 #include "base/memory/ref_counted.h"
 #include "base/strings/string_piece.h"
@@ -39,12 +41,23 @@ class EncryptionModule : public base::RefCountedThreadSafe<EncryptionModule> {
       Encryptor::PublicKeyId new_public_key_id,
       base::OnceCallback<void(Status)> response_cb);
 
+  // Returns `false` if encryption key has not been set yet, and `true`
+  // otherwise. The result is lazy: the method may return `false` for some time
+  // even after the key has already been set - this is harmless, since resetting
+  // or even changing the key is OK at any time.
+  bool has_encryption_key() const;
+
  protected:
   virtual ~EncryptionModule();
 
  private:
   friend base::RefCountedThreadSafe<EncryptionModule>;
 
+  // Lazy flag indicating public assymmetric key has been set.
+  // Initialized as `false`, set to `true` after |UpdateAsymmetricKey| is called
+  // for the first time.
+  std::atomic<bool> has_encryption_key_{false};
+
   // Encryptor.
   scoped_refptr<Encryptor> encryptor_;
 };

chrome/browser/policy/messaging_layer/public/report_client.cc

@@ -571,7 +571,8 @@ ReportingClient::BuildUploader(Priority priority) {
   DCHECK(instance->upload_client_);
   return Uploader::Create(
       base::BindOnce(&UploadClient::EnqueueUpload,
-                     base::Unretained(instance->upload_client_.get())));
+                     base::Unretained(instance->upload_client_.get()),
+                     !instance->config_->storage->has_encryption_key()));
 }
 
 ReportingClient::TestEnvironment::TestEnvironment(

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

@@ -150,7 +150,6 @@ void Storage::Create(
    public:
     StorageInitContext(
         const std::vector<std::pair<Priority, QueueOptions>>& queues_options,
-        scoped_refptr<EncryptionModule> encryption_module,
         scoped_refptr<Storage> storage,
         base::OnceCallback<void(StatusOr<scoped_refptr<Storage>>)> callback)
         : TaskRunnerContext<StatusOr<scoped_refptr<Storage>>>(
@@ -158,7 +157,6 @@ void Storage::Create(
               base::ThreadPool::CreateSequencedTaskRunner(
                   {base::TaskPriority::BEST_EFFORT, base::MayBlock()})),
           queues_options_(queues_options),
-          encryption_module_(encryption_module),
           storage_(std::move(storage)),
           count_(queues_options_.size()) {}
 
@@ -174,7 +172,7 @@ void Storage::Create(
             base::BindRepeating(&QueueUploaderInterface::ProvideUploader,
                                 /*priority=*/queue_options.first,
                                 storage_->start_upload_cb_),
-            encryption_module_,
+            storage_->encryption_module_,
             base::BindOnce(&StorageInitContext::ScheduleAddQueue,
                            base::Unretained(this),
                            /*priority=*/queue_options.first));
@@ -214,7 +212,6 @@ void Storage::Create(
     }
 
     const std::vector<std::pair<Priority, QueueOptions>> queues_options_;
-    scoped_refptr<EncryptionModule> encryption_module_;
     scoped_refptr<Storage> storage_;
     int32_t count_;
     Status final_status_;
@@ -222,17 +219,20 @@ void Storage::Create(
 
   // Create Storage object.
   // Cannot use base::MakeRefCounted<Storage>, because constructor is private.
-  scoped_refptr<Storage> storage =
-      base::WrapRefCounted(new Storage(options, std::move(start_upload_cb)));
+  scoped_refptr<Storage> storage = base::WrapRefCounted(
+      new Storage(options, encryption_module, std::move(start_upload_cb)));
 
   // Asynchronously run initialization.
   Start<StorageInitContext>(ExpectedQueues(storage->options_),
-                            encryption_module, std::move(storage),
-                            std::move(completion_cb));
+                            std::move(storage), std::move(completion_cb));
 }
 
-Storage::Storage(const StorageOptions& options, StartUploadCb start_upload_cb)
-    : options_(options), start_upload_cb_(std::move(start_upload_cb)) {}
+Storage::Storage(const StorageOptions& options,
+                 scoped_refptr<EncryptionModule> encryption_module,
+                 StartUploadCb start_upload_cb)
+    : options_(options),
+      encryption_module_(encryption_module),
+      start_upload_cb_(std::move(start_upload_cb)) {}
 
 Storage::~Storage() = default;
 
@@ -264,6 +264,10 @@ Status Storage::Flush(Priority priority) {
   return Status::StatusOK();
 }
 
+bool Storage::has_encryption_key() const {
+  return !encryption_module_->has_encryption_key();
+}
+
 StatusOr<scoped_refptr<StorageQueue>> Storage::GetQueue(Priority priority) {
   auto it = queues_.find(priority);
   if (it == queues_.end()) {

chrome/browser/policy/messaging_layer/storage/storage.h

@@ -65,6 +65,13 @@ class Storage : public base::RefCountedThreadSafe<Storage> {
   // Returns error if cannot start upload.
   Status Flush(Priority priority);
 
+  // Returns `false` if encryption key has not been found in the Storage during
+  // initialization and not received from the server yet, and `true` otherwise.
+  // The result is lazy: the method may return `false` for some time even after
+  // the key has already been set - this is harmless, since resetting or even
+  // changing the key is OK at any time.
+  bool has_encryption_key() const;
+
   Storage(const Storage& other) = delete;
   Storage& operator=(const Storage& other) = delete;
 
@@ -79,7 +86,9 @@ class Storage : public base::RefCountedThreadSafe<Storage> {
 
   // Private constructor, to be called by Create factory method only.
   // Queues need to be added afterwards.
-  Storage(const StorageOptions& options, StartUploadCb start_upload_cb);
+  Storage(const StorageOptions& options,
+          scoped_refptr<EncryptionModule> encryption_module,
+          StartUploadCb start_upload_cb);
 
   // Initializes the object by adding all queues for all priorities.
   // Must be called once and only once after construction.
@@ -92,8 +101,12 @@ class Storage : public base::RefCountedThreadSafe<Storage> {
   // need to protect or serialize access to it.
   StatusOr<scoped_refptr<StorageQueue>> GetQueue(Priority priority);
 
+  // Immutable options, stored at the time of creation.
   const StorageOptions options_;
 
+  // Encryption module.
+  scoped_refptr<EncryptionModule> encryption_module_;
+
   // Map priority->StorageQueue.
   base::flat_map<Priority, scoped_refptr<StorageQueue>> queues_;
 

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

@@ -67,4 +67,8 @@ void StorageModule::Create(
           std::move(instance), std::move(callback)));
 }
 
+bool StorageModule::has_encryption_key() const {
+  return storage_->has_encryption_key();
+}
+
 }  // namespace reporting

chrome/browser/policy/messaging_layer/storage/storage_module.h

@@ -44,6 +44,13 @@ class StorageModule : public base::RefCountedThreadSafe<StorageModule> {
   // can be passed back to the StorageModule here for record deletion.
   virtual void ReportSuccess(SequencingInformation sequencing_information);
 
+  // Returns `false` if encryption key has not been found in the Storage during
+  // initialization and not received from the server yet, and `true` otherwise.
+  // The result is lazy: the method may return `false` for some time even after
+  // the key has already been set - this is harmless, since resetting or even
+  // changing the key is OK at any time.
+  virtual bool has_encryption_key() const;
+
  protected:
   // Constructor can only be called by |Create| factory method.
   StorageModule();

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

@@ -43,5 +43,10 @@ void TestStorageModuleStrict::AddRecordSuccessfully(
   priority_ = priority;
   std::move(callback).Run(Status::StatusOK());
 }
+
+bool TestStorageModuleStrict::has_encryption_key() const {
+  return false;  // No encryption in test module.
+}
+
 }  // namespace test
 }  // namespace reporting

chrome/browser/policy/messaging_layer/storage/test_storage_module.h

@@ -34,6 +34,8 @@ class TestStorageModuleStrict : public StorageModule {
   Record record() const;
   Priority priority() const;
 
+  bool has_encryption_key() const override;
+
  protected:
   ~TestStorageModuleStrict() override;
 

chrome/browser/policy/messaging_layer/upload/dm_server_upload_service.cc

@@ -27,10 +27,6 @@
 
 namespace reporting {
 
-}  // namespace reporting
-
-namespace reporting {
-
 using DmServerUploader = DmServerUploadService::DmServerUploader;
 using ::policy::CloudPolicyClient;
 
@@ -73,12 +69,14 @@ DmServerUploadService::RecordHandler::RecordHandler(
     : client_(client) {}
 
 DmServerUploader::DmServerUploader(
+    bool need_encryption_key,
     std::unique_ptr<std::vector<EncryptedRecord>> records,
     RecordHandler* handler,
     CompletionCallback completion_cb,
     scoped_refptr<base::SequencedTaskRunner> sequenced_task_runner)
     : TaskRunnerContext<CompletionResponse>(std::move(completion_cb),
                                             sequenced_task_runner),
+      need_encryption_key_(need_encryption_key),
       encrypted_records_(std::move(records)),
       handler_(handler) {
   DETACH_FROM_SEQUENCE(sequence_checker_);
@@ -103,7 +101,6 @@ void DmServerUploader::OnStart() {
   ProcessRecords();
 }
 
-
 void DmServerUploader::ProcessRecords() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   Status process_status;
@@ -140,7 +137,7 @@ void DmServerUploader::ProcessRecords() {
 void DmServerUploader::HandleRecords() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   handler_->HandleRecords(
-      std::move(encrypted_records_),
+      need_encryption_key_, std::move(encrypted_records_),
       base::BindOnce(&DmServerUploader::Complete, base::Unretained(this)));
 }
 
@@ -194,9 +191,10 @@ DmServerUploadService::DmServerUploadService(
 DmServerUploadService::~DmServerUploadService() = default;
 
 Status DmServerUploadService::EnqueueUpload(
+    bool need_encryption_key,
     std::unique_ptr<std::vector<EncryptedRecord>> records) {
   Start<DmServerUploader>(
-      std::move(records), handler_.get(),
+      need_encryption_key, std::move(records), handler_.get(),
       base::BindOnce(&DmServerUploadService::UploadCompletion,
                      base::Unretained(this)),
       sequenced_task_runner_);

chrome/browser/policy/messaging_layer/upload/dm_server_upload_service.h

@@ -54,11 +54,17 @@ class DmServerUploadService {
     // Will iterate over |records| and ensure they are in ascending sequence
     // order, and within the same generation. Any out of order records will be
     // discarded.
+    // |need_encryption_key| is set to `true` if the client needs to request
+    // the encryption key from the server (either because it does not have it
+    // or because the one it has is old and may be outdated). In that case
+    // it is ok for |records| to be empty (otherwise at least one record must
+    // be present).
     // Once the server has responded |upload_complete| is called with either the
     // highest accepted SequencingInformation, or an error detailing the failure
     // cause.
     // Any errors will result in |upload_complete| being called with a Status.
     virtual void HandleRecords(
+        bool need_encryption_key,
         std::unique_ptr<std::vector<EncryptedRecord>> records,
         DmServerUploadService::CompletionCallback upload_complete) = 0;
 
@@ -76,6 +82,7 @@ class DmServerUploadService {
   class DmServerUploader : public TaskRunnerContext<CompletionResponse> {
    public:
     DmServerUploader(
+        bool need_encryption_key,
         std::unique_ptr<std::vector<EncryptedRecord>> records,
         RecordHandler* handler,
         CompletionCallback completion_cb,
@@ -120,6 +127,7 @@ class DmServerUploadService {
         base::RepeatingClosure done_cb,
         SequencingInformation sequencing_information);
 
+    const bool need_encryption_key_;
     std::unique_ptr<std::vector<EncryptedRecord>> encrypted_records_;
     RecordHandler* handler_;
 
@@ -143,7 +151,8 @@ class DmServerUploadService {
           created_cb);
   ~DmServerUploadService();
 
-  Status EnqueueUpload(std::unique_ptr<std::vector<EncryptedRecord>> record);
+  Status EnqueueUpload(bool need_encryption_key,
+                       std::unique_ptr<std::vector<EncryptedRecord>> record);
 
  private:
   DmServerUploadService(policy::CloudPolicyClient* client,

chrome/browser/policy/messaging_layer/upload/dm_server_upload_service_unittest.cc

@@ -124,14 +124,16 @@ class TestRecordHandler : public DmServerUploadService::RecordHandler {
   ~TestRecordHandler() override = default;
 
   void HandleRecords(
+      bool need_encryption_key,
       std::unique_ptr<std::vector<EncryptedRecord>> records,
       DmServerUploadService::CompletionCallback upload_complete) override {
-    HandleRecords_(records, upload_complete);
+    HandleRecords_(need_encryption_key, records, upload_complete);
   }
 
   MOCK_METHOD(void,
               HandleRecords_,
-              (std::unique_ptr<std::vector<EncryptedRecord>>&,
+              (bool,
+               std::unique_ptr<std::vector<EncryptedRecord>>&,
                DmServerUploadService::CompletionCallback&));
 };
 
@@ -158,8 +160,8 @@ TEST_F(DmServerUploaderTest, ProcessesRecord) {
   // Add an empty record.
   records_->emplace_back();
 
-  EXPECT_CALL(*handler_, HandleRecords_(_, _))
-      .WillOnce(WithArgs<1>(
+  EXPECT_CALL(*handler_, HandleRecords_(_, _, _))
+      .WillOnce(WithArgs<2>(
           Invoke([](DmServerUploadService::CompletionCallback& callback) {
             std::move(callback).Run(SequencingInformation());
           })));
@@ -169,9 +171,9 @@ TEST_F(DmServerUploaderTest, ProcessesRecord) {
       base::BindOnce(&TestCallbackWaiter::CompleteExpectSuccess,
                      base::Unretained(&callback_waiter));
 
-  Start<DmServerUploadService::DmServerUploader>(std::move(records_),
-                                                 handler_.get(), std::move(cb),
-                                                 sequenced_task_runner_);
+  Start<DmServerUploadService::DmServerUploader>(
+      /*need_encryption_key=*/false, std::move(records_), handler_.get(),
+      std::move(cb), sequenced_task_runner_);
 
   callback_waiter.Wait();
 }
@@ -192,8 +194,8 @@ TEST_F(DmServerUploaderTest, ProcessesRecords) {
     records_->push_back(std::move(encrypted_record));
   }
 
-  EXPECT_CALL(*handler_, HandleRecords_(_, _))
-      .WillOnce(WithArgs<1>(
+  EXPECT_CALL(*handler_, HandleRecords_(_, _, _))
+      .WillOnce(WithArgs<2>(
           Invoke([](DmServerUploadService::CompletionCallback& callback) {
             std::move(callback).Run(SequencingInformation());
           })));
@@ -203,9 +205,9 @@ TEST_F(DmServerUploaderTest, ProcessesRecords) {
       base::BindOnce(&TestCallbackWaiter::CompleteExpectSuccess,
                      base::Unretained(&callback_waiter));
 
-  Start<DmServerUploadService::DmServerUploader>(std::move(records_),
-                                                 handler_.get(), std::move(cb),
-                                                 sequenced_task_runner_);
+  Start<DmServerUploadService::DmServerUploader>(
+      /*need_encryption_key=*/false, std::move(records_), handler_.get(),
+      std::move(cb), sequenced_task_runner_);
 
   callback_waiter.Wait();
 }
@@ -214,8 +216,8 @@ TEST_F(DmServerUploaderTest, ReportsFailureToProcess) {
   // Add an empty record.
   records_->emplace_back();
 
-  EXPECT_CALL(*handler_, HandleRecords_(_, _))
-      .WillOnce(WithArgs<1>(
+  EXPECT_CALL(*handler_, HandleRecords_(_, _, _))
+      .WillOnce(WithArgs<2>(
           Invoke([](DmServerUploadService::CompletionCallback& callback) {
             std::move(callback).Run(
                 Status(error::FAILED_PRECONDITION, "Fail for test"));
@@ -226,9 +228,9 @@ TEST_F(DmServerUploaderTest, ReportsFailureToProcess) {
       base::BindOnce(&TestCallbackWaiter::CompleteExpectFailedPrecondition,
                      base::Unretained(&callback_waiter));
 
-  Start<DmServerUploadService::DmServerUploader>(std::move(records_),
-                                                 handler_.get(), std::move(cb),
-                                                 sequenced_task_runner_);
+  Start<DmServerUploadService::DmServerUploader>(
+      /*need_encryption_key=*/false, std::move(records_), handler_.get(),
+      std::move(cb), sequenced_task_runner_);
 
   callback_waiter.Wait();
 }
@@ -237,8 +239,8 @@ TEST_F(DmServerUploaderTest, ReportsFailureToUpload) {
   // Add an empty record.
   records_->emplace_back();
 
-  EXPECT_CALL(*handler_, HandleRecords_(_, _))
-      .WillOnce(WithArgs<1>(
+  EXPECT_CALL(*handler_, HandleRecords_(_, _, _))
+      .WillOnce(WithArgs<2>(
           Invoke([](DmServerUploadService::CompletionCallback& callback) {
             std::move(callback).Run(
                 Status(error::DEADLINE_EXCEEDED, "Fail for test"));
@@ -249,9 +251,9 @@ TEST_F(DmServerUploaderTest, ReportsFailureToUpload) {
       base::BindOnce(&TestCallbackWaiter::CompleteExpectDeadlineExceeded,
                      base::Unretained(&callback_waiter));
 
-  Start<DmServerUploadService::DmServerUploader>(std::move(records_),
-                                                 handler_.get(), std::move(cb),
-                                                 sequenced_task_runner_);
+  Start<DmServerUploadService::DmServerUploader>(
+      /*need_encryption_key=*/false, std::move(records_), handler_.get(),
+      std::move(cb), sequenced_task_runner_);
 
   callback_waiter.Wait();
 }
@@ -262,9 +264,9 @@ TEST_F(DmServerUploaderTest, FailWithZeroRecords) {
       base::BindOnce(&TestCallbackWaiter::CompleteExpectInvalidArgument,
                      base::Unretained(&callback_waiter));
 
-  Start<DmServerUploadService::DmServerUploader>(std::move(records_),
-                                                 handler_.get(), std::move(cb),
-                                                 sequenced_task_runner_);
+  Start<DmServerUploadService::DmServerUploader>(
+      /*need_encryption_key=*/false, std::move(records_), handler_.get(),
+      std::move(cb), sequenced_task_runner_);
 
   callback_waiter.Wait();
 }

chrome/browser/policy/messaging_layer/upload/record_handler_impl.cc

@@ -39,6 +39,7 @@ class RecordHandlerImpl::ReportUploader
     : public TaskRunnerContext<DmServerUploadService::CompletionResponse> {
  public:
   ReportUploader(
+      bool need_encryption_key,
       std::unique_ptr<std::vector<EncryptedRecord>> records,
       policy::CloudPolicyClient* client,
       DmServerUploadService::CompletionCallback upload_complete_cb,
@@ -49,13 +50,21 @@ class RecordHandlerImpl::ReportUploader
 
   void OnStart() override;
 
-  void StartUpload(const EncryptedRecord& encrypted_record);
+  void StartUpload(bool need_encryption_key,
+                   const EncryptedRecord& encrypted_record);
   void OnUploadComplete(base::Optional<base::Value> response);
   void HandleFailedUpload();
   void HandleSuccessfulUpload();
 
+  // Populates upload request. Returns JSON request base::Value or nullopt,
+  // if an error was detected.
+  base::Optional<base::Value> PopulateRequest(
+      bool need_encryption_key,
+      const EncryptedRecord& encrypted_record);
+
   void Complete(DmServerUploadService::CompletionResponse result);
 
+  const bool need_encryption_key_;
   std::unique_ptr<std::vector<EncryptedRecord>> records_;
   policy::CloudPolicyClient* client_;
 
@@ -70,6 +79,7 @@ class RecordHandlerImpl::ReportUploader
 };
 
 RecordHandlerImpl::ReportUploader::ReportUploader(
+    bool need_encryption_key,
     std::unique_ptr<std::vector<EncryptedRecord>> records,
     policy::CloudPolicyClient* client,
     DmServerUploadService::CompletionCallback client_cb,
@@ -77,6 +87,7 @@ RecordHandlerImpl::ReportUploader::ReportUploader(
     : TaskRunnerContext<DmServerUploadService::CompletionResponse>(
           std::move(client_cb),
           sequenced_task_runner),
+      need_encryption_key_(need_encryption_key),
       records_(std::move(records)),
       client_(client) {}
 
@@ -108,10 +119,11 @@ void RecordHandlerImpl::ReportUploader::OnStart() {
   // We'll be popping records off the back.
   std::reverse(records_->begin(), records_->end());
 
-  StartUpload(records_->back());
+  StartUpload(need_encryption_key_, records_->back());
 }
 
 void RecordHandlerImpl::ReportUploader::StartUpload(
+    bool need_encryption_key,
     const EncryptedRecord& encrypted_record) {
   auto response_cb =
       base::BindOnce(&RecordHandlerImpl::ReportUploader::OnUploadComplete,
@@ -212,7 +224,8 @@ void RecordHandlerImpl::ReportUploader::HandleSuccessfulUpload() {
     return;
   }
 
-  StartUpload(records_->back());
+  // Upload the next record but do not request encryption key again.
+  StartUpload(/*need_encryption_key=*/false, records_->back());
 }
 
 void RecordHandlerImpl::ReportUploader::Complete(
@@ -228,11 +241,12 @@ RecordHandlerImpl::RecordHandlerImpl(policy::CloudPolicyClient* client)
 RecordHandlerImpl::~RecordHandlerImpl() = default;
 
 void RecordHandlerImpl::HandleRecords(
+    bool need_encryption_key,
     std::unique_ptr<std::vector<EncryptedRecord>> records,
     DmServerUploadService::CompletionCallback upload_complete_cb) {
-  Start<RecordHandlerImpl::ReportUploader>(std::move(records), GetClient(),
-                                           std::move(upload_complete_cb),
-                                           sequenced_task_runner_);
+  Start<RecordHandlerImpl::ReportUploader>(
+      need_encryption_key, std::move(records), GetClient(),
+      std::move(upload_complete_cb), sequenced_task_runner_);
 }
 
 }  // namespace reporting

chrome/browser/policy/messaging_layer/upload/record_handler_impl.h

@@ -10,6 +10,7 @@
 
 #include "base/callback.h"
 #include "base/memory/ref_counted.h"
+#include "base/optional.h"
 #include "base/sequenced_task_runner.h"
 #include "base/task/post_task.h"
 #include "base/task_runner.h"
@@ -37,6 +38,7 @@ class RecordHandlerImpl : public DmServerUploadService::RecordHandler {
 
   // Base class RecordHandler method implementation.
   void HandleRecords(
+      bool need_encryption_key,
       std::unique_ptr<std::vector<EncryptedRecord>> record,
       DmServerUploadService::CompletionCallback upload_complete) override;
 

chrome/browser/policy/messaging_layer/upload/record_handler_impl_unittest.cc

@@ -164,7 +164,8 @@ TEST_F(RecordHandlerImplTest, ForwardsRecordsToCloudPolicyClient) {
   auto responder_callback = base::BindOnce(
       &TestCompletionResponder::RecordsHandled, base::Unretained(&responder));
 
-  handler.HandleRecords(std::move(test_records), std::move(responder_callback));
+  handler.HandleRecords(/*need_encryption_key=*/false, std::move(test_records),
+                        std::move(responder_callback));
 
   client_waiter.Wait();
   responder_waiter.Wait();
@@ -211,7 +212,8 @@ TEST_F(RecordHandlerImplTest, ReportsEarlyFailure) {
   auto responder_callback = base::BindOnce(
       &TestCompletionResponder::RecordsHandled, base::Unretained(&responder));
 
-  handler.HandleRecords(std::move(test_records), std::move(responder_callback));
+  handler.HandleRecords(/*need_encryption_key=*/false, std::move(test_records),
+                        std::move(responder_callback));
 
   client_waiter.Wait();
   responder_waiter.Wait();

chrome/browser/policy/messaging_layer/upload/record_upload_request_builder.cc

@@ -21,6 +21,7 @@ namespace {
 
 // UploadEncryptedReportingRequestBuilder list key
 constexpr char kEncryptedRecordListKey[] = "encryptedRecord";
+constexpr char kAttachEncryptionSettingsKey[] = "attachEncryptionSettings";
 
 // EncrypedRecordDictionaryBuilder strings
 constexpr char kEncryptedWrappedRecord[] = "encryptedWrappedRecord";
@@ -38,11 +39,14 @@ constexpr char kPublicKeyId[] = "publicKeyId";
 
 }  // namespace
 
-UploadEncryptedReportingRequestBuilder::
-    UploadEncryptedReportingRequestBuilder() {
+UploadEncryptedReportingRequestBuilder::UploadEncryptedReportingRequestBuilder(
+    bool attach_encryption_settings) {
   result_ = base::Value{base::Value::Type::DICTIONARY};
   result_.value().SetKey(GetEncryptedRecordListPath(),
                          base::Value{base::Value::Type::LIST});
+  if (attach_encryption_settings) {
+    result_.value().SetBoolKey(GetAttachEncryptionSettingsPath(), true);
+  }
 }
 
 UploadEncryptedReportingRequestBuilder::
@@ -83,6 +87,12 @@ UploadEncryptedReportingRequestBuilder::GetEncryptedRecordListPath() {
   return kEncryptedRecordListKey;
 }
 
+// static
+base::StringPiece
+UploadEncryptedReportingRequestBuilder::GetAttachEncryptionSettingsPath() {
+  return kAttachEncryptionSettingsKey;
+}
+
 EncryptedRecordDictionaryBuilder::EncryptedRecordDictionaryBuilder(
     const EncryptedRecord& record) {
   base::Value record_dictionary{base::Value::Type::DICTIONARY};

chrome/browser/policy/messaging_layer/upload/record_upload_request_builder.h

@@ -44,8 +44,8 @@ namespace reporting {
 //       }
 //     }
 //   ]
+//   "attachEncryptionSettings": true  // optional field
 // }
-// TODO(b/170054326): Add 'need_encryption_key' field, if requested.
 // TODO(b/159361496): Periodically add memory and disk space usage.
 //
 // This payload is added to the common payload of all reporting jobs, which
@@ -62,7 +62,8 @@ namespace reporting {
 
 class UploadEncryptedReportingRequestBuilder {
  public:
-  UploadEncryptedReportingRequestBuilder();
+  explicit UploadEncryptedReportingRequestBuilder(
+      bool attach_encryption_settings = false);
   ~UploadEncryptedReportingRequestBuilder();
 
   UploadEncryptedReportingRequestBuilder& AddRecord(
@@ -71,6 +72,7 @@ class UploadEncryptedReportingRequestBuilder {
   base::Optional<base::Value> Build();
 
   static base::StringPiece GetEncryptedRecordListPath();
+  static base::StringPiece GetAttachEncryptionSettingsPath();
 
   static const char kEncryptedRecordListKey_[];
 

chrome/browser/policy/messaging_layer/upload/record_upload_request_builder_unittest.cc

@@ -25,7 +25,7 @@ constexpr Priority kPriority = Priority::IMMEDIATE;
 constexpr char kEncryptionKey[] = "abcdef";
 constexpr uint64_t kPublicKeyId = 9876;
 
-class RecordUploadRequestBuilderTest : public ::testing::Test {
+class RecordUploadRequestBuilderTest : public ::testing::TestWithParam<bool> {
  public:
   RecordUploadRequestBuilderTest() = default;
 
@@ -52,9 +52,11 @@ class RecordUploadRequestBuilderTest : public ::testing::Test {
     static uint64_t sequencing_id = 0;
     return sequencing_id++;
   }
+
+  bool need_encryption_key() const { return GetParam(); }
 };
 
-TEST_F(RecordUploadRequestBuilderTest, AcceptEncryptedRecordsList) {
+TEST_P(RecordUploadRequestBuilderTest, AcceptEncryptedRecordsList) {
   const std::vector<std::string> kEncryptedWrappedRecords{
       "T", "E", "S", "T", "_", "I", "N", "F", "O"};
   static constexpr size_t kNumRecords = 10;
@@ -66,13 +68,17 @@ TEST_F(RecordUploadRequestBuilderTest, AcceptEncryptedRecordsList) {
         base::StrCat({"TEST_INFO_", base::NumberToString(counter)})));
   }
 
-  UploadEncryptedReportingRequestBuilder builder;
+  UploadEncryptedReportingRequestBuilder builder(need_encryption_key());
   for (const auto& record : records) {
     builder.AddRecord(record);
   }
   auto request_payload = builder.Build();
   ASSERT_TRUE(request_payload.has_value());
   ASSERT_TRUE(request_payload.value().is_dict());
+  const auto attach_ncryption_settings = request_payload.value().FindBoolKey(
+      UploadEncryptedReportingRequestBuilder::
+          GetAttachEncryptionSettingsPath());
+  EXPECT_EQ(need_encryption_key(), attach_ncryption_settings.has_value());
   base::Value* const record_list = request_payload.value().FindListKey(
       UploadEncryptedReportingRequestBuilder::GetEncryptedRecordListPath());
   ASSERT_TRUE(record_list);
@@ -87,7 +93,7 @@ TEST_F(RecordUploadRequestBuilderTest, AcceptEncryptedRecordsList) {
   }
 }
 
-TEST_F(RecordUploadRequestBuilderTest, BreakListOnSingleBadRecord) {
+TEST_P(RecordUploadRequestBuilderTest, BreakListOnSingleBadRecord) {
   const std::vector<std::string> kEncryptedWrappedRecords{
       "T", "E", "S", "T", "_", "I", "N", "F", "O"};
   static constexpr size_t kNumRecords = 10;
@@ -103,7 +109,7 @@ TEST_F(RecordUploadRequestBuilderTest, BreakListOnSingleBadRecord) {
       .mutable_sequencing_information()
       ->clear_generation_id();
 
-  UploadEncryptedReportingRequestBuilder builder;
+  UploadEncryptedReportingRequestBuilder builder(need_encryption_key());
   for (const auto& record : records) {
     builder.AddRecord(record);
   }
@@ -111,7 +117,7 @@ TEST_F(RecordUploadRequestBuilderTest, BreakListOnSingleBadRecord) {
   ASSERT_FALSE(request_payload.has_value()) << request_payload.value();
 }
 
-TEST_F(RecordUploadRequestBuilderTest, DenyPoorlyFormedEncryptedRecords) {
+TEST_P(RecordUploadRequestBuilderTest, DenyPoorlyFormedEncryptedRecords) {
   // Reject empty record.
   EncryptedRecord record;
 
@@ -145,5 +151,9 @@ TEST_F(RecordUploadRequestBuilderTest, DenyPoorlyFormedEncryptedRecords) {
   EXPECT_TRUE(EncryptedRecordDictionaryBuilder(record).Build().has_value());
 }
 
+INSTANTIATE_TEST_SUITE_P(NeedOrNoNeedKey,
+                         RecordUploadRequestBuilderTest,
+                         testing::Bool());
+
 }  // namespace
 }  // namespace reporting

chrome/browser/policy/messaging_layer/upload/upload_client.cc

@@ -42,6 +42,7 @@ void UploadClient::Create(
 }
 
 Status UploadClient::EnqueueUpload(
+    bool need_encryption_key,
     std::unique_ptr<std::vector<EncryptedRecord>> records) {
   DCHECK(records);
 
@@ -49,7 +50,8 @@ Status UploadClient::EnqueueUpload(
     return Status::StatusOK();
   }
 
-  return dm_server_upload_service_->EnqueueUpload(std::move(records));
+  return dm_server_upload_service_->EnqueueUpload(need_encryption_key,
+                                                  std::move(records));
 }
 
 UploadClient::UploadClient() = default;

chrome/browser/policy/messaging_layer/upload/upload_client.h

@@ -36,7 +36,8 @@ class UploadClient {
   UploadClient(const UploadClient& other) = delete;
   UploadClient& operator=(const UploadClient& other) = delete;
 
-  Status EnqueueUpload(std::unique_ptr<std::vector<EncryptedRecord>> record);
+  Status EnqueueUpload(bool need_encryption_key,
+                       std::unique_ptr<std::vector<EncryptedRecord>> record);
 
  private:
   UploadClient();

chrome/browser/policy/messaging_layer/upload/upload_client_unittest.cc

@@ -238,7 +238,8 @@ TEST_F(UploadClientTest, CreateUploadClientAndUploadRecords) {
   ASSERT_OK(upload_client_result) << upload_client_result.status();
 
   auto upload_client = std::move(upload_client_result.ValueOrDie());
-  auto enqueue_result = upload_client->EnqueueUpload(std::move(records));
+  auto enqueue_result = upload_client->EnqueueUpload(
+      /*need_encryption_key=*/false, std::move(records));
   EXPECT_TRUE(enqueue_result.ok());
 
   waiter.Wait();