(fake) session_manager: StorePolicy callback waits for owner key write

Makes sure that in-memory FakeSessionManagerClient does not run
StoreDevicePolicy callback until the owner key file was written (if new
owner key has to be written).
SessionManagerOperation tries to load the owner key in response to
StoreDevivePolicy, and the task to load the key might be run before the
task to write it (as they are not using the same sequence task runner).

BUG=952855

Change-Id: I41c584869741ba3d06245c378f2ef0bccd06b0a7
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1604216
Commit-Queue: Toni Baržić <tbarzic@chromium.org>
Reviewed-by: Steven Bennetts <stevenjb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#658756}

chromeos/dbus/session_manager/fake_session_manager_client.cc

@@ -498,7 +498,6 @@ void FakeSessionManagerClient::StorePolicy(
         base::BindOnce(std::move(callback), true /* success */));
   } else {
     policy_[GetMemoryStorageKey(descriptor)] = policy_blob;
-    PostReply(FROM_HERE, std::move(callback), true /* success */);
 
     if (IsChromeDevicePolicy(descriptor)) {
       // TODO(ljusten): For historical reasons, this code path only stores keys
@@ -508,19 +507,28 @@ void FakeSessionManagerClient::StorePolicy(
         GetStubPolicyFilePath(descriptor, &key_path);
         DCHECK(!key_path.empty());
 
-        base::PostTaskWithTraits(
+        base::PostTaskWithTraitsAndReply(
             FROM_HERE,
             {base::MayBlock(),
              base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN},
             base::BindOnce(StoreFiles,
                            std::map<base::FilePath, std::string>{
-                               {key_path, response.new_public_key()}}));
-        for (auto& observer : observers_)
-          observer.OwnerKeySet(true /* success */);
+                               {key_path, response.new_public_key()}}),
+            base::BindOnce(
+                &FakeSessionManagerClient::HandleOwnerKeySet,
+                weak_ptr_factory_.GetWeakPtr(),
+                base::BindOnce(std::move(callback), true /*success*/)));
       }
       for (auto& observer : observers_)
         observer.PropertyChangeComplete(true /* success */);
     }
+
+    // Run the callback if it hasn't been passed to
+    // PostTaskWithTraitsAndReply(), in which case it will be run after the
+    // owner key file was stored to disk.
+    if (callback) {
+      PostReply(FROM_HERE, std::move(callback), true /* success */);
+    }
   }
 }
 
@@ -717,4 +725,12 @@ void FakeSessionManagerClient::OnPropertyChangeComplete(bool success) {
     observer.PropertyChangeComplete(success);
 }
 
+void FakeSessionManagerClient::HandleOwnerKeySet(
+    base::OnceClosure callback_to_run) {
+  for (auto& observer : observers_)
+    observer.OwnerKeySet(true /* success */);
+
+  std::move(callback_to_run).Run();
+}
+
 }  // namespace chromeos

chromeos/dbus/session_manager/fake_session_manager_client.h

@@ -247,6 +247,11 @@ class COMPONENT_EXPORT(SESSION_MANAGER) FakeSessionManagerClient
   }
 
  private:
+  // Called in response to writing owner key file specified in new device
+  // policy - used for in-memory fake only.
+  // Notifies OwnerKeySet() observers, and runs |callback_to_run|.
+  void HandleOwnerKeySet(base::OnceClosure callback_to_run);
+
   // Whether browser restarts should be handled - intended for use in tests.
   bool supports_browser_restart_ = false;