PSM: Prevent parallel execution of PSM and Hash dance

In case the private set membership is in progress, and RetryStep has
been called, the PrivateSetMembershipRetryStep should return true to
prevent Hash dance protocol from executing while private set membership
is in progress. This CL is preventing that parallel execution from
occurring, along with its required regression tests.

BUG=chromium:1156119

Change-Id: I06d25b1fd1171d4b680b5d6b0be21bdb2282a518
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2577216
Commit-Queue: Amr Aboelkher <amraboelkher@google.com>
Reviewed-by: Pavol Marko <pmarko@chromium.org>
Reviewed-by: Maksim Ivanov <emaxx@chromium.org>
Cr-Commit-Position: refs/heads/master@{#835004}

chrome/browser/chromeos/policy/auto_enrollment_client_impl.cc

@@ -1027,12 +1027,17 @@ bool AutoEnrollmentClientImpl::RetryStep() {
 }
 
 bool AutoEnrollmentClientImpl::PrivateSetMembershipRetryStep() {
-  // Don't retry if the protocol is disabled, protocol is still running, or an
-  // error occurred while executing the protocol.
+  // Don't retry if the protocol is disabled, or an error occurred while
+  // executing the protocol.
   if (!private_set_membership_helper_ ||
-      private_set_membership_helper_->HasPrivateSetMembershipError() ||
-      private_set_membership_helper_->IsCheckMembershipInProgress())
+      private_set_membership_helper_->HasPrivateSetMembershipError()) {
     return false;
+  }
+
+  // If the private set membership protocol is in progress, signal to the caller
+  // that nothing else needs to be done.
+  if (private_set_membership_helper_->IsCheckMembershipInProgress())
+    return true;
 
   const base::Optional<bool> private_set_membership_server_state =
       private_set_membership_helper_->GetPrivateSetMembershipCachedDecision();

chrome/browser/chromeos/policy/auto_enrollment_client_impl.h

@@ -167,8 +167,8 @@ class AutoEnrollmentClientImpl
 
   // Retries running private set membership protocol, if the protocol
   // is enabled and it is possible to start. Returns true if the protocol is
-  // enabled and progress has been made, false if the protocol is done. Also,
-  // that protocol is being started only one time.
+  // enabled or it's in progress, false if the protocol is done. Note that the
+  // PSM protocol is only performed once per OOBE flow.
   bool PrivateSetMembershipRetryStep();
 
   // Cleans up and invokes |progress_callback_|.

chrome/browser/chromeos/policy/auto_enrollment_client_impl_unittest.cc

@@ -366,7 +366,8 @@ class AutoEnrollmentClientImplTest
 
   void ServerWillReplyAsync(DeviceManagementService::JobControl** job) {
     EXPECT_CALL(*service_, StartJob(_))
-        .WillOnce(service_->StartJobFullControl(job));
+        .WillOnce(DoAll(service_->CaptureJobType(&last_async_job_type_),
+                        service_->StartJobFullControl(job)));
   }
 
   bool HasCachedDecision() {
@@ -506,6 +507,8 @@ class AutoEnrollmentClientImplTest
   AutoEnrollmentState state_;
   DeviceManagementService::JobConfiguration::JobType failed_job_type_ =
       DeviceManagementService::JobConfiguration::TYPE_INVALID;
+  DeviceManagementService::JobConfiguration::JobType last_async_job_type_ =
+      DeviceManagementService::JobConfiguration::TYPE_INVALID;
   DeviceManagementService::JobConfiguration::JobType auto_enrollment_job_type_ =
       DeviceManagementService::JobConfiguration::TYPE_INVALID;
   DeviceManagementService::JobConfiguration::JobType state_retrieval_job_type_ =
@@ -1411,6 +1414,52 @@ class PrivateSetMembershipHelperTest : public AutoEnrollmentClientImplTest {
         .RetiresOnSaturation();
   }
 
+  // Holds the full control of the given job in |job| and captures the job type
+  // in |private_set_membership_last_job_type_|, and its request in
+  // |private_set_membership_last_request_|.
+  void ServerWillReplyAsyncForPrivateSetMembership(
+      DeviceManagementService::JobControl** job) {
+    EXPECT_CALL(*service_, StartJob(_))
+        .WillOnce(DoAll(
+            service_->CaptureJobType(&private_set_membership_last_job_type_),
+            service_->CaptureRequest(&private_set_membership_last_request_),
+            service_->StartJobFullControl(job)));
+  }
+
+  void ServerReplyForPrivateSetMembershipAsyncJobWithOprfResponse(
+      DeviceManagementService::JobControl** job) {
+    em::DeviceManagementResponse response =
+        GetPrivateSetMembershipOprfResponse();
+
+    ServerReplyForAsyncJob(job, net::OK, DeviceManagementService::kSuccess,
+                           response);
+  }
+
+  void ServerReplyForPrivateSetMembershipAsyncJobWithQueryResponse(
+      DeviceManagementService::JobControl** job) {
+    em::DeviceManagementResponse response =
+        GetPrivateSetMembershipQueryResponse();
+
+    ServerReplyForAsyncJob(job, net::OK, DeviceManagementService::kSuccess,
+                           response);
+  }
+
+  void ServerFailsForAsyncJob(DeviceManagementService::JobControl** job) {
+    em::DeviceManagementResponse dummy_response;
+    ServerReplyForAsyncJob(job, net::OK,
+                           DeviceManagementService::kServiceUnavailable,
+                           dummy_response);
+  }
+
+  // Mocks the server reply for the full controlled job |job|.
+  void ServerReplyForAsyncJob(
+      DeviceManagementService::JobControl** job,
+      int net_error,
+      int response_code,
+      const enterprise_management::DeviceManagementResponse& response) {
+    service_->DoURLCompletion(job, net_error, response_code, response);
+  }
+
   const em::PrivateSetMembershipRequest& private_set_membership_request()
       const {
     return private_set_membership_last_request_
@@ -2095,6 +2144,156 @@ TEST_P(PrivateSetMembershipHelperAndHashDanceTest,
       PrivateSetMembershipHashDanceComparison::kBothError);
 }
 
+TEST_P(PrivateSetMembershipHelperAndHashDanceTest,
+       RetryWhileWaitingForPrivateSetMembershipOprfResponseAndHashDanceFails) {
+  InSequence sequence;
+
+  DeviceManagementService::JobControl* psm_rlwe_oprf_job = nullptr;
+  DeviceManagementService::JobControl* hash_dance_job = nullptr;
+
+  // Expect two requests and capture them, in order, when available in
+  // |psm_rlwe_oprf_job| and |hash_dance_job|.
+  ServerWillReplyAsyncForPrivateSetMembership(&psm_rlwe_oprf_job);
+  ServerWillReplyAsync(&hash_dance_job);
+
+  // Expect none of the jobs have been captured.
+  EXPECT_FALSE(psm_rlwe_oprf_job);
+  EXPECT_FALSE(hash_dance_job);
+
+  client()->Start();
+  base::RunLoop().RunUntilIdle();
+
+  // Verify the only job that has been captured is the PSM RLWE OPRF request.
+  VerifyPrivateSetMembershipRlweOprfRequest();
+  VerifyPrivateSetMembershipLastRequestJobType();
+  ASSERT_TRUE(psm_rlwe_oprf_job);
+  EXPECT_FALSE(hash_dance_job);
+
+  // Trigger RetryStep.
+  client()->Retry();
+
+  // Verify hash dance job has not been triggered after RetryStep.
+  EXPECT_FALSE(hash_dance_job);
+
+  // Fail for private set membership RLWE OPRF request.
+  ServerFailsForAsyncJob(&psm_rlwe_oprf_job);
+
+  // Verify failure of private set membership protocol.
+  EXPECT_EQ(GetStateDiscoveryResult(), StateDiscoveryResult::kFailure);
+  ExpectPrivateSetMembershipHistograms({PrivateSetMembershipStatus::kAttempt,
+                                        PrivateSetMembershipStatus::kError},
+                                       /*success_time_recorded=*/false);
+
+  // Verify hash dance job has been captured.
+  ASSERT_TRUE(hash_dance_job);
+  EXPECT_EQ(DeviceManagementService::JobConfiguration::TYPE_AUTO_ENROLLMENT,
+            last_async_job_type_);
+
+  // Fail for DeviceAutoEnrollmentRequest i.e. hash dance request.
+  ServerFailsForAsyncJob(&hash_dance_job);
+
+  // Verify failure of Hash dance by inexistence of its cached decision.
+  EXPECT_FALSE(HasCachedDecision());
+
+  // Verify that no enrollment has been done, and no state has been retrieved.
+  EXPECT_EQ(state_, AUTO_ENROLLMENT_STATE_SERVER_ERROR);
+  EXPECT_FALSE(HasServerBackedState());
+
+  // Verify recorded comparison between private set membership and Hash
+  // dance.
+  ExpectPrivateSetMembershipHashDanceComparisonRecorded(
+      PrivateSetMembershipHashDanceComparison::kBothError);
+
+  // Verify both jobs have finished.
+  EXPECT_EQ(hash_dance_job, nullptr);
+  EXPECT_EQ(psm_rlwe_oprf_job, nullptr);
+}
+
+TEST_P(PrivateSetMembershipHelperAndHashDanceTest,
+       RetryWhileWaitingForPrivateSetMembershipQueryResponseAndHashDanceFails) {
+  InSequence sequence;
+
+  DeviceManagementService::JobControl* psm_rlwe_oprf_job = nullptr;
+  DeviceManagementService::JobControl* psm_rlwe_query_job = nullptr;
+  DeviceManagementService::JobControl* hash_dance_job = nullptr;
+
+  // Expect three requests and capture them, in order, when available in
+  // |psm_rlwe_oprf_job|, |psm_rlwe_query_job|, and |hash_dance_job|.
+  ServerWillReplyAsyncForPrivateSetMembership(&psm_rlwe_oprf_job);
+  ServerWillReplyAsyncForPrivateSetMembership(&psm_rlwe_query_job);
+  ServerWillReplyAsync(&hash_dance_job);
+
+  // Expect none of the jobs have been captured.
+  EXPECT_FALSE(psm_rlwe_oprf_job);
+  EXPECT_FALSE(psm_rlwe_query_job);
+  EXPECT_FALSE(hash_dance_job);
+
+  client()->Start();
+  base::RunLoop().RunUntilIdle();
+
+  // Verify the only job that has been captured is the PSM RLWE OPRF request.
+  VerifyPrivateSetMembershipRlweOprfRequest();
+  VerifyPrivateSetMembershipLastRequestJobType();
+  ASSERT_TRUE(psm_rlwe_oprf_job);
+  EXPECT_FALSE(psm_rlwe_query_job);
+  EXPECT_FALSE(hash_dance_job);
+
+  // Reply with PSM RLWE OPRF response.
+  ServerReplyForPrivateSetMembershipAsyncJobWithOprfResponse(
+      &psm_rlwe_oprf_job);
+
+  // Verify the only job that has been captured is the PSM RLWE Query request.
+  VerifyPrivateSetMembershipRlweQueryRequest();
+  VerifyPrivateSetMembershipLastRequestJobType();
+  ASSERT_TRUE(psm_rlwe_query_job);
+  EXPECT_FALSE(hash_dance_job);
+
+  // Trigger RetryStep.
+  client()->Retry();
+
+  // Verify hash dance job has not been triggered after RetryStep.
+  EXPECT_FALSE(hash_dance_job);
+
+  // Reply with PSM RLWE Query response.
+  ServerReplyForPrivateSetMembershipAsyncJobWithQueryResponse(
+      &psm_rlwe_query_job);
+
+  // Verify private set membership result.
+  EXPECT_EQ(GetStateDiscoveryResult(),
+            GetExpectedMembershipResult()
+                ? StateDiscoveryResult::kSuccessHasServerSideState
+                : StateDiscoveryResult::kSuccessNoServerSideState);
+  ExpectPrivateSetMembershipHistograms(
+      {PrivateSetMembershipStatus::kAttempt,
+       PrivateSetMembershipStatus::kSuccessfulDetermination},
+      /*success_time_recorded=*/true);
+
+  // Verify hash dance job has been captured.
+  ASSERT_TRUE(hash_dance_job);
+  EXPECT_EQ(DeviceManagementService::JobConfiguration::TYPE_AUTO_ENROLLMENT,
+            last_async_job_type_);
+
+  // Fail for DeviceAutoEnrollmentRequest i.e. hash dance request.
+  ServerFailsForAsyncJob(&hash_dance_job);
+
+  // Verify failure of Hash dance by inexistence of its cached decision.
+  EXPECT_FALSE(HasCachedDecision());
+
+  // Verify that no enrollment has been done, and no state has been retrieved.
+  EXPECT_EQ(state_, AUTO_ENROLLMENT_STATE_SERVER_ERROR);
+  EXPECT_FALSE(HasServerBackedState());
+
+  // Verify recorded comparison between private set membership and Hash
+  // dance.
+  ExpectPrivateSetMembershipHashDanceComparisonRecorded(
+      PrivateSetMembershipHashDanceComparison::kPSMSuccessHashDanceError);
+
+  // Verify all jobs have finished.
+  EXPECT_EQ(hash_dance_job, nullptr);
+  EXPECT_EQ(psm_rlwe_oprf_job, nullptr);
+  EXPECT_EQ(psm_rlwe_query_job, nullptr);
+}
+
 INSTANTIATE_TEST_SUITE_P(
     PrivateSetMembershipAndHashDance,
     PrivateSetMembershipHelperAndHashDanceTest,