Testing: Unify access to protobufs in FakeUserDataAuthClient

Bug: b:239430274 Change-Id: I1ffa74c59fcbb72ee22d6fa743166aaed3cebb16 Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4227737 Reviewed-by: Elie Maamari <emaamari@google.com> Reviewed-by: Leonard Grey <lgrey@chromium.org> Reviewed-by: Martin Bidlingmaier <mbid@google.com> Commit-Queue: Denis Kuznetsov <antrim@chromium.org> Cr-Commit-Position: refs/heads/main@{#1103280}
chromium · Feb 9, 2023 · 16cf368 · 16cf368
1 parent 0624f1a
commit 16cf368
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 63 deletions.
diff --git a/chrome/browser/ash/login/encryption_migration_browsertest.cc b/chrome/browser/ash/login/encryption_migration_browsertest.cc
@@ -24,6 +24,7 @@
 #include "chrome/browser/browser_process.h"
 #include "chrome/browser/ui/webui/ash/login/gaia_screen_handler.h"
 #include "chrome/browser/ui/webui/ash/login/oobe_ui.h"
+#include "chromeos/ash/components/dbus/cryptohome/UserDataAuth.pb.h"
 #include "chromeos/ash/components/dbus/cryptohome/account_identifier_operators.h"
 #include "chromeos/ash/components/dbus/cryptohome/rpc.pb.h"
 #include "chromeos/ash/components/dbus/userdataauth/fake_userdataauth_client.h"
@@ -58,6 +59,8 @@ const test::UIPath kInsufficientSpaceSkipButton = {
 const test::UIPath kInsufficientSpaceRestartButton = {
     kEncryptionMigrationId, "insufficient-space-restart-button"};
 
+using AuthOp = FakeUserDataAuthClient::Operation;
+
 }  // namespace
 
 // Base class for testing encryption migration during sign-in.
@@ -144,10 +147,11 @@ class EncryptionMigrationTestBase
     test::OobeJS().ExpectHiddenPath(kErrorDialog);
     test::OobeJS().ExpectHiddenPath(kInsufficientSpaceDialog);
 
-    EXPECT_EQ(
-        GetTestCryptohomeId(),
-        FakeUserDataAuthClient::Get()->get_id_for_disk_migrated_to_dircrypto());
-    EXPECT_FALSE(FakeUserDataAuthClient::Get()->minimal_migration());
+    auto migrate_request =
+        FakeUserDataAuthClient::Get()
+            ->GetLastRequest<AuthOp::kStartMigrateToDircrypto>();
+    EXPECT_EQ(GetTestCryptohomeId(), migrate_request.account_id());
+    EXPECT_FALSE(migrate_request.minimal_migration());
 
     EXPECT_EQ(
         0,
@@ -269,8 +273,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, SkipWithNoPolicySet) {
   WaitForActiveSession();
 
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 }
 
 IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, MigrateWithNoUserPolicySet) {
@@ -289,8 +292,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, MigrateWithNoUserPolicySet) {
   test::OobeJS().ExpectVisiblePath(kUpgradeButton);
 
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 
   test::OobeJS().TapOnPath(kUpgradeButton);
 
@@ -363,8 +365,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest,
 
   WaitForActiveSession();
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 }
 
 IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, MigrateWithInsuficientSpace) {
@@ -390,8 +391,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, MigrateWithInsuficientSpace) {
   EXPECT_EQ(
       1, chromeos::FakePowerManagerClient::Get()->num_request_restart_calls());
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 }
 
 IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, InsufficientSpaceOnResume) {
@@ -417,8 +417,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, InsufficientSpaceOnResume) {
   EXPECT_EQ(
       1, chromeos::FakePowerManagerClient::Get()->num_request_restart_calls());
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 }
 
 IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, MigrationFailure) {
@@ -433,9 +432,10 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, MigrationFailure) {
       .CreateWaiter(test::GetOobeElementPath(kMigratingDialog))
       ->Wait();
 
-  EXPECT_EQ(
-      GetTestCryptohomeId(),
-      FakeUserDataAuthClient::Get()->get_id_for_disk_migrated_to_dircrypto());
+  EXPECT_EQ(GetTestCryptohomeId(),
+            FakeUserDataAuthClient::Get()
+                ->GetLastRequest<AuthOp::kStartMigrateToDircrypto>()
+                .account_id());
   FakeUserDataAuthClient::Get()->NotifyDircryptoMigrationProgress(
       ::user_data_auth::DircryptoMigrationStatus::DIRCRYPTO_MIGRATION_FAILED,
       5 /*current*/, 5 /*total*/);
@@ -481,8 +481,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest, LowBattery) {
 
   WaitForActiveSession();
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 }
 
 IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest,
@@ -505,8 +504,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest,
   test::OobeJS().ExpectPathDisplayed(false, kUpgradeButton);
 
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 }
 
 IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest,
@@ -526,8 +524,7 @@ IN_PROC_BROWSER_TEST_F(EncryptionMigrationTest,
   test::OobeJS().ExpectHiddenPath(kErrorDialog);
 
   EXPECT_FALSE(FakeUserDataAuthClient::Get()
-                   ->get_id_for_disk_migrated_to_dircrypto()
-                   .has_account_id());
+                   ->WasCalled<AuthOp::kStartMigrateToDircrypto>());
 
   SetBatteryPercent(60);
 

diff --git a/chrome/browser/ash/login/oobe_browsertest.cc b/chrome/browser/ash/login/oobe_browsertest.cc
@@ -110,7 +110,8 @@ IN_PROC_BROWSER_TEST_F(OobeTest, NewUser) {
 
   // Verify the parameters that were passed to the latest AddAuthFactor call.
   const user_data_auth::AddAuthFactorRequest& request =
-      FakeUserDataAuthClient::Get()->get_last_add_authfactor_request();
+      FakeUserDataAuthClient::Get()
+          ->GetLastRequest<FakeUserDataAuthClient::Operation::kAddAuthFactor>();
   EXPECT_EQ(request.auth_factor().label(), kCryptohomeGaiaKeyLabel);
   EXPECT_FALSE(request.auth_input().password_input().secret().empty());
   EXPECT_EQ(user_data_auth::AUTH_FACTOR_TYPE_PASSWORD,

diff --git a/chromeos/ash/components/dbus/userdataauth/DEPS b/chromeos/ash/components/dbus/userdataauth/DEPS
@@ -0,0 +1,3 @@
+include_rules = [
+  "+third_party/protobuf/src/google/protobuf"
+]
diff --git a/chromeos/ash/components/dbus/userdataauth/fake_userdataauth_client.cc b/chromeos/ash/components/dbus/userdataauth/fake_userdataauth_client.cc
@@ -716,9 +716,9 @@ void FakeUserDataAuthClient::EndFingerprintAuthSession(
 void FakeUserDataAuthClient::StartMigrateToDircrypto(
     const ::user_data_auth::StartMigrateToDircryptoRequest& request,
     StartMigrateToDircryptoCallback callback) {
-  last_migrate_to_dircrypto_request_ = request;
   ::user_data_auth::StartMigrateToDircryptoReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kStartMigrateToDircrypto>(request);
 
   dircrypto_migration_progress_ = 0;
 
@@ -773,6 +773,7 @@ void FakeUserDataAuthClient::StartAuthSession(
     StartAuthSessionCallback callback) {
   ::user_data_auth::StartAuthSessionReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kStartAuthSession>(request);
 
   if (auto error = TakeOperationError(Operation::kStartAuthSession);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -847,6 +848,7 @@ void FakeUserDataAuthClient::ListAuthFactors(
     ListAuthFactorsCallback callback) {
   ::user_data_auth::ListAuthFactorsReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kListAuthFactors>(request);
 
   if (auto error = TakeOperationError(Operation::kListAuthFactors);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -910,6 +912,7 @@ void FakeUserDataAuthClient::PrepareGuestVault(
     PrepareGuestVaultCallback callback) {
   ::user_data_auth::PrepareGuestVaultReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kPrepareGuestVault>(request);
 
   if (auto error = TakeOperationError(Operation::kPrepareGuestVault);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -929,6 +932,7 @@ void FakeUserDataAuthClient::PrepareEphemeralVault(
     PrepareEphemeralVaultCallback callback) {
   ::user_data_auth::PrepareEphemeralVaultReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kPrepareEphemeralVault>(request);
 
   if (auto error = TakeOperationError(Operation::kPrepareEphemeralVault);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -977,6 +981,7 @@ void FakeUserDataAuthClient::CreatePersistentUser(
     CreatePersistentUserCallback callback) {
   ::user_data_auth::CreatePersistentUserReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kCreatePersistentUser>(request);
 
   if (auto error = TakeOperationError(Operation::kCreatePersistentUser);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -1017,6 +1022,7 @@ void FakeUserDataAuthClient::PreparePersistentVault(
     PreparePersistentVaultCallback callback) {
   ::user_data_auth::PreparePersistentVaultReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kPreparePersistentVault>(request);
 
   if (auto error = TakeOperationError(Operation::kPreparePersistentVault);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -1068,6 +1074,7 @@ void FakeUserDataAuthClient::PrepareVaultForMigration(
     PrepareVaultForMigrationCallback callback) {
   ::user_data_auth::PrepareVaultForMigrationReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kPrepareVaultForMigration>(request);
 
   if (auto error = TakeOperationError(Operation::kPrepareVaultForMigration);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -1122,7 +1129,7 @@ void FakeUserDataAuthClient::AddAuthFactor(
     AddAuthFactorCallback callback) {
   ::user_data_auth::AddAuthFactorReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
-  last_add_auth_factor_request_ = request;
+  RememberRequest<Operation::kAddAuthFactor>(request);
 
   if (auto error = TakeOperationError(Operation::kAddAuthFactor);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -1155,6 +1162,7 @@ void FakeUserDataAuthClient::AuthenticateAuthFactor(
     AuthenticateAuthFactorCallback callback) {
   ::user_data_auth::AuthenticateAuthFactorReply reply;
   ReplyOnReturn auto_reply(&reply, std::move(callback));
+  RememberRequest<Operation::kAuthenticateAuthFactor>(request);
 
   if (auto error = TakeOperationError(Operation::kAuthenticateAuthFactor);
       error != CryptohomeErrorCode::CRYPTOHOME_ERROR_NOT_SET) {
@@ -1492,8 +1500,8 @@ void FakeUserDataAuthClient::OnDircryptoMigrationProgressUpdated() {
     NotifyDircryptoMigrationProgress(
         ::user_data_auth::DircryptoMigrationStatus::DIRCRYPTO_MIGRATION_SUCCESS,
         dircrypto_migration_progress_, kDircryptoMigrationMaxProgress);
-    const auto user_it =
-        users_.find(last_migrate_to_dircrypto_request_.account_id());
+    const auto user_it = users_.find(
+        GetLastRequest<Operation::kStartMigrateToDircrypto>().account_id());
     DCHECK(user_it != std::end(users_))
         << "User for dircrypto migration does not exist";
 

diff --git a/chromeos/ash/components/dbus/userdataauth/fake_userdataauth_client.h b/chromeos/ash/components/dbus/userdataauth/fake_userdataauth_client.h
@@ -18,6 +18,7 @@
 #include "chromeos/ash/components/dbus/cryptohome/UserDataAuth.pb.h"
 #include "chromeos/ash/components/dbus/cryptohome/account_identifier_operators.h"
 #include "third_party/abseil-cpp/absl/types/optional.h"
+#include "third_party/protobuf/src/google/protobuf/message_lite.h"
 
 namespace ash {
 
@@ -37,6 +38,7 @@ class COMPONENT_EXPORT(USERDATAAUTH_CLIENT) FakeUserDataAuthClient
     kPrepareVaultForMigration,
     kAddAuthFactor,
     kListAuthFactors,
+    kStartMigrateToDircrypto,
   };
 
   // The method by which a user's home directory can be encrypted.
@@ -265,42 +267,67 @@ class COMPONENT_EXPORT(USERDATAAUTH_CLIENT) FakeUserDataAuthClient
       const ::user_data_auth::TerminateAuthFactorRequest& request,
       TerminateAuthFactorCallback callback) override;
 
-  // Sets the CryptohomeError value to return during next operation.
-  void SetNextOperationError(Operation operation,
-                             ::user_data_auth::CryptohomeErrorCode error);
-
   // Returns the `unlock_webauthn_secret` parameter passed in the last
   // CheckKeyEx call (either successful or not).
   bool get_last_unlock_webauthn_secret() {
     return last_unlock_webauthn_secret_;
   }
 
-  // Getter for the AccountIdentifier() that was passed to the last
-  // StartMigrateToDircrypto() call.
-  const cryptohome::AccountIdentifier& get_id_for_disk_migrated_to_dircrypto()
-      const {
-    return last_migrate_to_dircrypto_request_.account_id();
-  }
-  // Whether the last StartMigrateToDircrypto() call indicates minimal
-  // migration.
-  bool minimal_migration() const {
-    return last_migrate_to_dircrypto_request_.minimal_migration();
-  }
-
   int get_prepare_guest_request_count() const {
     return prepare_guest_request_count_;
   }
 
-  const ::user_data_auth::AddAuthFactorRequest&
-  get_last_add_authfactor_request() const {
-    return last_add_auth_factor_request_;
+  // Per-operation API:
+  template <Operation>
+  struct ProtobufTypes;
+
+  // Template magic to have mapping from operation to
+  // associated types for protobufs.
+#define FUDAC_OPERATION_TYPES(OPERATION, REQUEST)  \
+  template <>                                      \
+  struct ProtobufTypes<Operation::OPERATION> {     \
+    using RequestType = ::user_data_auth::REQUEST; \
+  }
+
+  FUDAC_OPERATION_TYPES(kStartAuthSession, StartAuthSessionRequest);
+  FUDAC_OPERATION_TYPES(kAuthenticateAuthFactor, AuthenticateAuthFactorRequest);
+  FUDAC_OPERATION_TYPES(kPrepareGuestVault, PrepareGuestVaultRequest);
+  FUDAC_OPERATION_TYPES(kPrepareEphemeralVault, PrepareEphemeralVaultRequest);
+  FUDAC_OPERATION_TYPES(kCreatePersistentUser, CreatePersistentUserRequest);
+  FUDAC_OPERATION_TYPES(kPreparePersistentVault, PreparePersistentVaultRequest);
+  FUDAC_OPERATION_TYPES(kPrepareVaultForMigration,
+                        PrepareVaultForMigrationRequest);
+  FUDAC_OPERATION_TYPES(kAddAuthFactor, AddAuthFactorRequest);
+  FUDAC_OPERATION_TYPES(kListAuthFactors, ListAuthFactorsRequest);
+  FUDAC_OPERATION_TYPES(kStartMigrateToDircrypto,
+                        StartMigrateToDircryptoRequest);
+
+#undef FUDAC_OPERATION_TYPES
+
+  // Sets the CryptohomeError value to return during next operation.
+  void SetNextOperationError(Operation operation,
+                             ::user_data_auth::CryptohomeErrorCode error);
+
+  // Checks if operation was called.
+  template <Operation Op>
+  bool WasCalled() {
+    const auto op_request = operation_requests_.find(Op);
+    return op_request != std::end(operation_requests_);
   }
 
-  const ::user_data_auth::AuthenticateAuthFactorRequest&
-  get_last_authenticate_auth_factor_request() const {
-    return last_authenticate_auth_factor_request_;
+  // Provides request protobuf passed to last call of operation.
+  // Would crash if operation was not called before, use `WasCalled()` to
+  // check.
+  template <Operation Op>
+  const typename ProtobufTypes<Op>::RequestType& GetLastRequest() {
+    const auto op_request = operation_requests_.find(Op);
+    CHECK(op_request != std::end(operation_requests_));
+    return *static_cast<const typename ProtobufTypes<Op>::RequestType*>(
+        op_request->second.get());
   }
 
+  // See also `RememberRequest` in private section.
+
   // Calls DircryptoMigrationProgress() on Observer instances.
   void NotifyDircryptoMigrationProgress(
       ::user_data_auth::DircryptoMigrationStatus status,
@@ -322,6 +349,14 @@ class COMPONENT_EXPORT(USERDATAAUTH_CLIENT) FakeUserDataAuthClient
     kAuthFailed,
   };
 
+  // Utility method to remember request passed to operation in a
+  // type-safe way.
+  template <Operation Op>
+  void RememberRequest(const typename ProtobufTypes<Op>::RequestType& request) {
+    operation_requests_[Op] =
+        std::make_unique<typename ProtobufTypes<Op>::RequestType>(request);
+  }
+
   // Helper that returns the protobuf reply.
   template <typename ReplyType>
   void ReturnProtobufMethodCallback(
@@ -369,6 +404,10 @@ class COMPONENT_EXPORT(USERDATAAUTH_CLIENT) FakeUserDataAuthClient
   base::flat_map<Operation, ::user_data_auth::CryptohomeErrorCode>
       operation_errors_;
 
+  // Remembered requests.
+  base::flat_map<Operation, std::unique_ptr<::google::protobuf::MessageLite>>
+      operation_requests_;
+
   // The collection of users we know about.
   base::flat_map<cryptohome::AccountIdentifier, UserCryptohomeState> users_;
 
@@ -379,19 +418,6 @@ class COMPONENT_EXPORT(USERDATAAUTH_CLIENT) FakeUserDataAuthClient
   // the migration progress signal.
   uint64_t dircrypto_migration_progress_ = 0;
 
-  // The StartMigrateToDircryptoRequest passed in for the last
-  // StartMigrateToDircrypto() call.
-  ::user_data_auth::StartMigrateToDircryptoRequest
-      last_migrate_to_dircrypto_request_;
-
-  // The AuthenticateAuthFactorRequest passed in for the last
-  // AuthenticateAuthFactor() call.
-  ::user_data_auth::AuthenticateAuthFactorRequest
-      last_authenticate_auth_factor_request_;
-
-  // The AddAuthFactorRequest passed in for the last AddAuthFactor() call.
-  ::user_data_auth::AddAuthFactorRequest last_add_auth_factor_request_;
-
   // The auth sessions on file.
   base::flat_map<std::string, AuthSessionData> auth_sessions_;