arc: data_snapshotd: LoadSnapshot impl

Before starting ARC container or VM try to load a snapshot if available.
ARC data snapshots are loaded only for MGS. If there is no snapshot or
the process fails, start ARC session as usual.

BUG=b:161221001
TEST=components_unittests
TEST=unit_tests

Change-Id: I7e222f9abd11b0471a40a475c73931886582bf55
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2498488
Auto-Submit: Polina Bondarenko <pbond@chromium.org>
Commit-Queue: Polina Bondarenko <pbond@chromium.org>
Reviewed-by: Oleksandr Peletskyi <peletskyi@chromium.org>
Reviewed-by: Yusuke Sato <yusukes@chromium.org>
Cr-Commit-Position: refs/heads/master@{#829117}

components/arc/BUILD.gn

@@ -261,6 +261,7 @@ static_library("arc_base") {
     "//chromeos/memory:memory",
     "//chromeos/system:system",
     "//components/account_id",
+    "//components/arc/enterprise",
     "//components/keyed_service/content",
     "//components/prefs",
     "//components/user_manager",

components/arc/enterprise/arc_data_snapshotd_bridge.cc

@@ -128,5 +128,19 @@ void ArcDataSnapshotdBridge::TakeSnapshot(
   std::move(callback).Run(true /* success */);
 }
 
+void ArcDataSnapshotdBridge::LoadSnapshot(
+    const std::string& account_id,
+    base::OnceCallback<void(bool, bool)> callback) {
+  if (!is_available_) {
+    LOG(ERROR) << "LoadSnapshot call when D-Bus service is not available.";
+    std::move(callback).Run(false /* success */, false /* last */);
+    return;
+  }
+  VLOG(1) << "LoadSnapshot via D-Bus";
+  NOTIMPLEMENTED();
+
+  std::move(callback).Run(true /* success */, true /* last */);
+}
+
 }  // namespace data_snapshotd
 }  // namespace arc

components/arc/enterprise/arc_data_snapshotd_bridge.h

@@ -31,6 +31,8 @@ class ArcDataSnapshotdBridge {
   void ClearSnapshot(bool last, base::OnceCallback<void(bool)> callback);
   void TakeSnapshot(const std::string& account_id,
                     base::OnceCallback<void(bool)> callback);
+  void LoadSnapshot(const std::string& account_id,
+                    base::OnceCallback<void(bool, bool)> callback);
 
   bool is_available_for_testing() { return is_available_; }
 

components/arc/enterprise/arc_data_snapshotd_manager.cc

@@ -320,6 +320,23 @@ void ArcDataSnapshotdManager::EnsureDaemonStopped(base::OnceClosure callback) {
   StopDaemon(std::move(callback));
 }
 
+void ArcDataSnapshotdManager::StartLoadingSnapshot(base::OnceClosure callback) {
+  // Do not load a snapshot if it's not a normal MGS setup.
+  if (state_ != State::kNone) {
+    std::move(callback).Run();
+    return;
+  }
+  std::string account_id = GetCryptohomeAccountId();
+  if (!account_id.empty() && IsSnapshotEnabled() &&
+      (snapshot_.last() || snapshot_.previous())) {
+    EnsureDaemonStarted(base::BindOnce(
+        &ArcDataSnapshotdManager::LoadSnapshot, weak_ptr_factory_.GetWeakPtr(),
+        std::move(account_id), std::move(callback)));
+    return;
+  }
+  std::move(callback).Run();
+}
+
 bool ArcDataSnapshotdManager::IsAutoLoginConfigured() {
   switch (state_) {
     case ArcDataSnapshotdManager::State::kBlockedUi:
@@ -328,6 +345,7 @@ bool ArcDataSnapshotdManager::IsAutoLoginConfigured() {
       return true;
     case ArcDataSnapshotdManager::State::kNone:
     case ArcDataSnapshotdManager::State::kRestored:
+    case ArcDataSnapshotdManager::State::kRunning:
       return false;
   }
 }
@@ -338,6 +356,7 @@ bool ArcDataSnapshotdManager::IsAutoLoginAllowed() {
       return false;
     case ArcDataSnapshotdManager::State::kNone:
     case ArcDataSnapshotdManager::State::kRestored:
+    case ArcDataSnapshotdManager::State::kRunning:
     case ArcDataSnapshotdManager::State::kMgsLaunched:
     case ArcDataSnapshotdManager::State::kMgsToLaunch:
       return true;
@@ -365,9 +384,17 @@ void ArcDataSnapshotdManager::OnSessionStateChanged() {
         return;
       }
       break;
-    case ArcDataSnapshotdManager::State::kBlockedUi:
-    case ArcDataSnapshotdManager::State::kNone:
-    case ArcDataSnapshotdManager::State::kRestored:
+    case State::kRunning:
+      if (user_manager::UserManager::Get() &&
+          user_manager::UserManager::Get()->IsLoggedInAsPublicAccount()) {
+        return;
+      }
+      // It is a correct state. Exit MGS withg loaded snapshot.
+      state_ = State::kNone;
+      return;
+    case State::kBlockedUi:
+    case State::kNone:
+    case State::kRestored:
       break;
   }
 }
@@ -438,6 +465,19 @@ void ArcDataSnapshotdManager::TakeSnapshot(const std::string& account_id) {
                                  weak_ptr_factory_.GetWeakPtr()));
 }
 
+void ArcDataSnapshotdManager::LoadSnapshot(const std::string& account_id,
+                                           base::OnceClosure callback) {
+  if (!bridge_) {
+    OnSnapshotLoaded(std::move(callback), false /* success */,
+                     false /* last */);
+    return;
+  }
+  bridge_->LoadSnapshot(
+      account_id,
+      base::BindOnce(&ArcDataSnapshotdManager::OnSnapshotLoaded,
+                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
+}
+
 void ArcDataSnapshotdManager::OnSnapshotsCleared(bool success) {
   switch (state_) {
     case State::kBlockedUi:
@@ -451,6 +491,7 @@ void ArcDataSnapshotdManager::OnSnapshotsCleared(bool success) {
       return;
     case State::kMgsToLaunch:
     case State::kMgsLaunched:
+    case State::kRunning:
       LOG(WARNING) << "Snapshots are cleared while in incorrect state";
       return;
   }
@@ -532,17 +573,8 @@ void ArcDataSnapshotdManager::OnArcInstanceStopped(bool success) {
     OnSnapshotTaken(false /* success */);
     return;
   }
-  std::string account_id = "";
-  // Take snapshots only for MGSs.
-  if (user_manager::UserManager::Get() &&
-      user_manager::UserManager::Get()->IsLoggedInAsPublicAccount() &&
-      state_ == State::kMgsLaunched &&
-      user_manager::UserManager::Get()->GetActiveUser()) {
-    cryptohome::Identification id(
-        user_manager::UserManager::Get()->GetActiveUser()->GetAccountId());
-    account_id = id.id();
-  }
-  if (account_id.empty()) {
+  std::string account_id = GetCryptohomeAccountId();
+  if (account_id.empty() || state_ != State::kMgsLaunched) {
     LOG(ERROR) << "Cryptohome account ID is empty.";
     OnSnapshotTaken(false /* success */);
     return;
@@ -565,5 +597,36 @@ void ArcDataSnapshotdManager::OnSnapshotTaken(bool success) {
   EnsureDaemonStopped(std::move(attempt_user_exit_callback_));
 }
 
+void ArcDataSnapshotdManager::OnSnapshotLoaded(base::OnceClosure callback,
+                                               bool success,
+                                               bool last) {
+  if (!success) {
+    LOG(ERROR) << "Failed to load ARC data directory snapshot.";
+    std::move(callback).Run();
+    return;
+  }
+  VLOG(1) << "Successfully loaded " << (last ? "last" : "previous")
+          << " snapshot";
+  state_ = State::kRunning;
+  // Clear last snapshot if the previous one was loaded.
+  if (!last && snapshot_.last())
+    snapshot_.ClearSnapshot(true /* last */);
+  EnsureDaemonStopped(base::DoNothing());
+  std::move(callback).Run();
+}
+
+std::string ArcDataSnapshotdManager::GetCryptohomeAccountId() {
+  // Take snapshots only for MGSs.
+  if (user_manager::UserManager::Get() &&
+      user_manager::UserManager::Get()->IsLoggedInAsPublicAccount() &&
+      user_manager::UserManager::Get()->GetActiveUser()) {
+    return cryptohome::Identification(user_manager::UserManager::Get()
+                                          ->GetActiveUser()
+                                          ->GetAccountId())
+        .id();
+  }
+  return "";
+}
+
 }  // namespace data_snapshotd
 }  // namespace arc

components/arc/enterprise/arc_data_snapshotd_manager.h

@@ -36,6 +36,8 @@ class ArcDataSnapshotdManager final
     kNone,
     // Blocked UI mode is ON.
     kBlockedUi,
+    // Running with a snapshot.
+    kRunning,
     // In blocked UI mode, MGS can be launched.
     kMgsToLaunch,
     // MGS is launched to create a snapshot.
@@ -187,6 +189,10 @@ class ArcDataSnapshotdManager final
   // Stops arc-data-snapshotd.
   void EnsureDaemonStopped(base::OnceClosure callback);
 
+  // Starts loading a snapshot to android-data directory.
+  // |callback| is called once the process is over.
+  void StartLoadingSnapshot(base::OnceClosure callback);
+
   // Returns true if autologin to public account should be performed.
   bool IsAutoLoginConfigured();
   // Returns true if autologin is allowed to be performed and manager is not
@@ -231,6 +237,7 @@ class ArcDataSnapshotdManager final
   void GenerateKeyPair();
   void ClearSnapshot(bool last, base::OnceCallback<void(bool)> callback);
   void TakeSnapshot(const std::string& account_id);
+  void LoadSnapshot(const std::string& account_id, base::OnceClosure callback);
 
   // Called once the outdated snapshots were removed or ensured that there are
   // no outdated snapshots.
@@ -254,6 +261,13 @@ class ArcDataSnapshotdManager final
   // Called once a snapshot is taken.
   void OnSnapshotTaken(bool success);
 
+  // Called once a snapshot is taken.
+  void OnSnapshotLoaded(base::OnceClosure callback, bool success, bool last);
+
+  // Returns non-empty account ID string if a MGS is active.
+  // Otherwise returns an empty string.
+  std::string GetCryptohomeAccountId();
+
   static bool is_snapshot_enabled_for_testing_;
   State state_ = State::kNone;
   Snapshot snapshot_;

components/arc/session/arc_session_impl.cc

@@ -32,6 +32,7 @@
 #include "chromeos/system/scheduler_configuration_manager_base.h"
 #include "components/arc/arc_features.h"
 #include "components/arc/arc_util.h"
+#include "components/arc/enterprise/arc_data_snapshotd_manager.h"
 #include "components/arc/session/arc_bridge_host_impl.h"
 #include "components/user_manager/user_manager.h"
 #include "components/version_info/channel.h"
@@ -531,7 +532,23 @@ void ArcSessionImpl::OnSocketCreated(base::ScopedFD socket_fd) {
     return;
   }
 
-  VLOG(2) << "Socket is created. Starting ARC container";
+  VLOG(2) << "Socket is created. Start loading ARC data snapshot";
+  StartLoadingDataSnapshot(base::BindOnce(&ArcSessionImpl::OnDataSnapshotLoaded,
+                                          weak_factory_.GetWeakPtr(),
+                                          std::move(socket_fd)));
+}
+
+void ArcSessionImpl::StartLoadingDataSnapshot(base::OnceClosure callback) {
+  auto* arc_data_snapshotd_manager =
+      arc::data_snapshotd::ArcDataSnapshotdManager::Get();
+  if (arc_data_snapshotd_manager)
+    arc_data_snapshotd_manager->StartLoadingSnapshot(std::move(callback));
+  else
+    std::move(callback).Run();
+}
+
+void ArcSessionImpl::OnDataSnapshotLoaded(base::ScopedFD socket_fd) {
+  VLOG(2) << "Starting ARC container";
   client_->UpgradeArc(
       std::move(upgrade_params_),
       base::BindOnce(&ArcSessionImpl::OnUpgraded, weak_factory_.GetWeakPtr(),

components/arc/session/arc_session_impl.h

@@ -217,6 +217,15 @@ class ArcSessionImpl
   // Called when arcbridge socket is created.
   void OnSocketCreated(base::ScopedFD fd);
 
+  // Loads ARC data/ snapshot if necessary.
+  // |callback| is called once the load process is finished.
+  void StartLoadingDataSnapshot(base::OnceClosure callback);
+
+  // Called when ARC data/ snapshot step is done: either snapshot is loaded or
+  // skipped.
+  // |socket_fd| should be a socket to be passed to OnUpgraded.
+  void OnDataSnapshotLoaded(base::ScopedFD scoped_fd);
+
   // D-Bus callback for UpgradeArcContainer(). |socket_fd| should be a socket
   // which should be accept(2)ed to connect ArcBridgeService Mojo channel.
   void OnUpgraded(base::ScopedFD socket_fd, bool result);