Modify CloudPolicyClient to support multiple concurrent requests.

components/policy/core/common/cloud/cloud_policy_client.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "components/policy/core/common/cloud/cloud_policy_client.h"
 
 #include "base/bind.h"
 #include "base/guid.h"
 #include "base/logging.h"
 #include "base/stl_util.h"
@@ skipping 61 matching lines @@
 }
 
 void CloudPolicyClient::SetupRegistration(const std::string& dm_token,
                                           const std::string& client_id) {
   DCHECK(!dm_token.empty());
   DCHECK(!client_id.empty());
   DCHECK(!is_registered());
 
   dm_token_ = dm_token;
   client_id_ = client_id;
-  request_job_.reset();
+  request_jobs_.clear();
+  policy_fetch_request_job_.reset();
   STLDeleteValues(&responses_);
 
   NotifyRegistrationStateChanged();
 }
 
 void CloudPolicyClient::Register(em::DeviceRegisterRequest::Type type,
                                  em::DeviceRegisterRequest::Flavor flavor,
                                  const std::string& auth_token,
                                  const std::string& client_id,
                                  const std::string& requisition,
                                  const std::string& current_state_key) {
   DCHECK(service_);
   DCHECK(!auth_token.empty());
   DCHECK(!is_registered());
 
   if (client_id.empty()) {
     // Generate a new client ID. This is intentionally done on each new
     // registration request in order to preserve privacy. Reusing IDs would mean
     // the server could track clients by their registration attempts.
     client_id_ = base::GenerateGUID();
   } else {
     client_id_ = client_id;
   }
 
-  request_job_.reset(
+  policy_fetch_request_job_.reset(
       service_->CreateJob(DeviceManagementRequestJob::TYPE_REGISTRATION,
                           GetRequestContext()));
-  request_job_->SetOAuthToken(auth_token);
-  request_job_->SetClientID(client_id_);
+  policy_fetch_request_job_->SetOAuthToken(auth_token);
+  policy_fetch_request_job_->SetClientID(client_id_);
 
   em::DeviceRegisterRequest* request =
-      request_job_->GetRequest()->mutable_register_request();
+      policy_fetch_request_job_->GetRequest()->mutable_register_request();
   if (!client_id.empty())
     request->set_reregister(true);
   request->set_type(type);
   if (!machine_id_.empty())
     request->set_machine_id(machine_id_);
   if (!machine_model_.empty())
     request->set_machine_model(machine_model_);
   if (!requisition.empty())
     request->set_requisition(requisition);
   if (!current_state_key.empty())
     request->set_server_backed_state_key(current_state_key);
   request->set_flavor(flavor);
 
-  request_job_->SetRetryCallback(
+  policy_fetch_request_job_->SetRetryCallback(
       base::Bind(&CloudPolicyClient::OnRetryRegister, base::Unretained(this)));
 
-  request_job_->Start(base::Bind(&CloudPolicyClient::OnRegisterCompleted,
-                                 base::Unretained(this)));
+  policy_fetch_request_job_->Start(
+      base::Bind(&CloudPolicyClient::OnRegisterCompleted,
+                 base::Unretained(this)));
 }
 
 void CloudPolicyClient::SetInvalidationInfo(
     int64 version,
     const std::string& payload) {
   invalidation_version_ = version;
   invalidation_payload_ = payload;
 }
 
 void CloudPolicyClient::FetchPolicy() {
   CHECK(is_registered());
   CHECK(!types_to_fetch_.empty());
 
-  request_job_.reset(
+  policy_fetch_request_job_.reset(
       service_->CreateJob(DeviceManagementRequestJob::TYPE_POLICY_FETCH,
                           GetRequestContext()));
-  request_job_->SetDMToken(dm_token_);
-  request_job_->SetClientID(client_id_);
-  request_job_->SetUserAffiliation(user_affiliation_);
+  policy_fetch_request_job_->SetDMToken(dm_token_);
+  policy_fetch_request_job_->SetClientID(client_id_);
+  policy_fetch_request_job_->SetUserAffiliation(user_affiliation_);
 
-  em::DeviceManagementRequest* request = request_job_->GetRequest();
+  em::DeviceManagementRequest* request =
+      policy_fetch_request_job_->GetRequest();
 
   // Build policy fetch requests.
   em::DevicePolicyRequest* policy_request = request->mutable_policy_request();
   for (const auto& type_to_fetch : types_to_fetch_) {
     em::PolicyFetchRequest* fetch_request = policy_request->add_request();
     fetch_request->set_policy_type(type_to_fetch.first);
     if (!type_to_fetch.second.empty())
       fetch_request->set_settings_entity_id(type_to_fetch.second);
 
     // Request signed policy blobs to help prevent tampering on the client.
@@ skipping 30 matching lines @@
          ++key) {
       key_update_request->add_server_backed_state_key(*key);
     }
   }
 
   // Set the fetched invalidation version to the latest invalidation version
   // since it is now the invalidation version used for the latest fetch.
   fetched_invalidation_version_ = invalidation_version_;
 
   // Fire the job.
-  request_job_->Start(base::Bind(&CloudPolicyClient::OnPolicyFetchCompleted,
-                                 base::Unretained(this)));
+  policy_fetch_request_job_->Start(
+      base::Bind(&CloudPolicyClient::OnPolicyFetchCompleted,
+                 base::Unretained(this)));
 }
 
 void CloudPolicyClient::FetchRobotAuthCodes(const std::string& auth_token) {
   CHECK(is_registered());
   DCHECK(!auth_token.empty());
 
-  request_job_.reset(service_->CreateJob(
+  policy_fetch_request_job_.reset(service_->CreateJob(
       DeviceManagementRequestJob::TYPE_API_AUTH_CODE_FETCH,
       GetRequestContext()));
   // The credentials of a domain user are needed in order to mint a new OAuth2
   // authorization token for the robot account.
-  request_job_->SetOAuthToken(auth_token);
-  request_job_->SetDMToken(dm_token_);
-  request_job_->SetClientID(client_id_);
+  policy_fetch_request_job_->SetOAuthToken(auth_token);
+  policy_fetch_request_job_->SetDMToken(dm_token_);
+  policy_fetch_request_job_->SetClientID(client_id_);
 
   em::DeviceServiceApiAccessRequest* request =
-      request_job_->GetRequest()->mutable_service_api_access_request();
+      policy_fetch_request_job_->GetRequest()->
+      mutable_service_api_access_request();
   request->set_oauth2_client_id(
       GaiaUrls::GetInstance()->oauth2_chrome_client_id());
   request->add_auth_scope(GaiaConstants::kAnyApiOAuth2Scope);
 
-  request_job_->Start(
+  policy_fetch_request_job_->Start(
       base::Bind(&CloudPolicyClient::OnFetchRobotAuthCodesCompleted,
                  base::Unretained(this)));
 }
 
 void CloudPolicyClient::Unregister() {
   DCHECK(service_);
-  request_job_.reset(
+  policy_fetch_request_job_.reset(
       service_->CreateJob(DeviceManagementRequestJob::TYPE_UNREGISTRATION,
                           GetRequestContext()));
-  request_job_->SetDMToken(dm_token_);
-  request_job_->SetClientID(client_id_);
-  request_job_->GetRequest()->mutable_unregister_request();
-  request_job_->Start(base::Bind(&CloudPolicyClient::OnUnregisterCompleted,
-                                 base::Unretained(this)));
+  policy_fetch_request_job_->SetDMToken(dm_token_);
+  policy_fetch_request_job_->SetClientID(client_id_);
+  policy_fetch_request_job_->GetRequest()->mutable_unregister_request();
+  policy_fetch_request_job_->Start(
+      base::Bind(&CloudPolicyClient::OnUnregisterCompleted,
+                 base::Unretained(this)));
 }
 
 void CloudPolicyClient::UploadCertificate(
     const std::string& certificate_data,
     const CloudPolicyClient::StatusCallback& callback) {
   CHECK(is_registered());
-  request_job_.reset(
+  scoped_ptr<DeviceManagementRequestJob> request_job(
       service_->CreateJob(DeviceManagementRequestJob::TYPE_UPLOAD_CERTIFICATE,
                           GetRequestContext()));
-  request_job_->SetDMToken(dm_token_);
-  request_job_->SetClientID(client_id_);
+  request_job->SetDMToken(dm_token_);
+  request_job->SetClientID(client_id_);
 
-  em::DeviceManagementRequest* request = request_job_->GetRequest();
+  em::DeviceManagementRequest* request = request_job->GetRequest();
   request->mutable_cert_upload_request()->set_device_certificate(
       certificate_data);
 
   DeviceManagementRequestJob::Callback job_callback = base::Bind(
       &CloudPolicyClient::OnCertificateUploadCompleted,
       base::Unretained(this),
+      request_job.get(),
       callback);
-  request_job_->Start(job_callback);
+  request_jobs_.push_back(request_job.Pass());
+  request_jobs_.back()->Start(job_callback);
 }
 
 void CloudPolicyClient::UploadDeviceStatus(
     const em::DeviceStatusReportRequest* device_status,
     const em::SessionStatusReportRequest* session_status,
     const CloudPolicyClient::StatusCallback& callback) {
   CHECK(is_registered());
   // Should pass in at least one type of status.
   DCHECK(device_status || session_status);
-  request_job_.reset(
+  scoped_ptr<DeviceManagementRequestJob> request_job(
       service_->CreateJob(DeviceManagementRequestJob::TYPE_UPLOAD_STATUS,
                           GetRequestContext()));
-  request_job_->SetDMToken(dm_token_);
-  request_job_->SetClientID(client_id_);
+  request_job->SetDMToken(dm_token_);
+  request_job->SetClientID(client_id_);
 
-  em::DeviceManagementRequest* request = request_job_->GetRequest();
+  em::DeviceManagementRequest* request = request_job->GetRequest();
   if (device_status)
     *request->mutable_device_status_report_request() = *device_status;
   if (session_status)
     *request->mutable_session_status_report_request() = *session_status;
 
   DeviceManagementRequestJob::Callback job_callback = base::Bind(
       &CloudPolicyClient::OnStatusUploadCompleted,
       base::Unretained(this),
+      request_job.get(),
       callback);
 
-  // TODO(atwilson): Change CloudPolicyClient to support multiple requests in
-  // parallel, so status upload requests don't get cancelled by things like
-  // policy fetches (http://crbug.com/452563).
-  request_job_->Start(job_callback);
+  request_jobs_.push_back(request_job.Pass());
+  request_jobs_.back()->Start(job_callback);
 }
 
 void CloudPolicyClient::AddObserver(Observer* observer) {
   observers_.AddObserver(observer);
 }
 
 void CloudPolicyClient::RemoveObserver(Observer* observer) {
   observers_.RemoveObserver(observer);
 }
 
@@ skipping 20 matching lines @@
   ResponseMap::const_iterator it =
       responses_.find(std::make_pair(policy_type, settings_entity_id));
   return it == responses_.end() ? nullptr : it->second;
 }
 
 scoped_refptr<net::URLRequestContextGetter>
 CloudPolicyClient::GetRequestContext() {
   return request_context_;
 }
 
+int CloudPolicyClient::GetActiveRequestCountForTest() const {
+  return request_jobs_.size();
+}
+
 void CloudPolicyClient::OnRetryRegister(DeviceManagementRequestJob* job) {
-  DCHECK_EQ(request_job_.get(), job);
+  DCHECK_EQ(policy_fetch_request_job_.get(), job);
   // If the initial request managed to get to the server but the response didn't
   // arrive at the client then retrying with the same client ID will fail.
   // Set the re-registration flag so that the server accepts it.
   // If the server hasn't seen the client ID before then it will also accept
   // the re-registration.
   job->GetRequest()->mutable_register_request()->set_reregister(true);
 }
 
 void CloudPolicyClient::OnRegisterCompleted(
     DeviceManagementStatus status,
@@ skipping 98 matching lines @@
     int net_error,
     const em::DeviceManagementResponse& response) {
   if (status == DM_STATUS_SUCCESS && !response.has_unregister_response()) {
     // Assume unregistration has succeeded either way.
     LOG(WARNING) << "Empty unregistration response.";
   }
 
   status_ = status;
   if (status == DM_STATUS_SUCCESS) {
     dm_token_.clear();
+    // Cancel all outstanding jobs.
+    request_jobs_.clear();
     NotifyRegistrationStateChanged();
   } else {
     NotifyClientError();
   }
 }
 
 void CloudPolicyClient::OnCertificateUploadCompleted(
+    const DeviceManagementRequestJob* job,
     const CloudPolicyClient::StatusCallback& callback,
     DeviceManagementStatus status,
     int net_error,
     const enterprise_management::DeviceManagementResponse& response) {
-  if (status == DM_STATUS_SUCCESS && !response.has_cert_upload_response()) {
-    LOG(WARNING) << "Empty upload certificate response.";
-    callback.Run(false);
-    return;
-  }
-
+  bool success = true;
   status_ = status;
   if (status != DM_STATUS_SUCCESS) {
+    success = false;
     NotifyClientError();
-    callback.Run(false);
-    return;
+  } else if (!response.has_cert_upload_response()) {
+    LOG(WARNING) << "Empty upload certificate response.";
+    success = false;
   }
-  callback.Run(true);
+  callback.Run(success);
+  // Must call RemoveJob() last, because it frees |callback|.
+  RemoveJob(job);
+}
+
+void CloudPolicyClient::RemoveJob(const DeviceManagementRequestJob* job) {
+  for (auto it = request_jobs_.begin(); it != request_jobs_.end(); ++it) {
+    if (*it == job) {
+      request_jobs_.erase(it);
+      return;
+    }
+  }
+  // This job was already deleted from our list, somehow. This shouldn't
+  // happen since deleting the job should cancel the callback.
+  NOTREACHED();
 }
 
 void CloudPolicyClient::OnStatusUploadCompleted(
+    const DeviceManagementRequestJob* job,
     const CloudPolicyClient::StatusCallback& callback,
     DeviceManagementStatus status,
     int net_error,
     const enterprise_management::DeviceManagementResponse& response) {
   status_ = status;
   if (status != DM_STATUS_SUCCESS)
     NotifyClientError();
 
   callback.Run(status == DM_STATUS_SUCCESS);
+  // Must call RemoveJob() last, because it frees |callback|.
+  RemoveJob(job);
 }
 
 void CloudPolicyClient::NotifyPolicyFetched() {
   FOR_EACH_OBSERVER(Observer, observers_, OnPolicyFetched(this));
 }
 
 void CloudPolicyClient::NotifyRegistrationStateChanged() {
   FOR_EACH_OBSERVER(Observer, observers_, OnRegistrationStateChanged(this));
 }
 
 void CloudPolicyClient::NotifyRobotAuthCodesFetched() {
   FOR_EACH_OBSERVER(Observer, observers_, OnRobotAuthCodesFetched(this));
 }
 
 void CloudPolicyClient::NotifyClientError() {
   FOR_EACH_OBSERVER(Observer, observers_, OnClientError(this));
 }
 
 }  // namespace policy